databricks-sqlalchemy 0.0.1b1__py3-none-any.whl → 1.0.0__py3-none-any.whl

databricks/sqlalchemy/base.py

@@ -0,0 +1,436 @@
+from typing import Any, List, Optional, Dict, Union
+
+import databricks.sqlalchemy._ddl as dialect_ddl_impl
+import databricks.sqlalchemy._types as dialect_type_impl
+from databricks import sql
+from databricks.sqlalchemy._parse import (
+    _describe_table_extended_result_to_dict_list,
+    _match_table_not_found_string,
+    build_fk_dict,
+    build_pk_dict,
+    get_fk_strings_from_dte_output,
+    get_pk_strings_from_dte_output,
+    get_comment_from_dte_output,
+    parse_column_info_from_tgetcolumnsresponse,
+)
+
+import sqlalchemy
+from sqlalchemy import DDL, event
+from sqlalchemy.engine import Connection, Engine, default, reflection
+from sqlalchemy.engine.interfaces import (
+    ReflectedForeignKeyConstraint,
+    ReflectedPrimaryKeyConstraint,
+    ReflectedColumn,
+    ReflectedTableComment,
+)
+from sqlalchemy.engine.reflection import ReflectionDefaults
+from sqlalchemy.exc import DatabaseError, SQLAlchemyError
+
+try:
+    import alembic
+except ImportError:
+    pass
+else:
+    from alembic.ddl import DefaultImpl
+
+    class DatabricksImpl(DefaultImpl):
+        __dialect__ = "databricks"
+
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class DatabricksDialect(default.DefaultDialect):
+    """This dialect implements only those methods required to pass our e2e tests"""
+
+    # See sqlalchemy.engine.interfaces for descriptions of each of these properties
+    name: str = "databricks"
+    driver: str = "databricks"
+    default_schema_name: str = "default"
+    preparer = dialect_ddl_impl.DatabricksIdentifierPreparer  # type: ignore
+    ddl_compiler = dialect_ddl_impl.DatabricksDDLCompiler
+    statement_compiler = dialect_ddl_impl.DatabricksStatementCompiler
+    supports_statement_cache: bool = True
+    supports_multivalues_insert: bool = True
+    supports_native_decimal: bool = True
+    supports_sane_rowcount: bool = False
+    non_native_boolean_check_constraint: bool = False
+    supports_identity_columns: bool = True
+    supports_schemas: bool = True
+    default_paramstyle: str = "named"
+    div_is_floordiv: bool = False
+    supports_default_values: bool = False
+    supports_server_side_cursors: bool = False
+    supports_sequences: bool = False
+    supports_native_boolean: bool = True
+
+    colspecs = {
+        sqlalchemy.types.DateTime: dialect_type_impl.TIMESTAMP_NTZ,
+        sqlalchemy.types.Time: dialect_type_impl.DatabricksTimeType,
+        sqlalchemy.types.String: dialect_type_impl.DatabricksStringType,
+    }
+
+    # SQLAlchemy requires that a table with no primary key
+    # constraint return a dictionary that looks like this.
+    EMPTY_PK: Dict[str, Any] = {"constrained_columns": [], "name": None}
+
+    # SQLAlchemy requires that a table with no foreign keys
+    # defined return an empty list. Same for indexes.
+    EMPTY_FK: List
+    EMPTY_INDEX: List
+    EMPTY_FK = EMPTY_INDEX = []
+
+    @classmethod
+    def import_dbapi(cls):
+        return sql
+
+    def _force_paramstyle_to_native_mode(self):
+        """This method can be removed after databricks-sql-connector wholly switches to NATIVE ParamApproach.
+
+        This is a hack to trick SQLAlchemy into using a different paramstyle
+        than the one declared by this module in src/databricks/sql/__init__.py
+
+        This method is called _after_ the dialect has been initialised, which is important because otherwise
+        our users would need to include a `paramstyle` argument in their SQLAlchemy connection string.
+
+        This dialect is written to support NATIVE queries. Although the INLINE approach can technically work,
+        the same behaviour can be achieved within SQLAlchemy itself using its literal_processor methods.
+        """
+
+        self.paramstyle = self.default_paramstyle
+
+    def create_connect_args(self, url):
+        # TODO: can schema be provided after HOST?
+        # Expected URI format is: databricks+thrift://token:dapi***@***.cloud.databricks.com?http_path=/sql/***
+
+        kwargs = {
+            "server_hostname": url.host,
+            "access_token": url.password,
+            "http_path": url.query.get("http_path"),
+            "catalog": url.query.get("catalog"),
+            "schema": url.query.get("schema"),
+            "use_inline_params": False,
+        }
+
+        self.schema = kwargs["schema"]
+        self.catalog = kwargs["catalog"]
+
+        self._force_paramstyle_to_native_mode()
+
+        return [], kwargs
+
+    def get_columns(
+        self, connection, table_name, schema=None, **kwargs
+    ) -> List[ReflectedColumn]:
+        """Return information about columns in `table_name`."""
+
+        with self.get_connection_cursor(connection) as cur:
+            resp = cur.columns(
+                catalog_name=self.catalog,
+                schema_name=schema or self.schema,
+                table_name=table_name,
+            ).fetchall()
+
+        if not resp:
+            # TGetColumnsRequest will not raise an exception if passed a table that doesn't exist
+            # But Databricks supports tables with no columns. So if the result is an empty list,
+            # we need to check if the table exists (and raise an exception if not) or simply return
+            # an empty list.
+            self._describe_table_extended(
+                connection,
+                table_name,
+                self.catalog,
+                schema or self.schema,
+                expect_result=False,
+            )
+            return resp
+        columns = []
+        for col in resp:
+            row_dict = parse_column_info_from_tgetcolumnsresponse(col)
+            columns.append(row_dict)
+
+        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,
+    ) -> Union[List[Dict[str, str]], None]:
+        """Run DESCRIBE TABLE EXTENDED on a table and return a list of dictionaries 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)
+        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_list(result)
+        return fmt_result
+
+    @reflection.cache
+    def get_pk_constraint(
+        self,
+        connection,
+        table_name: str,
+        schema: Optional[str] = None,
+        **kw: Any,
+    ) -> ReflectedPrimaryKeyConstraint:
+        """Fetch information about the primary key constraint on table_name.
+
+        Returns a dictionary with these 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
+        """
+
+        result = self._describe_table_extended(
+            connection=connection,
+            table_name=table_name,
+            schema_name=schema,
+        )
+
+        # Type ignore is because mypy knows that self._describe_table_extended *can*
+        # return None (even though it never will since expect_result defaults to True)
+        raw_pk_constraints: List = get_pk_strings_from_dte_output(result)  # type: ignore
+        if not any(raw_pk_constraints):
+            return self.EMPTY_PK  # type: ignore
+
+        if len(raw_pk_constraints) > 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."
+            )
+
+        first_pk_constraint = raw_pk_constraints[0]
+        pk_name = first_pk_constraint.get("col_name")
+        pk_constraint_string = first_pk_constraint.get("data_type")
+
+        # TODO: figure out how to return sqlalchemy.interfaces in a way that mypy respects
+        return build_pk_dict(pk_name, pk_constraint_string)  # type: ignore
+
+    def get_foreign_keys(
+        self, connection, table_name, schema=None, **kw
+    ) -> List[ReflectedForeignKeyConstraint]:
+        """Return information about foreign_keys in `table_name`."""
+
+        result = self._describe_table_extended(
+            connection=connection,
+            table_name=table_name,
+            schema_name=schema,
+        )
+
+        # Type ignore is because mypy knows that self._describe_table_extended *can*
+        # return None (even though it never will since expect_result defaults to True)
+        raw_fk_constraints: List = get_fk_strings_from_dte_output(result)  # type: ignore
+
+        if not any(raw_fk_constraints):
+            return self.EMPTY_FK
+
+        fk_constraints = []
+        for constraint_dict in raw_fk_constraints:
+            fk_name = constraint_dict.get("col_name")
+            fk_constraint_string = constraint_dict.get("data_type")
+            this_constraint_dict = build_fk_dict(
+                fk_name, fk_constraint_string, schema_name=schema
+            )
+            fk_constraints.append(this_constraint_dict)
+
+        # TODO: figure out how to return sqlalchemy.interfaces in a way that mypy respects
+        return fk_constraints  # type: ignore
+
+    def get_indexes(self, connection, table_name, schema=None, **kw):
+        """SQLAlchemy requires this method. Databricks doesn't support indexes."""
+        return self.EMPTY_INDEX
+
+    @reflection.cache
+    def get_table_names(self, connection: Connection, schema=None, **kwargs):
+        """Return a list of tables in the current schema."""
+
+        _target_catalog = self.catalog
+        _target_schema = schema or self.schema
+        _target = f"`{_target_catalog}`.`{_target_schema}`"
+
+        stmt = DDL(f"SHOW TABLES FROM {_target}")
+
+        tables_result = connection.execute(stmt).all()
+        views_result = self.get_view_names(connection=connection, schema=schema)
+
+        # In Databricks, SHOW TABLES FROM <schema> returns both tables and views.
+        # Potential optimisation: rewrite this to instead query information_schema
+        tables_minus_views = [
+            row.tableName for row in tables_result if row.tableName not in views_result
+        ]
+
+        return tables_minus_views
+
+    @reflection.cache
+    def get_view_names(
+        self,
+        connection,
+        schema=None,
+        only_materialized=False,
+        only_temp=False,
+        **kwargs,
+    ) -> List[str]:
+        """Returns a list of string view names contained in the schema, if any."""
+
+        _target_catalog = self.catalog
+        _target_schema = schema or self.schema
+        _target = f"`{_target_catalog}`.`{_target_schema}`"
+
+        stmt = DDL(f"SHOW VIEWS FROM {_target}")
+        result = connection.execute(stmt).all()
+
+        return [
+            row.viewName
+            for row in result
+            if (not only_materialized or row.isMaterialized)
+            and (not only_temp or row.isTemporary)
+        ]
+
+    @reflection.cache
+    def get_materialized_view_names(
+        self, connection: Connection, schema: Optional[str] = None, **kw: Any
+    ) -> List[str]:
+        """A wrapper around get_view_names that fetches only the names of materialized views"""
+        return self.get_view_names(connection, schema, only_materialized=True)
+
+    @reflection.cache
+    def get_temp_view_names(
+        self, connection: Connection, schema: Optional[str] = None, **kw: Any
+    ) -> List[str]:
+        """A wrapper around get_view_names that fetches only the names of temporary views"""
+        return self.get_view_names(connection, schema, only_temp=True)
+
+    def do_rollback(self, dbapi_connection):
+        # Databricks SQL Does not support transactions
+        pass
+
+    @reflection.cache
+    def has_table(
+        self, connection, table_name, schema=None, catalog=None, **kwargs
+    ) -> bool:
+        """For internal dialect use, check the existence of a particular table
+        or view in the database.
+        """
+
+        try:
+            self._describe_table_extended(
+                connection=connection,
+                table_name=table_name,
+                catalog_name=catalog,
+                schema_name=schema,
+            )
+            return True
+        except sqlalchemy.exc.NoSuchTableError as e:
+            return False
+
+    def get_connection_cursor(self, connection):
+        """Added for backwards compatibility with 1.3.x"""
+        if hasattr(connection, "_dbapi_connection"):
+            return connection._dbapi_connection.dbapi_connection.cursor()
+        elif hasattr(connection, "raw_connection"):
+            return connection.raw_connection().cursor()
+        elif hasattr(connection, "connection"):
+            return connection.connection.cursor()
+
+        raise SQLAlchemyError(
+            "Databricks dialect can't obtain a cursor context manager from the dbapi"
+        )
+
+    @reflection.cache
+    def get_schema_names(self, connection, **kw):
+        """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
+
+    @reflection.cache
+    def get_table_comment(
+        self,
+        connection: Connection,
+        table_name: str,
+        schema: Optional[str] = None,
+        **kw: Any,
+    ) -> ReflectedTableComment:
+        result = self._describe_table_extended(
+            connection=connection,
+            table_name=table_name,
+            schema_name=schema,
+        )
+
+        if result is None:
+            return ReflectionDefaults.table_comment()
+
+        comment = get_comment_from_dte_output(result)
+
+        if comment:
+            return dict(text=comment)
+        else:
+            return ReflectionDefaults.table_comment()
+
+
+@event.listens_for(Engine, "do_connect")
+def receive_do_connect(dialect, conn_rec, cargs, cparams):
+    """Helpful for DS on traffic from clients using SQLAlchemy in particular"""
+
+    # Ignore connect invocations that don't use our dialect
+    if not dialect.name == "databricks":
+        return
+
+    ua = cparams.get("_user_agent_entry", "")
+
+    def add_sqla_tag_if_not_present(val: str):
+        if not val:
+            output = "sqlalchemy"
+
+        if val and "sqlalchemy" in val:
+            output = val
+
+        else:
+            output = f"sqlalchemy + {val}"
+
+        return output
+
+    cparams["_user_agent_entry"] = add_sqla_tag_if_not_present(ua)
+
+    if sqlalchemy.__version__.startswith("1.3"):
+        # SQLAlchemy 1.3.x fails to parse the http_path, catalog, and schema from our connection string
+        # These should be passed in as connect_args when building the Engine
+
+        if "schema" in cparams:
+            dialect.schema = cparams["schema"]
+
+        if "catalog" in cparams:
+            dialect.catalog = cparams["catalog"]

databricks/sqlalchemy/dependency_test/test_dependency.py

@@ -0,0 +1,22 @@
+import pytest
+
+class DatabricksImportError(Exception):
+    pass
+
+class TestLibraryDependencySuite:
+
+    @pytest.mark.skipif(pytest.importorskip("databricks_sql_connector_core"), reason="databricks_sql_connector_core is present")
+    def test_sql_core(self):
+        with pytest.raises(DatabricksImportError, match="databricks_sql_connector_core module is not available"):
+            try:
+                import databricks
+            except ImportError:
+                raise DatabricksImportError("databricks_sql_connector_core module is not available")
+
+    @pytest.mark.skipif(pytest.importorskip("sqlalchemy"), reason="SQLAlchemy is present")
+    def test_sqlalchemy(self):
+        with pytest.raises(DatabricksImportError, match="sqlalchemy module is not available"):
+            try:
+                import sqlalchemy
+            except ImportError:
+                raise DatabricksImportError("sqlalchemy module is not available")

databricks/sqlalchemy/pytest.ini

@@ -0,0 +1,4 @@
+
+[sqla_testing]
+requirement_cls=databricks.sqlalchemy.requirements:Requirements
+profile_file=profiles.txt