From ee8a0f25910fa3b892b1045559bc10b216a35f41 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Thu, 12 Oct 2023 00:46:27 -0400 Subject: [PATCH 1/7] Immediately skip any tests on temp tables because we don't support them 5 skipped, 3 warnings, 727 errors in 40.01s Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/test/test_suite.py | 182 ++----------------- 1 file changed, 12 insertions(+), 170 deletions(-) diff --git a/src/databricks/sqlalchemy/test/test_suite.py b/src/databricks/sqlalchemy/test/test_suite.py index a16e97096..4d4d15701 100644 --- a/src/databricks/sqlalchemy/test/test_suite.py +++ b/src/databricks/sqlalchemy/test/test_suite.py @@ -402,188 +402,30 @@ def test_empty_insert_multiple(self): """ +@pytest.mark.reviewed class ComponentReflectionTest(ComponentReflectionTest): - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_autoincrement_col(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_dialect_initialize(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_columns(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_comments(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_comments_with_schema(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_default_schema_name(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_foreign_keys(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_foreign_keys(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_indexes(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_inter_schema_foreign_keys(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_noncol_index(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_pk_constraint(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_schema_names(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_schema_names_w_translate_map(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_table_names(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_table_oid(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_table_oid(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") + @pytest.mark.skip(reason="Databricks doesn't support temp tables.") def test_get_temp_table_columns(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ + pass - @pytest.mark.skip(reason="Error during execution. Requires investigation.") + @pytest.mark.skip(reason="Databricks doesn't support temp tables.") def test_get_temp_table_indexes(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ + pass - @pytest.mark.skip(reason="Error during execution. Requires investigation.") + @pytest.mark.skip(reason="Databricks doesn't support temp tables.") def test_get_temp_table_names(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ + pass - @pytest.mark.skip(reason="Error during execution. Requires investigation.") + @pytest.mark.skip(reason="Databricks doesn't support temp tables.") def test_get_temp_table_unique_constraints(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_temp_view_columns(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_temp_view_names(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ + pass - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_unique_constraints(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ + @pytest.mark.skip(reason="Databricks doesn't support temp tables.") + def test_reflect_table_temp_table(self): + pass - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_view_definition(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ - @pytest.mark.skip(reason="Error during execution. Requires investigation.") - def test_get_view_definition(self): - """ - Exception: - - NotImplementedError: no temp table keyword args routine for cfg: databricks+databricks://token:***redacted***@e2-dogfood.staging.cloud.databricks.com?catalog=main&http_path=%2Fsql%2F1.0%2Fwarehouses%2F5c89f447c476a5a8&schema=pysql_sqlalchemy - """ @pytest.mark.reviewed From 3527b043af4c00fc96ca81d0c7fb74e138a58800 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Thu, 12 Oct 2023 23:01:07 -0400 Subject: [PATCH 2/7] This commit gets the test_not_existing_table test to pass It requires the implementation of raising NoSuchTableError for the different methods it checks. And hooboy was this difficult to track down! I found that these tests were mostly failing during setup and teardown because I didn't understand that two different schemas are required. This is covered in the SQLAlchemy tests documentation here: https://github.com/sqlalchemy/sqlalchemy/blob/main/README.unittests.rst#database-configuration The skipped tests happen because of exclusions, which is why the skip log information isn't as pleasant as our other tests. test_suite.py::ComponentReflectionTest_databricks+databricks::test_not_existing_table[get_check_constraints-_exclusions_06] SKIPPED test_suite.py::ComponentReflectionTest_databricks+databricks::test_not_existing_table[get_columns-_exclusions_01] b'table_does_not_exists' PASSED test_suite.py::ComponentReflectionTest_databricks+databricks::test_not_existing_table[get_foreign_keys-_exclusions_03] b'No such table table_does_not_exists' PASSED test_suite.py::ComponentReflectionTest_databricks+databricks::test_not_existing_table[get_indexes-_exclusions_04] SKIPPED test_suite.py::ComponentReflectionTest_databricks+databricks::test_not_existing_table[get_pk_constraint-_exclusions_02] b'No such table table_does_not_exists' PASSED test_suite.py::ComponentReflectionTest_databricks+databricks::test_not_existing_table[get_table_comment-_exclusions_07] SKIPPED test_suite.py::ComponentReflectionTest_databricks+databricks::test_not_existing_table[get_table_options-_exclusions_00] SKIPPED test_suite.py::ComponentReflectionTest_databricks+databricks::test_not_existing_table[get_unique_constraints-_exclusions_05] SKIPPED Signed-off-by: Jesse Whitehouse --- CONTRIBUTING.md | 4 +++ src/databricks/sqlalchemy/__init__.py | 33 +++++++++++++---- src/databricks/sqlalchemy/requirements.py | 38 ++++++++++++++++++-- src/databricks/sqlalchemy/test/test_suite.py | 12 +++++-- 4 files changed, 74 insertions(+), 13 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c3e64a4a6..eebf3353b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -158,6 +158,10 @@ python -m pytest test_suite.py --dburi \ Some of these of these tests fail currently. We're working on getting relevant tests passing and others skipped. The tests that we've already reviewed and verified are decorated with a pytest marker called `reviewed`. To only run these tests and check for regressions, you can add `-m reviewed` to the invocation command above. +These tests require two schemas exist in your catalog: +- An empty schema which can have an arbitrary name. It is configured in the SQLAlchemy --dburi +- An empty schema named `test_schema` + ### Code formatting This project uses [Black](https://pypi.org/project/black/). diff --git a/src/databricks/sqlalchemy/__init__.py b/src/databricks/sqlalchemy/__init__.py index 9a81bda56..11ba3b13b 100644 --- a/src/databricks/sqlalchemy/__init__.py +++ b/src/databricks/sqlalchemy/__init__.py @@ -11,6 +11,7 @@ from sqlalchemy.exc import DatabaseError, SQLAlchemyError import databricks.sqlalchemy._ddl as dialect_ddl_impl +from databricks.sql.exc import ServerOperationError # This import is required to process our @compiles decorators import databricks.sqlalchemy._types as dialect_type_impl @@ -30,6 +31,8 @@ class DatabricksImpl(DefaultImpl): __dialect__ = "databricks" +DBR_LTE_12_NOT_FOUND_STRING = "Table or view not found" +DBR_GT_12_NOT_FOUND_STRING = "TABLE_OR_VIEW_NOT_FOUND" class DatabricksDialect(default.DefaultDialect): """This dialect implements only those methods required to pass our e2e tests""" @@ -135,6 +138,8 @@ def get_columns(self, connection, table_name, schema=None, **kwargs): table_name=table_name, ).fetchall() + if not resp: + raise sqlalchemy.exc.NoSuchTableError(table_name) columns = [] for col in resp: @@ -164,8 +169,16 @@ def get_pk_constraint( """ with self.get_connection_cursor(connection) as cursor: - # DESCRIBE TABLE EXTENDED doesn't support parameterised inputs :( - result = cursor.execute(f"DESCRIBE TABLE EXTENDED {table_name}").fetchall() + + try: + # DESCRIBE TABLE EXTENDED doesn't support parameterised inputs :( + result = cursor.execute(f"DESCRIBE TABLE EXTENDED {table_name}").fetchall() + except ServerOperationError as e: + if DBR_GT_12_NOT_FOUND_STRING in str( + e + ) or DBR_LTE_12_NOT_FOUND_STRING in str(e): + raise sqlalchemy.exc.NoSuchTableError(f"No such table {table_name}") from e + # DESCRIBE TABLE EXTENDED doesn't give a deterministic name to the field where # a primary key constraint will be found in its output. So we cycle through its @@ -214,11 +227,17 @@ def get_foreign_keys( table_name`. """ - with self.get_connection_cursor(connection) as cursor: - # DESCRIBE TABLE EXTENDED doesn't support parameterised inputs :( - result = cursor.execute( - f"DESCRIBE TABLE EXTENDED {schema + '.' if schema else ''}{table_name}" - ).fetchall() + try: + with self.get_connection_cursor(connection) as cursor: + # DESCRIBE TABLE EXTENDED doesn't support parameterised inputs :( + result = cursor.execute( + f"DESCRIBE TABLE EXTENDED {schema + '.' if schema else ''}{table_name}" + ).fetchall() + except ServerOperationError as e: + if DBR_GT_12_NOT_FOUND_STRING in str( + e + ) or DBR_LTE_12_NOT_FOUND_STRING in str(e): + raise sqlalchemy.exc.NoSuchTableError(f"No such table {table_name}") from e # DESCRIBE TABLE EXTENDED doesn't give a deterministic name to the field where # a foreign key constraint will be found in its output. So we cycle through its diff --git a/src/databricks/sqlalchemy/requirements.py b/src/databricks/sqlalchemy/requirements.py index d8229dace..4d4013927 100644 --- a/src/databricks/sqlalchemy/requirements.py +++ b/src/databricks/sqlalchemy/requirements.py @@ -133,14 +133,46 @@ def has_temp_table(self): we're in a bind. """ return sqlalchemy.testing.exclusions.closed() - + @property def temporary_views(self): """target database supports temporary views""" return sqlalchemy.testing.exclusions.open() - + @property def views(self): """Target database must support VIEWs.""" - return sqlalchemy.testing.exclusions.open() \ No newline at end of file + return sqlalchemy.testing.exclusions.open() + + @property + def temporary_tables(self): + """target database supports temporary tables + + ComponentReflection test is intricate and simply cannot function without this exclusion being defined here. + This happens because we cannot skip individual combinations used in ComponentReflection test. + """ + return sqlalchemy.testing.exclusions.closed() + + @property + def temp_table_reflection(self): + """ComponentReflection test is intricate and simply cannot function without this exclusion being defined here. + This happens because we cannot skip individual combinations used in ComponentReflection test. + """ + return sqlalchemy.testing.exclusions.closed() + + @property + def index_reflection(self): + """ComponentReflection test is intricate and simply cannot function without this exclusion being defined here. + This happens because we cannot skip individual combinations used in ComponentReflection test. + """ + return sqlalchemy.testing.exclusions.closed() + + @property + def unique_constraint_reflection(self): + """ComponentReflection test is intricate and simply cannot function without this exclusion being defined here. + This happens because we cannot skip individual combinations used in ComponentReflection test. + + Databricks supports unique constraints but they are not implemented in this dialect. + """ + return sqlalchemy.testing.exclusions.closed() diff --git a/src/databricks/sqlalchemy/test/test_suite.py b/src/databricks/sqlalchemy/test/test_suite.py index 4d4d15701..38cf33be1 100644 --- a/src/databricks/sqlalchemy/test/test_suite.py +++ b/src/databricks/sqlalchemy/test/test_suite.py @@ -228,7 +228,9 @@ def test_drop_index_if_exists(self): @pytest.mark.reviewed -@pytest.mark.skip(reason="Identity works. Test needs rewrite for Databricks. See comments in test_suite.py") +@pytest.mark.skip( + reason="Identity works. Test needs rewrite for Databricks. See comments in test_suite.py" +) class IdentityColumnTest(IdentityColumnTest): """The setup for these tests tries to create a table with a DELTA IDENTITY column but has two problems: 1. It uses an Integer() type for the column. Whereas DELTA IDENTITY columns must be BIGINT. @@ -239,6 +241,7 @@ class IdentityColumnTest(IdentityColumnTest): I'm satisified through manual testing that our implementation of visit_identity_column works but a better test is needed. """ + pass @@ -326,6 +329,7 @@ class LastrowidTest(LastrowidTest): class CompositeKeyReflectionTest(CompositeKeyReflectionTest): pass + class ComponentReflectionTestExtra(ComponentReflectionTestExtra): @pytest.mark.skip(reason="Test setup needs adjustment.") def test_varchar_reflection(self): @@ -404,6 +408,10 @@ def test_empty_insert_multiple(self): @pytest.mark.reviewed class ComponentReflectionTest(ComponentReflectionTest): + """This test requires two schemas be present in the target Databricks workspace: + - The schema set in --dburi + - A second schema named "test_schema" + """ @pytest.mark.skip(reason="Databricks doesn't support temp tables.") def test_get_temp_table_columns(self): @@ -426,8 +434,6 @@ def test_reflect_table_temp_table(self): pass - - @pytest.mark.reviewed class TableDDLTest(TableDDLTest): @pytest.mark.skip(reason="Databricks does not support indexes.") From c957e7dd56e265ae0992b5acc6dd264b19506622 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Thu, 12 Oct 2023 23:11:35 -0400 Subject: [PATCH 3/7] Make test_get_schema_names PASS Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/__init__.py | 12 +++++++----- src/databricks/sqlalchemy/test/test_suite.py | 4 ++++ 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/src/databricks/sqlalchemy/__init__.py b/src/databricks/sqlalchemy/__init__.py index 11ba3b13b..7b49b737e 100644 --- a/src/databricks/sqlalchemy/__init__.py +++ b/src/databricks/sqlalchemy/__init__.py @@ -2,7 +2,7 @@ from typing import Any, Optional import sqlalchemy -from sqlalchemy import event +from sqlalchemy import event, DDL from sqlalchemy.engine import Engine, default, reflection from sqlalchemy.engine.interfaces import ( ReflectedForeignKeyConstraint, @@ -372,10 +372,12 @@ def get_connection_cursor(self, connection): @reflection.cache def get_schema_names(self, connection, **kw): - # Equivalent to SHOW DATABASES - - # TODO: replace with call to cursor.schemas() once its performance matches raw SQL - return [row[0] for row in connection.execute("SHOW SCHEMAS")] + """Return a list of all schema names available in the database. + """ + stmt = DDL("SHOW SCHEMAS") + result = connection.execute(stmt) + schema_list = [row[0] for row in result] + return schema_list @event.listens_for(Engine, "do_connect") diff --git a/src/databricks/sqlalchemy/test/test_suite.py b/src/databricks/sqlalchemy/test/test_suite.py index 38cf33be1..93096b509 100644 --- a/src/databricks/sqlalchemy/test/test_suite.py +++ b/src/databricks/sqlalchemy/test/test_suite.py @@ -413,6 +413,10 @@ class ComponentReflectionTest(ComponentReflectionTest): - A second schema named "test_schema" """ + # We've reviewed these tests: + # test_get_schema_names + # test_not_existing_table + @pytest.mark.skip(reason="Databricks doesn't support temp tables.") def test_get_temp_table_columns(self): pass From 48604ef411d9cd46ae3060b4411feda9adca56a8 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Fri, 13 Oct 2023 11:57:39 -0400 Subject: [PATCH 4/7] Move COLUMN_TYPE_MAP into module so it won't be instantiated on every call to get_columns Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/__init__.py | 42 +++++++++++++-------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/src/databricks/sqlalchemy/__init__.py b/src/databricks/sqlalchemy/__init__.py index 7b49b737e..14e0d7486 100644 --- a/src/databricks/sqlalchemy/__init__.py +++ b/src/databricks/sqlalchemy/__init__.py @@ -34,6 +34,26 @@ class DatabricksImpl(DefaultImpl): DBR_LTE_12_NOT_FOUND_STRING = "Table or view not found" DBR_GT_12_NOT_FOUND_STRING = "TABLE_OR_VIEW_NOT_FOUND" +COLUMN_TYPE_MAP = { + "boolean": sqlalchemy.types.Boolean, + "smallint": sqlalchemy.types.SmallInteger, + "int": sqlalchemy.types.Integer, + "bigint": sqlalchemy.types.BigInteger, + "float": sqlalchemy.types.Float, + "double": sqlalchemy.types.Float, + "string": sqlalchemy.types.String, + "varchar": sqlalchemy.types.String, + "char": sqlalchemy.types.String, + "binary": sqlalchemy.types.String, + "array": sqlalchemy.types.String, + "map": sqlalchemy.types.String, + "struct": sqlalchemy.types.String, + "uniontype": sqlalchemy.types.String, + "decimal": sqlalchemy.types.Numeric, + "timestamp": sqlalchemy.types.DateTime, + "date": sqlalchemy.types.Date, +} + class DatabricksDialect(default.DefaultDialect): """This dialect implements only those methods required to pass our e2e tests""" @@ -111,26 +131,6 @@ def get_columns(self, connection, table_name, schema=None, **kwargs): Additional column attributes may be present. """ - _type_map = { - "boolean": sqlalchemy.types.Boolean, - "smallint": sqlalchemy.types.SmallInteger, - "int": sqlalchemy.types.Integer, - "bigint": sqlalchemy.types.BigInteger, - "float": sqlalchemy.types.Float, - "double": sqlalchemy.types.Float, - "string": sqlalchemy.types.String, - "varchar": sqlalchemy.types.String, - "char": sqlalchemy.types.String, - "binary": sqlalchemy.types.String, - "array": sqlalchemy.types.String, - "map": sqlalchemy.types.String, - "struct": sqlalchemy.types.String, - "uniontype": sqlalchemy.types.String, - "decimal": sqlalchemy.types.Numeric, - "timestamp": sqlalchemy.types.DateTime, - "date": sqlalchemy.types.Date, - } - with self.get_connection_cursor(connection) as cur: resp = cur.columns( catalog_name=self.catalog, @@ -147,7 +147,7 @@ def get_columns(self, connection, table_name, schema=None, **kwargs): _col_type = re.search(r"^\w+", col.TYPE_NAME).group(0) this_column = { "name": col.COLUMN_NAME, - "type": _type_map[_col_type.lower()], + "type": COLUMN_TYPE_MAP[_col_type.lower()], "nullable": bool(col.NULLABLE), "default": col.COLUMN_DEF, "autoincrement": False if col.IS_AUTO_INCREMENT == "NO" else True, From ad8edbcff91731c588cf2f64399ee4a96467583c Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Fri, 13 Oct 2023 12:54:25 -0400 Subject: [PATCH 5/7] Refactor dialect's has_table method to extract _describe_table_extended This passes the HasTableTest test group in the dialect compliance test suite Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/__init__.py | 115 ++++++++++++++++++++------ 1 file changed, 88 insertions(+), 27 deletions(-) diff --git a/src/databricks/sqlalchemy/__init__.py b/src/databricks/sqlalchemy/__init__.py index 14e0d7486..a59a4d04a 100644 --- a/src/databricks/sqlalchemy/__init__.py +++ b/src/databricks/sqlalchemy/__init__.py @@ -3,7 +3,7 @@ import sqlalchemy from sqlalchemy import event, DDL -from sqlalchemy.engine import Engine, default, reflection +from sqlalchemy.engine import Engine, default, reflection, Connection, Row, CursorResult from sqlalchemy.engine.interfaces import ( ReflectedForeignKeyConstraint, ReflectedPrimaryKeyConstraint, @@ -31,9 +31,35 @@ class DatabricksImpl(DefaultImpl): __dialect__ = "databricks" + DBR_LTE_12_NOT_FOUND_STRING = "Table or view not found" DBR_GT_12_NOT_FOUND_STRING = "TABLE_OR_VIEW_NOT_FOUND" + +def _match_table_not_found_string(message: str) -> bool: + """Return True if the message contains a substring indicating that a table was not found""" + return any( + [ + DBR_LTE_12_NOT_FOUND_STRING in message, + DBR_GT_12_NOT_FOUND_STRING in message, + ] + ) + + +def _describe_table_extended_result_to_dict(result: CursorResult) -> dict: + """Transform the output of DESCRIBE TABLE EXTENDED into a dictionary + + The output from DESCRIBE TABLE EXTENDED puts all values in the `data_type` column + Even CONSTRAINT descriptions are contained in the `data_type` column + Some rows have an empty string for their col_name. These are present only for spacing + so we ignore them. + """ + + result_dict = {row.col_name: row.data_type for row in result if row.col_name != ""} + + return result_dict + + COLUMN_TYPE_MAP = { "boolean": sqlalchemy.types.Boolean, "smallint": sqlalchemy.types.SmallInteger, @@ -54,6 +80,7 @@ class DatabricksImpl(DefaultImpl): "date": sqlalchemy.types.Date, } + class DatabricksDialect(default.DefaultDialect): """This dialect implements only those methods required to pass our e2e tests""" @@ -156,6 +183,46 @@ def get_columns(self, connection, table_name, schema=None, **kwargs): return columns + def _describe_table_extended( + self, + connection: Connection, + table_name: str, + catalog_name: Optional[str] = None, + schema_name: Optional[str] = None, + expect_result=True, + ): + """Run DESCRIBE TABLE EXTENDED on a table and return a dictionary of the result. + + This method is the fastest way to check for the presence of a table in a schema. + + If expect_result is False, this method returns None as the output dict isn't required. + + Raises NoSuchTableError if the table is not present in the schema. + """ + + _target_catalog = catalog_name or self.catalog + _target_schema = schema_name or self.schema + _target = f"`{_target_catalog}`.`{_target_schema}`.`{table_name}`" + + # sql injection risk? + # DESCRIBE TABLE EXTENDED in DBR doesn't support parameterised inputs :( + stmt = DDL(f"DESCRIBE TABLE EXTENDED {_target}") + + try: + result = connection.execute(stmt).all() + except DatabaseError as e: + if _match_table_not_found_string(str(e)): + raise sqlalchemy.exc.NoSuchTableError( + f"No such table {table_name}" + ) from e + raise e + + if not expect_result: + return None + + fmt_result = _describe_table_extended_result_to_dict(result) + return fmt_result + @reflection.cache def get_pk_constraint( self, @@ -169,16 +236,18 @@ def get_pk_constraint( """ with self.get_connection_cursor(connection) as cursor: - try: # DESCRIBE TABLE EXTENDED doesn't support parameterised inputs :( - result = cursor.execute(f"DESCRIBE TABLE EXTENDED {table_name}").fetchall() + result = cursor.execute( + f"DESCRIBE TABLE EXTENDED {table_name}" + ).fetchall() except ServerOperationError as e: if DBR_GT_12_NOT_FOUND_STRING in str( e ) or DBR_LTE_12_NOT_FOUND_STRING in str(e): - raise sqlalchemy.exc.NoSuchTableError(f"No such table {table_name}") from e - + raise sqlalchemy.exc.NoSuchTableError( + f"No such table {table_name}" + ) from e # DESCRIBE TABLE EXTENDED doesn't give a deterministic name to the field where # a primary key constraint will be found in its output. So we cycle through its @@ -237,7 +306,9 @@ def get_foreign_keys( if DBR_GT_12_NOT_FOUND_STRING in str( e ) or DBR_LTE_12_NOT_FOUND_STRING in str(e): - raise sqlalchemy.exc.NoSuchTableError(f"No such table {table_name}") from e + raise sqlalchemy.exc.NoSuchTableError( + f"No such table {table_name}" + ) from e # DESCRIBE TABLE EXTENDED doesn't give a deterministic name to the field where # a foreign key constraint will be found in its output. So we cycle through its @@ -333,29 +404,20 @@ def do_rollback(self, dbapi_connection): 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" + """For internal dialect use, check the existence of a particular table + or view in the database. + """ try: - res = connection.execute( - sqlalchemy.text( - f"DESCRIBE TABLE `{_catalog}`.`{_schema}`.`{table_name}`" - ) + self._describe_table_extended( + connection=connection, + table_name=table_name, + catalog_name=catalog, + schema_name=schema, ) return True - except DatabaseError as e: - if DBR_GT_12_NOT_FOUND_STRING in str( - e - ) or DBR_LTE_12_NOT_FOUND_STRING in str(e): - return False - else: - raise e + except sqlalchemy.exc.NoSuchTableError as e: + return False def get_connection_cursor(self, connection): """Added for backwards compatibility with 1.3.x""" @@ -372,8 +434,7 @@ def get_connection_cursor(self, connection): @reflection.cache def get_schema_names(self, connection, **kw): - """Return a list of all schema names available in the database. - """ + """Return a list of all schema names available in the database.""" stmt = DDL("SHOW SCHEMAS") result = connection.execute(stmt) schema_list = [row[0] for row in result] From e15542f5d4dbafcb1d7a3b85c52ce6307d254d89 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Fri, 13 Oct 2023 14:35:46 -0400 Subject: [PATCH 6/7] Refactor primary and foreign key parsing This now PASSES test_get_primary_keys and test_get_pk_constraint Needs further refactor into actual parse methods... Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/__init__.py | 261 ++++++++++-------- src/databricks/sqlalchemy/requirements.py | 6 + .../sqlalchemy/test_local/test_utils.py | 12 + 3 files changed, 165 insertions(+), 114 deletions(-) diff --git a/src/databricks/sqlalchemy/__init__.py b/src/databricks/sqlalchemy/__init__.py index a59a4d04a..47d3ad6e5 100644 --- a/src/databricks/sqlalchemy/__init__.py +++ b/src/databricks/sqlalchemy/__init__.py @@ -1,5 +1,5 @@ import re -from typing import Any, Optional +from typing import Any, Optional, List import sqlalchemy from sqlalchemy import event, DDL @@ -19,6 +19,7 @@ from databricks.sqlalchemy.utils import ( extract_identifier_groups_from_string, extract_identifiers_from_string, + extract_three_level_identifier_from_constraint_string ) try: @@ -32,6 +33,10 @@ class DatabricksImpl(DefaultImpl): __dialect__ = "databricks" +import logging + +logger = logging.getLogger(__name__) + DBR_LTE_12_NOT_FOUND_STRING = "Table or view not found" DBR_GT_12_NOT_FOUND_STRING = "TABLE_OR_VIEW_NOT_FOUND" @@ -60,6 +65,133 @@ def _describe_table_extended_result_to_dict(result: CursorResult) -> dict: return result_dict +def _extract_pk_from_dte_result(result: dict) -> ReflectedPrimaryKeyConstraint: + """Return a dictionary with the keys: + + constrained_columns + a list of column names that make up the primary key. Results is an empty list + if no PRIMARY KEY is defined. + + name + the name of the primary key constraint + + Today, DESCRIBE TABLE EXTENDED doesn't give a deterministic name to the field where + a primary key constraint will be found in its output. So we cycle through its + output looking for a match that includes "PRIMARY KEY". This is brittle. We + could optionally make two roundtrips: the first would query information_schema + for the name of the primary key constraint on this table, and a second to + DESCRIBE TABLE EXTENDED, at which point we would know the name of the constraint. + But for now we instead assume that Python list comprehension is faster than a + network roundtrip. + """ + + # find any rows that contain "PRIMARY KEY" as the `data_type` + filtered_rows = [(k, v) for k, v in result.items() if "PRIMARY KEY" in v] + + # bail if no primary key was found + if not filtered_rows: + return {"constrained_columns": [], "name": None} + + # there should only ever be one PRIMARY KEY that matches + if len(filtered_rows) > 1: + logger.warning( + "Found more than one primary key constraint in DESCRIBE TABLE EXTENDED output. " + "This is unexpected. Please report this as a bug. " + "Only the first primary key constraint will be returned." + ) + + # target is a tuple of (constraint_name, constraint_string) + target = filtered_rows[0] + name = target[0] + _constraint_string = target[1] + column_list = extract_identifiers_from_string(_constraint_string) + + return {"constrained_columns": column_list, "name": name} + + +def _extract_single_fk_dict_from_dte_result_row( + table_name: str, schema_name: Optional[str], fk_name: str, fk_constraint_string: str +) -> dict: + """ + """ + + # SQLAlchemy's ComponentReflectionTest::test_get_foreign_keys is strange in that it + # expects the `referred_schema` member of the outputted dictionary to be None if + # a `schema` argument was not passed to the dialect's `get_foreign_keys` method + referred_table_dict = extract_three_level_identifier_from_constraint_string(fk_constraint_string) + referred_table = referred_table_dict["table"] + if schema_name: + referred_schema = referred_table_dict["schema"] + else: + referred_schema = None + + _extracted = extract_identifier_groups_from_string(fk_constraint_string) + constrained_columns_str, referred_columns_str = ( + _extracted[0], + _extracted[1], + ) + + constrainted_columns = extract_identifiers_from_string(constrained_columns_str) + referred_columns = extract_identifiers_from_string(referred_columns_str) + + return { + "constrained_columns": constrainted_columns, + "name": fk_name, + "referred_table": referred_table, + "referred_columns": referred_columns, + "referred_schema": referred_schema, + } + + +def _extract_fk_from_dte_result( + table_name: str, schema_name: Optional[str], result: dict +) -> ReflectedForeignKeyConstraint: + """Return a list of dictionaries with the keys: + + constrained_columns + a list of column names that make up the foreign key + + name + the name of the foreign key constraint + + referred_table + the name of the table that the foreign key references + + referred_columns + a list of column names that are referenced by the foreign key + + Returns an empty list if no foreign key is defined. + + Today, DESCRIBE TABLE EXTENDED doesn't give a deterministic name to the field where + a foreign key constraint will be found in its output. So we cycle through its + output looking for a match that includes "FOREIGN KEY". This is brittle. We + could optionally make two roundtrips: the first would query information_schema + for the name of the foreign key constraint on this table, and a second to + DESCRIBE TABLE EXTENDED, at which point we would know the name of the constraint. + But for now we instead assume that Python list comprehension is faster than a + network roundtrip. + """ + + # find any rows that contain "FOREIGN_KEY" as the `data_type` + filtered_rows = [(k, v) for k, v in result.items() if "FOREIGN KEY" in v] + + # bail if no foreign key was found + if not filtered_rows: + return [] + + constraint_list = [] + + # target is a tuple of (constraint_name, constraint_string) + for target in filtered_rows: + _constraint_name, _constraint_string = target + this_constraint_dict = _extract_single_fk_dict_from_dte_result_row( + table_name, schema_name, _constraint_name, _constraint_string + ) + constraint_list.append(this_constraint_dict) + + return constraint_list + + COLUMN_TYPE_MAP = { "boolean": sqlalchemy.types.Boolean, "smallint": sqlalchemy.types.SmallInteger, @@ -235,125 +367,26 @@ def get_pk_constraint( table_name`. """ - with self.get_connection_cursor(connection) as cursor: - try: - # DESCRIBE TABLE EXTENDED doesn't support parameterised inputs :( - result = cursor.execute( - f"DESCRIBE TABLE EXTENDED {table_name}" - ).fetchall() - except ServerOperationError as e: - if DBR_GT_12_NOT_FOUND_STRING in str( - e - ) or DBR_LTE_12_NOT_FOUND_STRING in str(e): - raise sqlalchemy.exc.NoSuchTableError( - f"No such table {table_name}" - ) from e - - # DESCRIBE TABLE EXTENDED doesn't give a deterministic name to the field where - # a primary key constraint will be found in its output. So we cycle through its - # output looking for a match that includes "PRIMARY KEY". This is brittle. We - # could optionally make two roundtrips: the first would query information_schema - # for the name of the primary key constraint on this table, and a second to - # DESCRIBE TABLE EXTENDED, at which point we would know the name of the constraint. - # But for now we instead assume that Python list comprehension is faster than a - # network roundtrip. - dte_dict = {row["col_name"]: row["data_type"] for row in result} - target = [(k, v) for k, v in dte_dict.items() if "PRIMARY KEY" in v] - if target: - name, _constraint_string = target[0] - column_list = extract_identifiers_from_string(_constraint_string) - else: - name, column_list = None, None - - return {"constrained_columns": column_list, "name": name} + result = self._describe_table_extended( + connection=connection, + table_name=table_name, + schema_name=schema, + ) + + return _extract_pk_from_dte_result(result) def get_foreign_keys( self, connection, table_name, schema=None, **kw ) -> ReflectedForeignKeyConstraint: - """Return information about foreign_keys in `table_name`. - - Given a :class:`_engine.Connection`, a string - `table_name`, and an optional string `schema`, return foreign - key information as a list of dicts with these keys: - - name - the constraint's name + """Return information about foreign_keys in `table_name`.""" - constrained_columns - a list of column names that make up the foreign key - - referred_schema - the name of the referred schema - - referred_table - the name of the referred table - - referred_columns - a list of column names in the referred table that correspond to - constrained_columns - """ - """Return information about the primary key constraint on - table_name`. - """ - - try: - with self.get_connection_cursor(connection) as cursor: - # DESCRIBE TABLE EXTENDED doesn't support parameterised inputs :( - result = cursor.execute( - f"DESCRIBE TABLE EXTENDED {schema + '.' if schema else ''}{table_name}" - ).fetchall() - except ServerOperationError as e: - if DBR_GT_12_NOT_FOUND_STRING in str( - e - ) or DBR_LTE_12_NOT_FOUND_STRING in str(e): - raise sqlalchemy.exc.NoSuchTableError( - f"No such table {table_name}" - ) from e - - # DESCRIBE TABLE EXTENDED doesn't give a deterministic name to the field where - # a foreign key constraint will be found in its output. So we cycle through its - # output looking for a match that includes "FOREIGN KEY". This is brittle. We - # could optionally make two roundtrips: the first would query information_schema - # for the name of the foreign key constraint on this table, and a second to - # DESCRIBE TABLE EXTENDED, at which point we would know the name of the constraint. - # But for now we instead assume that Python list comprehension is faster than a - # network roundtrip. - dte_dict = {row["col_name"]: row["data_type"] for row in result} - target = [(k, v) for k, v in dte_dict.items() if "FOREIGN KEY" in v] - - def extract_constraint_dict_from_target(target): - if target: - name, _constraint_string = target - _extracted = extract_identifier_groups_from_string(_constraint_string) - constrained_columns_str, referred_columns_str = ( - _extracted[0], - _extracted[1], - ) - - constrained_columns = extract_identifiers_from_string( - constrained_columns_str - ) - referred_columns = extract_identifiers_from_string(referred_columns_str) - referred_table = str(table_name) - else: - name, constrained_columns, referred_columns, referred_table = ( - None, - None, - None, - None, - ) - - return { - "constrained_columns": constrained_columns, - "name": name, - "referred_table": referred_table, - "referred_columns": referred_columns, - } + result = self._describe_table_extended( + connection=connection, + table_name=table_name, + schema_name=schema, + ) - if target: - return [extract_constraint_dict_from_target(i) for i in target] - else: - return [] + return _extract_fk_from_dte_result(table_name, schema, result) def get_indexes(self, connection, table_name, schema=None, **kw): """Return information about indexes in `table_name`. diff --git a/src/databricks/sqlalchemy/requirements.py b/src/databricks/sqlalchemy/requirements.py index 4d4013927..614eea2e6 100644 --- a/src/databricks/sqlalchemy/requirements.py +++ b/src/databricks/sqlalchemy/requirements.py @@ -176,3 +176,9 @@ def unique_constraint_reflection(self): Databricks supports unique constraints but they are not implemented in this dialect. """ return sqlalchemy.testing.exclusions.closed() + + @property + def reflects_pk_names(self): + """Target driver reflects the name of primary key constraints.""" + + return sqlalchemy.testing.exclusions.open() \ No newline at end of file diff --git a/src/databricks/sqlalchemy/test_local/test_utils.py b/src/databricks/sqlalchemy/test_local/test_utils.py index ecb9fd437..b31e0cbcf 100644 --- a/src/databricks/sqlalchemy/test_local/test_utils.py +++ b/src/databricks/sqlalchemy/test_local/test_utils.py @@ -2,6 +2,7 @@ from databricks.sqlalchemy.utils import ( extract_identifiers_from_string, extract_identifier_groups_from_string, + extract_three_level_identifier_from_constraint_string ) @@ -36,3 +37,14 @@ def test_extract_identifer_batches(input, expected): assert ( extract_identifier_groups_from_string(input) == expected ), "Failed to extract identifier groups from string" + +def test_extract_3l_namespace_from_constraint_string(): + + input = "FOREIGN KEY (`parent_user_id`) REFERENCES `main`.`pysql_dialect_compliance`.`users` (`user_id`)" + expected = { + "catalog": "main", + "schema": "pysql_dialect_compliance", + "table": "users" + } + + assert extract_three_level_identifier_from_constraint_string(input) == expected, "Failed to extract 3L namespace from constraint string" \ No newline at end of file From bea9f0b6b73ab371695e884d0e87e6fcb65c6fbe Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Fri, 13 Oct 2023 19:00:02 -0400 Subject: [PATCH 7/7] Refactor foreign key parsing completely into the _parse.py file Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/__init__.py | 67 ++------ src/databricks/sqlalchemy/_parse.py | 145 ++++++++++++++++++ .../{test_utils.py => test_parsing.py} | 22 ++- src/databricks/sqlalchemy/utils.py | 23 --- 4 files changed, 175 insertions(+), 82 deletions(-) create mode 100644 src/databricks/sqlalchemy/_parse.py rename src/databricks/sqlalchemy/test_local/{test_utils.py => test_parsing.py} (69%) delete mode 100644 src/databricks/sqlalchemy/utils.py diff --git a/src/databricks/sqlalchemy/__init__.py b/src/databricks/sqlalchemy/__init__.py index 47d3ad6e5..325f6c513 100644 --- a/src/databricks/sqlalchemy/__init__.py +++ b/src/databricks/sqlalchemy/__init__.py @@ -1,5 +1,5 @@ import re -from typing import Any, Optional, List +from typing import Any, Optional, List, Tuple import sqlalchemy from sqlalchemy import event, DDL @@ -16,10 +16,10 @@ # This import is required to process our @compiles decorators import databricks.sqlalchemy._types as dialect_type_impl from databricks import sql -from databricks.sqlalchemy.utils import ( - extract_identifier_groups_from_string, +from databricks.sqlalchemy._parse import ( + build_fk_dict, extract_identifiers_from_string, - extract_three_level_identifier_from_constraint_string + extract_three_level_identifier_from_constraint_string, ) try: @@ -109,56 +109,11 @@ def _extract_pk_from_dte_result(result: dict) -> ReflectedPrimaryKeyConstraint: return {"constrained_columns": column_list, "name": name} -def _extract_single_fk_dict_from_dte_result_row( - table_name: str, schema_name: Optional[str], fk_name: str, fk_constraint_string: str -) -> dict: - """ - """ - - # SQLAlchemy's ComponentReflectionTest::test_get_foreign_keys is strange in that it - # expects the `referred_schema` member of the outputted dictionary to be None if - # a `schema` argument was not passed to the dialect's `get_foreign_keys` method - referred_table_dict = extract_three_level_identifier_from_constraint_string(fk_constraint_string) - referred_table = referred_table_dict["table"] - if schema_name: - referred_schema = referred_table_dict["schema"] - else: - referred_schema = None - - _extracted = extract_identifier_groups_from_string(fk_constraint_string) - constrained_columns_str, referred_columns_str = ( - _extracted[0], - _extracted[1], - ) - - constrainted_columns = extract_identifiers_from_string(constrained_columns_str) - referred_columns = extract_identifiers_from_string(referred_columns_str) - - return { - "constrained_columns": constrainted_columns, - "name": fk_name, - "referred_table": referred_table, - "referred_columns": referred_columns, - "referred_schema": referred_schema, - } - - def _extract_fk_from_dte_result( - table_name: str, schema_name: Optional[str], result: dict + result: dict, schema_name: Optional[str] ) -> ReflectedForeignKeyConstraint: - """Return a list of dictionaries with the keys: - - constrained_columns - a list of column names that make up the foreign key - - name - the name of the foreign key constraint - - referred_table - the name of the table that the foreign key references - - referred_columns - a list of column names that are referenced by the foreign key + """Extract a list of foreign key information dictionaries from the result + of a DESCRIBE TABLE EXTENDED call. Returns an empty list if no foreign key is defined. @@ -173,7 +128,7 @@ def _extract_fk_from_dte_result( """ # find any rows that contain "FOREIGN_KEY" as the `data_type` - filtered_rows = [(k, v) for k, v in result.items() if "FOREIGN KEY" in v] + filtered_rows: List[Tuple] = [(k, v) for k, v in result.items() if "FOREIGN KEY" in v] # bail if no foreign key was found if not filtered_rows: @@ -184,8 +139,8 @@ def _extract_fk_from_dte_result( # target is a tuple of (constraint_name, constraint_string) for target in filtered_rows: _constraint_name, _constraint_string = target - this_constraint_dict = _extract_single_fk_dict_from_dte_result_row( - table_name, schema_name, _constraint_name, _constraint_string + this_constraint_dict = build_fk_dict( + _constraint_name, _constraint_string, schema_name ) constraint_list.append(this_constraint_dict) @@ -386,7 +341,7 @@ def get_foreign_keys( schema_name=schema, ) - return _extract_fk_from_dte_result(table_name, schema, result) + return _extract_fk_from_dte_result(result, schema) def get_indexes(self, connection, table_name, schema=None, **kw): """Return information about indexes in `table_name`. diff --git a/src/databricks/sqlalchemy/_parse.py b/src/databricks/sqlalchemy/_parse.py new file mode 100644 index 000000000..587b1381d --- /dev/null +++ b/src/databricks/sqlalchemy/_parse.py @@ -0,0 +1,145 @@ +from typing import List, Optional +import re + +""" +This module contains helper functions that can parse the contents +of DESCRIBE TABLE EXTENDED calls. Mostly wrappers around regexes. +""" + +def extract_identifiers_from_string(input_str: str) -> List[str]: + """For a string input resembling (`a`, `b`, `c`) return a list of identifiers ['a', 'b', 'c']""" + + # This matches the valid character list contained in DatabricksIdentifierPreparer + pattern = re.compile(r"`([A-Za-z0-9_]+)`") + matches = pattern.findall(input_str) + return [i for i in matches] + + +def extract_identifier_groups_from_string(input_str: str) -> List[str]: + """For a string input resembling : + + FOREIGN KEY (`pname`, `pid`, `pattr`) REFERENCES `main`.`pysql_sqlalchemy`.`tb1` (`name`, `id`, `attr`) + + Return ['(`pname`, `pid`, `pattr`)', '(`name`, `id`, `attr`)'] + """ + pattern = re.compile(r"\([`A-Za-z0-9_,\s]*\)") + matches = pattern.findall(input_str) + return [i for i in matches] + + +def extract_three_level_identifier_from_constraint_string(input_str: str) -> dict: + """For a string input resembling : + FOREIGN KEY (`parent_user_id`) REFERENCES `main`.`pysql_dialect_compliance`.`users` (`user_id`) + + Return a dict like + { + "catalog": "main", + "schema": "pysql_dialect_compliance", + "table": "users" + } + """ + pat = re.compile(r"REFERENCES\s+(.*?)\s*\(") + matches = pat.findall(input_str) + + if not matches: + return None + + first_match = matches[0] + parts = first_match.split(".") + + def strip_backticks(input:str): + return input.replace("`", "") + + return { + "catalog": strip_backticks(parts[0]), + "schema": strip_backticks(parts[1]), + "table": strip_backticks(parts[2]) + } + +def _parse_fk_from_constraint_string(constraint_str: str) -> dict: + """Build a dictionary of foreign key constraint information from a constraint string. + + For example: + + ``` + FOREIGN KEY (`pname`, `pid`, `pattr`) REFERENCES `main`.`pysql_dialect_compliance`.`tb1` (`name`, `id`, `attr`) + ``` + + Return a dictionary like: + + ``` + { + "constrained_columns": ["pname", "pid", "pattr"], + "referred_table": "tb1", + "referred_schema": "pysql_dialect_compliance", + "referred_columns": ["name", "id", "attr"] + } + ``` + + Note that the constraint name doesn't appear in the constraint string so it will not + be present in the output of this function. + """ + + referred_table_dict = extract_three_level_identifier_from_constraint_string( + constraint_str + ) + referred_table = referred_table_dict["table"] + referred_schema = referred_table_dict["schema"] + + # _extracted is a tuple of two lists of identifiers + # we assume the first immediately follows "FOREIGN KEY" and the second + # immediately follows REFERENCES $tableName + _extracted = extract_identifier_groups_from_string(constraint_str) + constrained_columns_str, referred_columns_str = ( + _extracted[0], + _extracted[1], + ) + + constrained_columns = extract_identifiers_from_string(constrained_columns_str) + referred_columns = extract_identifiers_from_string(referred_columns_str) + + return { + "constrained_columns": constrained_columns, + "referred_table": referred_table, + "referred_columns": referred_columns, + "referred_schema": referred_schema, + } + +def build_fk_dict( + fk_name: str, fk_constraint_string: str, schema_name: Optional[str] +) -> dict: + """ + Given a foriegn key name and a foreign key constraint string, return a dictionary + with the following keys: + + name + the name of the foreign key constraint + constrained_columns + a list of column names that make up the foreign key + referred_table + the name of the table that the foreign key references + referred_columns + a list of column names that are referenced by the foreign key + referred_schema + the name of the schema that the foreign key references. + + referred schema will be None if the schema_name argument is None. + This is required by SQLAlchey's ComponentReflectionTest::test_get_foreign_keys + """ + + # The foreign key name is not contained in the constraint string so we + # need to add it manually + base_fk_dict = _parse_fk_from_constraint_string(fk_constraint_string) + + if not schema_name: + schema_override_dict = dict(referred_schema=None) + else: + schema_override_dict = {} + + complete_foreign_key_dict = { + "name": fk_name, + **base_fk_dict, + **schema_override_dict, + } + + return complete_foreign_key_dict \ No newline at end of file diff --git a/src/databricks/sqlalchemy/test_local/test_utils.py b/src/databricks/sqlalchemy/test_local/test_parsing.py similarity index 69% rename from src/databricks/sqlalchemy/test_local/test_utils.py rename to src/databricks/sqlalchemy/test_local/test_parsing.py index b31e0cbcf..3ebb8616e 100644 --- a/src/databricks/sqlalchemy/test_local/test_utils.py +++ b/src/databricks/sqlalchemy/test_local/test_parsing.py @@ -1,8 +1,9 @@ import pytest -from databricks.sqlalchemy.utils import ( +from databricks.sqlalchemy._parse import ( extract_identifiers_from_string, extract_identifier_groups_from_string, - extract_three_level_identifier_from_constraint_string + extract_three_level_identifier_from_constraint_string, + build_fk_dict ) @@ -47,4 +48,19 @@ def test_extract_3l_namespace_from_constraint_string(): "table": "users" } - assert extract_three_level_identifier_from_constraint_string(input) == expected, "Failed to extract 3L namespace from constraint string" \ No newline at end of file + assert extract_three_level_identifier_from_constraint_string(input) == expected, "Failed to extract 3L namespace from constraint string" + +@pytest.mark.parametrize("schema", [None, "some_schema"]) +def test_build_fk_dict(schema): + fk_constraint_string = "FOREIGN KEY (`parent_user_id`) REFERENCES `main`.`some_schema`.`users` (`user_id`)" + + result = build_fk_dict("some_fk_name", fk_constraint_string, schema_name=schema) + + assert result == { + "name": "some_fk_name", + "constrained_columns": ["parent_user_id"], + "referred_schema": schema, + "referred_table": "users", + "referred_columns": ["user_id"], + } + diff --git a/src/databricks/sqlalchemy/utils.py b/src/databricks/sqlalchemy/utils.py deleted file mode 100644 index d13dbdd1f..000000000 --- a/src/databricks/sqlalchemy/utils.py +++ /dev/null @@ -1,23 +0,0 @@ -from typing import List -import re - - -def extract_identifiers_from_string(input_str: str) -> List[str]: - """For a string input resembling (`a`, `b`, `c`) return a list of identifiers ['a', 'b', 'c']""" - - # This matches the valid character list contained in DatabricksIdentifierPreparer - pattern = re.compile(r"`([A-Za-z0-9_]+)`") - matches = pattern.findall(input_str) - return [i for i in matches] - - -def extract_identifier_groups_from_string(input_str: str) -> List[str]: - """For a string input resembling : - - FOREIGN KEY (`pname`, `pid`, `pattr`) REFERENCES `main`.`pysql_sqlalchemy`.`tb1` (`name`, `id`, `attr`) - - Return ['(`pname`, `pid`, `pattr`)', '(`name`, `id`, `attr`)'] - """ - pattern = re.compile(r"\([`A-Za-z0-9_,\s]*\)") - matches = pattern.findall(input_str) - return [i for i in matches]