From 159355d281e112a06eff06b21157be17d867a92e Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Mon, 10 Jul 2023 20:23:33 -0500 Subject: [PATCH 1/4] Update has_table to honor schema selection and allow for catalog select Closes #115 Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/dialect/__init__.py | 9 +++--- tests/e2e/sqlalchemy/test_basic.py | 31 +++++++++++++++++++ 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/src/databricks/sqlalchemy/dialect/__init__.py b/src/databricks/sqlalchemy/dialect/__init__.py index 0f96c2bc8..f888aff4f 100644 --- a/src/databricks/sqlalchemy/dialect/__init__.py +++ b/src/databricks/sqlalchemy/dialect/__init__.py @@ -267,17 +267,18 @@ def do_rollback(self, dbapi_connection): # Databricks SQL Does not support transactions pass - def has_table(self, connection, table_name, schema=None, **kwargs) -> bool: + def has_table(self, connection, table_name, schema=None, catalog=None, **kwargs) -> bool: """SQLAlchemy docstrings say dialect providers must implement this method""" - schema = schema or "default" - + _schema = schema or self.schema + _catalog = catalog or self.catalog + # DBR >12.x uses underscores in error messages DBR_LTE_12_NOT_FOUND_STRING = "Table or view not found" DBR_GT_12_NOT_FOUND_STRING = "TABLE_OR_VIEW_NOT_FOUND" try: - res = connection.execute(f"DESCRIBE TABLE {table_name}") + res = connection.execute(f"DESCRIBE TABLE {_catalog}.{_schema}.{table_name}") return True except DatabaseError as e: if DBR_GT_12_NOT_FOUND_STRING in str( diff --git a/tests/e2e/sqlalchemy/test_basic.py b/tests/e2e/sqlalchemy/test_basic.py index 89ceb07e3..b0bafb3be 100644 --- a/tests/e2e/sqlalchemy/test_basic.py +++ b/tests/e2e/sqlalchemy/test_basic.py @@ -340,3 +340,34 @@ def test_get_table_names_smoke_test(samples_engine: Engine): with samples_engine.connect() as conn: _names = samples_engine.table_names(schema="nyctaxi", connection=conn) _names is not None, "get_table_names did not succeed" + +def test_has_table_across_schemas(db_engine: Engine, samples_engine: Engine): + """For this test to pass these conditions must be met: + - Table samples.nyctaxi.trips must exist + - Table samples.tpch.customer must exist + - The `catalog` and `schema` environment variables must be set and valid + """ + + with samples_engine.connect() as conn: + + # Check for table within schema declared at engine creation time + assert samples_engine.dialect.has_table(connection=conn, table_name="trips") + + # Check for table within another schema in the same catalog + assert samples_engine.dialect.has_table(connection=conn, table_name="customer", schema="tpch") + + # Check for a table within a different catalog + other_catalog = os.environ.get("catalog") + other_schema = os.environ.get("schema") + + # Create a table in a different catalog + with db_engine.connect() as conn: + conn.execute("CREATE TABLE test_has_table (numbers_are_cool INT);") + + try: + # Verify with a different engine that this table is not found in the samples catalog + assert not samples_engine.dialect.has_table(connection=conn, table_name="test_has_table") + # Verify with a different engine that this table is found in a separate catalog + assert samples_engine.dialect.has_table(connection=conn, table_name="test_has_table", schema=other_schema, catalog=other_catalog) + finally: + conn.execute("DROP TABLE test_has_table;") \ No newline at end of file From 2d2add75d7c0bc47e0cc264dc2c8d91b9fca6f62 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Mon, 10 Jul 2023 20:24:53 -0500 Subject: [PATCH 2/4] Formatting and update comments Signed-off-by: Jesse Whitehouse --- tests/e2e/sqlalchemy/test_basic.py | 54 ++++++++++++++++++------------ 1 file changed, 32 insertions(+), 22 deletions(-) diff --git a/tests/e2e/sqlalchemy/test_basic.py b/tests/e2e/sqlalchemy/test_basic.py index b0bafb3be..1d3125f2b 100644 --- a/tests/e2e/sqlalchemy/test_basic.py +++ b/tests/e2e/sqlalchemy/test_basic.py @@ -341,6 +341,7 @@ def test_get_table_names_smoke_test(samples_engine: Engine): _names = samples_engine.table_names(schema="nyctaxi", connection=conn) _names is not None, "get_table_names did not succeed" + def test_has_table_across_schemas(db_engine: Engine, samples_engine: Engine): """For this test to pass these conditions must be met: - Table samples.nyctaxi.trips must exist @@ -349,25 +350,34 @@ def test_has_table_across_schemas(db_engine: Engine, samples_engine: Engine): """ with samples_engine.connect() as conn: - - # Check for table within schema declared at engine creation time - assert samples_engine.dialect.has_table(connection=conn, table_name="trips") - - # Check for table within another schema in the same catalog - assert samples_engine.dialect.has_table(connection=conn, table_name="customer", schema="tpch") - - # Check for a table within a different catalog - other_catalog = os.environ.get("catalog") - other_schema = os.environ.get("schema") - - # Create a table in a different catalog - with db_engine.connect() as conn: - conn.execute("CREATE TABLE test_has_table (numbers_are_cool INT);") - - try: - # Verify with a different engine that this table is not found in the samples catalog - assert not samples_engine.dialect.has_table(connection=conn, table_name="test_has_table") - # Verify with a different engine that this table is found in a separate catalog - assert samples_engine.dialect.has_table(connection=conn, table_name="test_has_table", schema=other_schema, catalog=other_catalog) - finally: - conn.execute("DROP TABLE test_has_table;") \ No newline at end of file + + # 1) Check for table within schema declared at engine creation time + assert samples_engine.dialect.has_table(connection=conn, table_name="trips") + + # 2) Check for table within another schema in the same catalog + assert samples_engine.dialect.has_table( + connection=conn, table_name="customer", schema="tpch" + ) + + # 3) Check for a table within a different catalog + other_catalog = os.environ.get("catalog") + other_schema = os.environ.get("schema") + + # Create a table in a different catalog + with db_engine.connect() as conn: + conn.execute("CREATE TABLE test_has_table (numbers_are_cool INT);") + + try: + # Verify that this table is not found in the samples catalog + assert not samples_engine.dialect.has_table( + connection=conn, table_name="test_has_table" + ) + # Verify that this table is found in a separate catalog + assert samples_engine.dialect.has_table( + connection=conn, + table_name="test_has_table", + schema=other_schema, + catalog=other_catalog, + ) + finally: + conn.execute("DROP TABLE test_has_table;") From a3a05d6ad401c08a511f561b9bb5e417f1284f49 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Mon, 10 Jul 2023 21:41:30 -0500 Subject: [PATCH 3/4] Black the codebase Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/dialect/__init__.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/databricks/sqlalchemy/dialect/__init__.py b/src/databricks/sqlalchemy/dialect/__init__.py index f888aff4f..cfb7d8573 100644 --- a/src/databricks/sqlalchemy/dialect/__init__.py +++ b/src/databricks/sqlalchemy/dialect/__init__.py @@ -267,18 +267,22 @@ def do_rollback(self, dbapi_connection): # Databricks SQL Does not support transactions pass - def has_table(self, connection, table_name, schema=None, catalog=None, **kwargs) -> bool: + def has_table( + self, connection, table_name, schema=None, catalog=None, **kwargs + ) -> bool: """SQLAlchemy docstrings say dialect providers must implement this method""" _schema = schema or self.schema _catalog = catalog or self.catalog - + # DBR >12.x uses underscores in error messages DBR_LTE_12_NOT_FOUND_STRING = "Table or view not found" DBR_GT_12_NOT_FOUND_STRING = "TABLE_OR_VIEW_NOT_FOUND" try: - res = connection.execute(f"DESCRIBE TABLE {_catalog}.{_schema}.{table_name}") + res = connection.execute( + f"DESCRIBE TABLE {_catalog}.{_schema}.{table_name}" + ) return True except DatabaseError as e: if DBR_GT_12_NOT_FOUND_STRING in str( From 2e78f9c0c0c91737decb9894d7922f5004edf24e Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Tue, 11 Jul 2023 20:06:06 -0500 Subject: [PATCH 4/4] Update changelog Signed-off-by: Jesse Whitehouse --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8be708c9e..a7693b8a9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ ## 2.7.x (Unreleased) - Add support for Cloud Fetch +- SQLAlchemy has_table function now honours schema= argument and adds catalog= argument - Fix: Revised SQLAlchemy dialect and examples for compatibility with SQLAlchemy==1.3.x - Fix: oauth would fail if expired credentials appeared in ~/.netrc - Fix: Python HTTP proxies were broken after switch to urllib3