databricks-sqlalchemy 1.0.0__py3-none-any.whl → 1.0.2__py3-none-any.whl

CHANGELOG.md

@@ -1,274 +1,5 @@
 # Release History
 
-# 3.3.0 (2024-07-18)
+# 1.0.1 
 
-- Don't retry requests that fail with HTTP code 401 (databricks/databricks-sql-python#408 by @Hodnebo)
-- Remove username/password (aka "basic") auth option (databricks/databricks-sql-python#409 by @jackyhu-db)
-- Refactor CloudFetch handler to fix numerous issues with it (databricks/databricks-sql-python#405 by @kravets-levko)
-- Add option to disable SSL verification for CloudFetch links (databricks/databricks-sql-python#414 by @kravets-levko)
-
-Databricks-managed passwords reached end of life on July 10, 2024. Therefore, Basic auth support was removed from
-the library. See https://docs.databricks.com/en/security/auth-authz/password-deprecation.html
-
-The existing option `_tls_no_verify=True` of `sql.connect(...)` will now also disable SSL cert verification
-(but not the SSL itself) for CloudFetch links. This option should be used as a workaround only, when other ways
-to fix SSL certificate errors didn't work.
-
-# 3.2.0 (2024-06-06)
-
-- Update proxy authentication (databricks/databricks-sql-python#354 by @amir-haroun)
-- Relax `pyarrow` pin (databricks/databricks-sql-python#389 by @dhirschfeld)
-- Fix error logging in OAuth manager (databricks/databricks-sql-python#269 by @susodapop)
-- SQLAlchemy: enable delta.feature.allowColumnDefaults for all tables (databricks/databricks-sql-python#343 by @dhirschfeld)
-- Update `thrift` dependency (databricks/databricks-sql-python#397 by @m1n0)
-
-# 3.1.2 (2024-04-18)
-
-- Remove broken cookie code (#379)
-- Small typing fixes (#382, #384 thanks @wyattscarpenter)
-
-# 3.1.1 (2024-03-19)
-
-- Don't retry requests that fail with code 403 (#373)
-- Assume a default retry-after for 429/503 (#371)
-- Fix boolean literals (#357)
-
-# 3.1.0 (2024-02-16)
-
-- Revert retry-after behavior to be exponential backoff (#349)
-- Support Databricks OAuth on Azure (#351)
-- Support Databricks OAuth on GCP (#338)
-
-# 3.0.3 (2024-02-02)
-
-- Revised docstrings and examples for OAuth (#339)
-- Redact the URL query parameters from the urllib3.connectionpool logs (#341)
-
-# 3.0.2 (2024-01-25)
-
-- SQLAlchemy dialect now supports table and column comments (thanks @cbornet!)
-- Fix: SQLAlchemy dialect now correctly reflects TINYINT types (thanks @TimTheinAtTabs!)
-- Fix: `server_hostname` URIs that included `https://` would raise an exception
-- Other: pinned to `pandas<=2.1` and `urllib3>=1.26` to avoid runtime errors in dbt-databricks (#330)
-
-## 3.0.1 (2023-12-01)
-
-- Other: updated docstring comment about default parameterization approach (#287)
-- Other: added tests for reading complex types and revised docstrings and type hints (#293)
-- Fix: SQLAlchemy dialect raised DeprecationWarning due to `dbapi` classmethod (#294)
-- Fix: SQLAlchemy dialect could not reflect TIMESTAMP_NTZ columns (#296)
-
-## 3.0.0 (2023-11-17)
-
-- Remove support for Python 3.7
-- Add support for native parameterized SQL queries. Requires DBR 14.2 and above. See docs/parameters.md for more info.
-- Completely rewritten SQLAlchemy dialect
-  - Adds support for SQLAlchemy >= 2.0 and drops support for SQLAlchemy 1.x
-  - Full e2e test coverage of all supported features
-  - Detailed usage notes in `README.sqlalchemy.md`
-  - Adds support for:
-    - New types: `TIME`, `TIMESTAMP`, `TIMESTAMP_NTZ`, `TINYINT`
-    - `Numeric` type scale and precision, like `Numeric(10,2)`
-    - Reading and writing `PrimaryKeyConstraint` and `ForeignKeyConstraint`
-    - Reading and writing composite keys
-    - Reading and writing from views
-    - Writing `Identity` to tables (i.e. autoincrementing primary keys)
-    - `LIMIT` and `OFFSET` for paging through results
-    - Caching metadata calls
-- Enable cloud fetch by default. To disable, set `use_cloud_fetch=False` when building `databricks.sql.client`.
-- Add integration tests for Databricks UC Volumes ingestion queries
-- Retries:
-  - Add `_retry_max_redirects` config
-  - Set `_enable_v3_retries=True` and warn if users override it
-- Security: bump minimum pyarrow version to 14.0.1 (CVE-2023-47248)
-
-## 2.9.3 (2023-08-24)
-
-- Fix: Connections failed when urllib3~=1.0.0 is installed (#206)
-
-## 2.9.2 (2023-08-17)
-
-**Note: this release was yanked from Pypi on 13 September 2023 due to compatibility issues with environments where `urllib3<=2.0.0` were installed. The log changes are incorporated into version 2.9.3 and greater.**
-
-- Other: Add `examples/v3_retries_query_execute.py` (#199)
-- Other: suppress log message when `_enable_v3_retries` is not `True` (#199)
-- Other: make this connector backwards compatible with `urllib3>=1.0.0` (#197)
-
-## 2.9.1 (2023-08-11)
-
-**Note: this release was yanked from Pypi on 13 September 2023 due to compatibility issues with environments where `urllib3<=2.0.0` were installed.**
-
-- Other: Explicitly pin urllib3 to ^2.0.0 (#191)
-
-## 2.9.0 (2023-08-10)
-
-- Replace retry handling with DatabricksRetryPolicy. This is disabled by default. To enable, set `_enable_v3_retries=True` when creating `databricks.sql.client` (#182)
-- Other: Fix typo in README quick start example (#186)
-- Other: Add autospec to Client mocks and tidy up `make_request` (#188)
-
-## 2.8.0 (2023-07-21)
-
-- Add support for Cloud Fetch. Disabled by default. Set `use_cloud_fetch=True` when building `databricks.sql.client` to enable it (#146, #151, #154)
-- SQLAlchemy has_table function now honours schema= argument and adds catalog= argument (#174)
-- SQLAlchemy set non_native_boolean_check_constraint False as it's not supported by Databricks (#120)
-- Fix: Revised SQLAlchemy dialect and examples for compatibility with SQLAlchemy==1.3.x (#173)
-- Fix: oauth would fail if expired credentials appeared in ~/.netrc (#122)
-- Fix: Python HTTP proxies were broken after switch to urllib3 (#158)
-- Other: remove unused import in SQLAlchemy dialect
-- Other: Relax pandas dependency constraint to allow ^2.0.0 (#164)
-- Other: Connector now logs operation handle guids as hexadecimal instead of bytes (#170)
-- Other: test_socket_timeout_user_defined e2e test was broken (#144)
-
-## 2.7.0 (2023-06-26)
-
-- Fix: connector raised exception when calling close() on a closed Thrift session
-- Improve e2e test development ergonomics
-- Redact logged thrift responses by default
-- Add support for OAuth on Databricks Azure
-
-## 2.6.2 (2023-06-14)
-
-- Fix: Retry GetOperationStatus requests for http errors
-
-## 2.6.1 (2023-06-08)
-
-- Fix: http.client would raise a BadStatusLine exception in some cases
-
-## 2.6.0 (2023-06-07)
-
-- Add support for HTTP 1.1 connections (connection pools)
-- Add a default socket timeout for thrift RPCs
-
-## 2.5.2 (2023-05-08)
-
-- Fix: SQLAlchemy adapter could not reflect TIMESTAMP or DATETIME columns
-- Other: Relax pandas and alembic dependency specifications
-
-## 2.5.1 (2023-04-28)
-
-- Other: Relax sqlalchemy required version as it was unecessarily strict.
-
-## 2.5.0 (2023-04-14)
-
-- Add support for External Auth providers
-- Fix: Python HTTP proxies were broken
-- Other: All Thrift requests that timeout during connection will be automatically retried
-
-## 2.4.1 (2023-03-21)
-
-- Less strict numpy and pyarrow dependencies
-- Update examples in README to use security best practices
-- Update docstring for client.execute() for clarity
-
-## 2.4.0 (2023-02-21)
-
-- Improve compatibility when installed alongside other Databricks namespace Python packages
-- Add SQLAlchemy dialect
-
-## 2.3.0 (2023-01-10)
-
-- Support staging ingestion commands for DBR 12+
-
-## 2.2.2 (2023-01-03)
-
-- Support custom oauth client id and redirect port
-- Fix: Add none check on \_oauth_persistence in DatabricksOAuthProvider
-
-## 2.2.1 (2022-11-29)
-
-- Add support for Python 3.11
-
-## 2.2.0 (2022-11-15)
-
-- Bump thrift version to address https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2020-13949
-- Add support for lz4 compression
-
-## 2.1.0 (2022-09-30)
-
-- Introduce experimental OAuth support while Bring Your Own IDP is in Public Preview on AWS
-- Add functional examples
-
-## 2.0.5 (2022-08-23)
-
-- Fix: closing a connection now closes any open cursors from that connection at the server
-- Other: Add project links to pyproject.toml (helpful for visitors from PyPi)
-
-## 2.0.4 (2022-08-17)
-
-- Add support for Python 3.10
-- Add unit test matrix for supported Python versions
-
-Huge thanks to @dbaxa for contributing this change!
-
-## 2.0.3 (2022-08-05)
-
-- Add retry logic for `GetOperationStatus` requests that fail with an `OSError`
-- Reorganised code to use Poetry for dependency management.
-
-## 2.0.2 (2022-05-04)
-
-- Better exception handling in automatic connection close
-
-## 2.0.1 (2022-04-21)
-
-- Fixed Pandas dependency in setup.cfg to be >= 1.2.0
-
-## 2.0.0 (2022-04-19)
-
-- Initial stable release of V2
-- Added better support for complex types, so that in Databricks runtime 10.3+, Arrays, Maps and Structs will get
-  deserialized as lists, lists of tuples and dicts, respectively.
-- Changed the name of the metadata arg to http_headers
-
-## 2.0.b2 (2022-04-04)
-
-- Change import of collections.Iterable to collections.abc.Iterable to make the library compatible with Python 3.10
-- Fixed bug with .tables method so that .tables works as expected with Unity-Catalog enabled endpoints
-
-## 2.0.0b1 (2022-03-04)
-
-- Fix packaging issue (dependencies were not being installed properly)
-- Fetching timestamp results will now return aware instead of naive timestamps
-- The client will now default to using simplified error messages
-
-## 2.0.0b (2022-02-08)
-
-- Initial beta release of V2. V2 is an internal re-write of large parts of the connector to use Databricks edge features. All public APIs from V1 remain.
-- Added Unity Catalog support (pass catalog and / or schema key word args to the .connect method to select initial schema and catalog)
-
----
-
-**Note**: The code for versions prior to `v2.0.0b` is not contained in this repository. The below entries are included for reference only.
-
----
-
-## 1.0.0 (2022-01-20)
-
-- Add operations for retrieving metadata
-- Add the ability to access columns by name on result rows
-- Add the ability to provide configuration settings on connect
-
-## 0.9.4 (2022-01-10)
-
-- Improved logging and error messages.
-
-## 0.9.3 (2021-12-08)
-
-- Add retries for 429 and 503 HTTP responses.
-
-## 0.9.2 (2021-12-02)
-
-- (Bug fix) Increased Thrift requirement from 0.10.0 to 0.13.0 as 0.10.0 was in fact incompatible
-- (Bug fix) Fixed error message after query execution failed -SQLSTATE and Error message were misplaced
-
-## 0.9.1 (2021-09-01)
-
-- Public Preview release, Experimental tag removed
-- minor updates in internal build/packaging
-- no functional changes
-
-## 0.9.0 (2021-08-04)
-
-- initial (Experimental) release of pyhive-forked connector
-- Python DBAPI 2.0 (PEP-0249), thrift based
-- see docs for more info: https://docs.databricks.com/dev-tools/python-sql-connector.html
+- This is databricks-sqlalchemy plugin based on sqlalchemy v1 and has all the databricks-sql-python v2.9.6 needed sqlalchemy features

databricks/sqlalchemy/__init__.py

@@ -1,4 +1 @@
-from databricks.sqlalchemy.base import DatabricksDialect
-from databricks.sqlalchemy._types import TINYINT, TIMESTAMP, TIMESTAMP_NTZ
-
-__all__ = ["TINYINT", "TIMESTAMP", "TIMESTAMP_NTZ"]
+from databricks.sqlalchemy.dialect import DatabricksDialect

databricks/sqlalchemy/dialect/__init__.py

@@ -0,0 +1,340 @@
+"""This module's layout loosely follows example of SQLAlchemy's postgres dialect
+"""
+
+import decimal, re, datetime
+from dateutil.parser import parse
+
+import sqlalchemy
+from sqlalchemy import types, event
+from sqlalchemy.engine import default, Engine
+from sqlalchemy.exc import DatabaseError, SQLAlchemyError
+from sqlalchemy.engine import reflection
+
+from databricks import sql
+
+
+from databricks.sqlalchemy.dialect.base import (
+    DatabricksDDLCompiler,
+    DatabricksIdentifierPreparer,
+)
+from databricks.sqlalchemy.dialect.compiler import DatabricksTypeCompiler
+
+try:
+    import alembic
+except ImportError:
+    pass
+else:
+    from alembic.ddl import DefaultImpl
+
+    class DatabricksImpl(DefaultImpl):
+        __dialect__ = "databricks"
+
+
+class DatabricksDecimal(types.TypeDecorator):
+    """Translates strings to decimals"""
+
+    impl = types.DECIMAL
+
+    def process_result_value(self, value, dialect):
+        if value is not None:
+            return decimal.Decimal(value)
+        else:
+            return None
+
+
+class DatabricksTimestamp(types.TypeDecorator):
+    """Translates timestamp strings to datetime objects"""
+
+    impl = types.TIMESTAMP
+
+    def process_result_value(self, value, dialect):
+        return value
+
+    def adapt(self, impltype, **kwargs):
+        return self.impl
+
+
+class DatabricksDate(types.TypeDecorator):
+    """Translates date strings to date objects"""
+
+    impl = types.DATE
+
+    def process_result_value(self, value, dialect):
+        return value
+
+    def adapt(self, impltype, **kwargs):
+        return self.impl
+
+
+class DatabricksDialect(default.DefaultDialect):
+    """This dialect implements only those methods required to pass our e2e tests"""
+
+    # Possible attributes are defined here: https://docs.sqlalchemy.org/en/14/core/internals.html#sqlalchemy.engine.Dialect
+    name: str = "databricks"
+    driver: str = "databricks-sql-python"
+    default_schema_name: str = "default"
+
+    preparer = DatabricksIdentifierPreparer  # type: ignore
+    type_compiler = DatabricksTypeCompiler
+    ddl_compiler = DatabricksDDLCompiler
+    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
+
+    @classmethod
+    def dbapi(cls):
+        return sql
+
+    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"),
+        }
+
+        self.schema = kwargs["schema"]
+        self.catalog = kwargs["catalog"]
+
+        return [], kwargs
+
+    def get_columns(self, connection, table_name, schema=None, **kwargs):
+        """Return information about columns in `table_name`.
+
+        Given a :class:`_engine.Connection`, a string
+        `table_name`, and an optional string `schema`, return column
+        information as a list of dictionaries with these keys:
+
+        name
+          the column's name
+
+        type
+          [sqlalchemy.types#TypeEngine]
+
+        nullable
+          boolean
+
+        default
+          the column's default value
+
+        autoincrement
+          boolean
+
+        sequence
+          a dictionary of the form
+              {'name' : str, 'start' :int, 'increment': int, 'minvalue': int,
+               'maxvalue': int, 'nominvalue': bool, 'nomaxvalue': bool,
+               'cycle': bool, 'cache': int, 'order': bool}
+
+        Additional column attributes may be present.
+        """
+
+        _type_map = {
+            "boolean": types.Boolean,
+            "smallint": types.SmallInteger,
+            "int": types.Integer,
+            "bigint": types.BigInteger,
+            "float": types.Float,
+            "double": types.Float,
+            "string": types.String,
+            "varchar": types.String,
+            "char": types.String,
+            "binary": types.String,
+            "array": types.String,
+            "map": types.String,
+            "struct": types.String,
+            "uniontype": types.String,
+            "decimal": DatabricksDecimal,
+            "timestamp": DatabricksTimestamp,
+            "date": DatabricksDate,
+        }
+
+        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()
+
+        columns = []
+
+        for col in resp:
+
+            # Taken from PyHive. This removes added type info from decimals and maps
+            _col_type = re.search(r"^\w+", col.TYPE_NAME).group(0)
+            this_column = {
+                "name": col.COLUMN_NAME,
+                "type": _type_map[_col_type.lower()],
+                "nullable": bool(col.NULLABLE),
+                "default": col.COLUMN_DEF,
+                "autoincrement": False if col.IS_AUTO_INCREMENT == "NO" else True,
+            }
+            columns.append(this_column)
+
+        return columns
+
+    def get_pk_constraint(self, connection, table_name, schema=None, **kw):
+        """Return information about the primary key constraint on
+        table_name`.
+
+        Given a :class:`_engine.Connection`, a string
+        `table_name`, and an optional string `schema`, return primary
+        key information as a dictionary with these keys:
+
+        constrained_columns
+          a list of column names that make up the primary key
+
+        name
+          optional name of the primary key constraint.
+
+        """
+        # TODO: implement this behaviour
+        return {"constrained_columns": []}
+
+    def get_foreign_keys(self, connection, table_name, schema=None, **kw):
+        """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
+
+        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
+        """
+        # TODO: Implement this behaviour
+        return []
+
+    def get_indexes(self, connection, table_name, schema=None, **kw):
+        """Return information about indexes in `table_name`.
+
+        Given a :class:`_engine.Connection`, a string
+        `table_name` and an optional string `schema`, return index
+        information as a list of dictionaries with these keys:
+
+        name
+          the index's name
+
+        column_names
+          list of column names in order
+
+        unique
+          boolean
+        """
+        # TODO: Implement this behaviour
+        return []
+
+    def get_table_names(self, connection, schema=None, **kwargs):
+        TABLE_NAME = 1
+        with self.get_connection_cursor(connection) as cur:
+            sql_str = "SHOW TABLES FROM {}".format(
+                ".".join([self.catalog, schema or self.schema])
+            )
+            data = cur.execute(sql_str).fetchall()
+            _tables = [i[TABLE_NAME] for i in data]
+
+        return _tables
+
+    def get_view_names(self, connection, schema=None, **kwargs):
+        VIEW_NAME = 1
+        with self.get_connection_cursor(connection) as cur:
+            sql_str = "SHOW VIEWS FROM {}".format(
+                ".".join([self.catalog, schema or self.schema])
+            )
+            data = cur.execute(sql_str).fetchall()
+            _tables = [i[VIEW_NAME] for i in data]
+
+        return _tables
+
+    def do_rollback(self, dbapi_connection):
+        # Databricks SQL Does not support transactions
+        pass
+
+    def has_table(
+        self, connection, table_name, schema=None, catalog=None, **kwargs
+    ) -> bool:
+        """SQLAlchemy docstrings say dialect providers must implement this method"""
+
+        _schema = schema or self.schema
+        _catalog = catalog or self.catalog
+
+        # DBR >12.x uses underscores in error messages
+        DBR_LTE_12_NOT_FOUND_STRING = "Table or view not found"
+        DBR_GT_12_NOT_FOUND_STRING = "TABLE_OR_VIEW_NOT_FOUND"
+
+        try:
+            res = connection.execute(
+                f"DESCRIBE TABLE {_catalog}.{_schema}.{table_name}"
+            )
+            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
+
+    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):
+        # 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")]
+
+
+@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
+
+    if "_user_agent_entry" in cparams:
+        new_user_agent = f"sqlalchemy + {cparams['_user_agent_entry']}"
+    else:
+        new_user_agent = "sqlalchemy"
+
+    cparams["_user_agent_entry"] = new_user_agent
+
+    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/dialect/base.py

@@ -0,0 +1,17 @@
+import re
+from sqlalchemy.sql import compiler
+
+
+class DatabricksIdentifierPreparer(compiler.IdentifierPreparer):
+    # SparkSQL identifier specification:
+    # ref: https://spark.apache.org/docs/latest/sql-ref-identifier.html
+
+    legal_characters = re.compile(r"^[A-Z0-9_]+$", re.I)
+
+    def __init__(self, dialect):
+        super().__init__(dialect, initial_quote="`")
+
+
+class DatabricksDDLCompiler(compiler.DDLCompiler):
+    def post_create_table(self, table):
+        return " USING DELTA"

databricks/sqlalchemy/dialect/compiler.py

@@ -0,0 +1,38 @@
+from sqlalchemy.sql import compiler
+
+
+class DatabricksTypeCompiler(compiler.GenericTypeCompiler):
+    """Originally forked from pyhive"""
+
+    def visit_INTEGER(self, type_):
+        return "INT"
+
+    def visit_NUMERIC(self, type_):
+        return "DECIMAL"
+
+    def visit_CHAR(self, type_):
+        return "STRING"
+
+    def visit_VARCHAR(self, type_):
+        return "STRING"
+
+    def visit_NCHAR(self, type_):
+        return "STRING"
+
+    def visit_TEXT(self, type_):
+        return "STRING"
+
+    def visit_CLOB(self, type_):
+        return "STRING"
+
+    def visit_BLOB(self, type_):
+        return "BINARY"
+
+    def visit_TIME(self, type_):
+        return "TIMESTAMP"
+
+    def visit_DATE(self, type_):
+        return "DATE"
+
+    def visit_DATETIME(self, type_):
+        return "TIMESTAMP"