Add SQLAlchemy Support

pyproject.toml

@@ -14,10 +14,14 @@ thrift = "^0.13.0"
 pyarrow = "^5.0.0"
 pandas = "^1.3.0"
 
+[tool.poetry.plugins."sqlalchemy.dialects"]
+"databricks.thrift" = "databricks.sqlalchemy:DatabricksDialect"
+
 [tool.poetry.dev-dependencies]
 pytest = "^7.1.2"
 mypy = "^0.950"
 black = "^22.3.0"
+SQLAlchemy = "^1.4.39"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]
@@ -28,4 +32,5 @@ ignore_missing_imports = "true"
 exclude = ['ttypes\.py$', 'TCLIService\.py$']
 
 [tool.black]
-exclude = '/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck-out|build|dist|thrift_api)/'
+exclude = '/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck-out|build|dist|thrift_api)/'
+

setup.cfg

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

src/databricks/sqlalchemy/__init__.py

@@ -0,0 +1 @@
+from databricks.sqlalchemy.dialect import DatabricksDialect

src/databricks/sqlalchemy/dialect.py

@@ -0,0 +1,59 @@
+from databricks import sql
+from typing import AnyStr
+
+from sqlalchemy import util, exc, types
+from sqlalchemy.engine import default
+
+
+class DatabricksDialect(default.DefaultDialect):
+
+    # Possible attributes are defined here: https://docs.sqlalchemy.org/en/14/core/internals.html#sqlalchemy.engine.Dialect
+    name: str = "databricks"
+    driver: str = "thrift"
+    default_schema_name: str = "default"
+
+    @classmethod
+    def dbapi(cls):
+        return sql
+
+    def create_connect_args(self, url):
+        # 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"),
+        }
+
+        return [], kwargs
+
+    def get_table_names(self, *args, **kwargs):
+
+        # TODO: Implement with native driver `.tables()` call
+        return super().get_table_names(*args, **kwargs)
+
+    def get_columns(self, *args, **kwargs):
+
+        # TODO: Implement with native driver `.columns()` call
+
+        return super().get_columns(*args, **kwargs)
+
+    def do_rollback(self, dbapi_connection):
+        # Databricks SQL Does not support transaction
+        pass
+
+    def has_table(self, connection, table_name, schema=None, **kwargs) -> bool:
+        """Required for `tests.sqlalchemy.integration.test_create_table` to pass.
+        """
+        try:
+            COLUMN_NAME = 3
+            with self.get_driver_connection(
+                connection
+            )._dbapi_connection.dbapi_connection.cursor() as cur:
+                data = cur.columns(
+                    schema_name=schema or "default", table_name=table_name
+                ).fetchmany(1)
+            # the table exists as long as there's a non-zero number of columns
+            return len(data) > 0
+        except exc.NoSuchTableError:
+            return False

src/databricks/sqlalchemy/requirements.py

@@ -0,0 +1,18 @@
+# Following official SQLAlchemy guide:
+#
+# https://github.com/sqlalchemy/sqlalchemy/blob/main/README.dialects.rst#dialect-layout
+#
+# The full group of requirements is available here:
+#
+# https://github.com/sqlalchemy/sqlalchemy/blob/a453256afc334acabee25ec275de555ef7287144/test/requirements.py
+
+
+from sqlalchemy.testing.requirements import SuiteRequirements
+from sqlalchemy.testing import exclusions
+
+class Requirements(SuiteRequirements):
+
+    @property
+    def two_phase_transactions(self):
+        # Databricks SQL doesn't support transactions
+        return exclusions.closed()

tests/sqlalchemy/integration.py

@@ -0,0 +1,46 @@
+import os, datetime, pytest
+import sqlalchemy
+
+
+@pytest.fixture
+def db_engine():
+
+    host = os.getenv("host")
+    http_path = os.getenv("http_path")
+    api_token = os.getenv("api_token")
+    engine = sqlalchemy.create_engine(f"databricks+thrift://token:{api_token}@{host}?http_path={http_path}")
+    return engine
+
+def test_basic_connection(db_engine):
+    """Make sure we can connect and run basic query
+    """
+
+    curs = db_engine.execute("SELECT id FROM RANGE(100)")
+    result = curs.fetchall()
+    assert len(result) == 100
+
+def test_create_and_drop_table(db_engine):
+    """Make sure we can automatically create and drop a table defined with SQLAlchemy's MetaData object
+    """
+
+    mdo = sqlalchemy.MetaData()
+    this_moment = datetime.datetime.utcnow().strftime("%s")
+
+    tname = f"integration_test_table_{this_moment}"
+
+    t1 = sqlalchemy.Table(
+        tname,
+        mdo,
+        sqlalchemy.Column('f_short', sqlalchemy.types.SMALLINT),
+        sqlalchemy.Column('f_int', sqlalchemy.types.Integer),
+        sqlalchemy.Column('f_long', sqlalchemy.types.BigInteger),
+        sqlalchemy.Column('f_float', sqlalchemy.types.Float),
+        sqlalchemy.Column('f_decimal', sqlalchemy.types.DECIMAL),
+        sqlalchemy.Column('f_boolean', sqlalchemy.types.BOOLEAN)
+    )
+
+    mdo.create_all(bind=db_engine,checkfirst=True)
+
+    check_it_exists = db_engine.execute(f"DESCRIBE TABLE EXTENDED {tname}")
+
+    mdo.drop_all(db_engine, checkfirst=True)