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..325f6c513 100644 --- a/src/databricks/sqlalchemy/__init__.py +++ b/src/databricks/sqlalchemy/__init__.py @@ -1,9 +1,9 @@ import re -from typing import Any, Optional +from typing import Any, Optional, List, Tuple import sqlalchemy -from sqlalchemy import event -from sqlalchemy.engine import Engine, default, reflection +from sqlalchemy import event, DDL +from sqlalchemy.engine import Engine, default, reflection, Connection, Row, CursorResult from sqlalchemy.engine.interfaces import ( ReflectedForeignKeyConstraint, ReflectedPrimaryKeyConstraint, @@ -11,13 +11,15 @@ 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 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, ) try: @@ -31,6 +33,141 @@ 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" + + +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 + + +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_fk_from_dte_result( + result: dict, schema_name: Optional[str] +) -> ReflectedForeignKeyConstraint: + """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. + + 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: 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: + 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 = build_fk_dict( + _constraint_name, _constraint_string, schema_name + ) + constraint_list.append(this_constraint_dict) + + return constraint_list + + +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""" @@ -108,26 +245,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, @@ -135,6 +252,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: @@ -142,7 +261,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, @@ -151,6 +270,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, @@ -163,107 +322,26 @@ def get_pk_constraint( 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 {table_name}").fetchall() - - # 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`. - """ - - 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() - - # 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(result, schema) def get_indexes(self, connection, table_name, schema=None, **kw): """Return information about indexes in `table_name`. @@ -314,29 +392,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""" @@ -353,10 +422,11 @@ 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/_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/requirements.py b/src/databricks/sqlalchemy/requirements.py index d8229dace..614eea2e6 100644 --- a/src/databricks/sqlalchemy/requirements.py +++ b/src/databricks/sqlalchemy/requirements.py @@ -133,14 +133,52 @@ 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() + + @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() + + @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/test_suite.py b/src/databricks/sqlalchemy/test/test_suite.py index a16e97096..93096b509 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): @@ -402,188 +406,36 @@ 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 - """ + """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="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 - """ + # We've reviewed these tests: + # test_get_schema_names + # test_not_existing_table - @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 - """ - - @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="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 - """ + 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="Databricks doesn't support temp tables.") + def test_reflect_table_temp_table(self): + pass @pytest.mark.reviewed diff --git a/src/databricks/sqlalchemy/test_local/test_utils.py b/src/databricks/sqlalchemy/test_local/test_parsing.py similarity index 50% rename from src/databricks/sqlalchemy/test_local/test_utils.py rename to src/databricks/sqlalchemy/test_local/test_parsing.py index ecb9fd437..3ebb8616e 100644 --- a/src/databricks/sqlalchemy/test_local/test_utils.py +++ b/src/databricks/sqlalchemy/test_local/test_parsing.py @@ -1,7 +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, + build_fk_dict ) @@ -36,3 +38,29 @@ 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" + +@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]