Scaffold basic file structure for a new dialect

https://github.com/sqlalchemy/sqlalchemy/blob/main/README.dialects.rst

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: susodapop

poetry.lock

@@ -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)/'
+

pyproject.toml

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

setup.cfg

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

src/databricks/sqlalchemy/__init__.py

@@ -0,0 +1,45 @@
+from databricks import sql
+from typing import AnyStr
+
+from sqlalchemy import types
+from sqlalchemy import util
+
+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

src/databricks/sqlalchemy/dialect.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()

src/databricks/sqlalchemy/requirements.py

@@ -0,0 +1,7 @@
+from sqlalchemy.dialects import registry
+import pytest
+
+registry.register("databricks.thrift", "databricks.sqlalchemy", "DatabricksDialect")
+pytest.register_assert_rewrite("sqlalchemy.testing.assertions")
+
+from sqlalchemy.testing.plugin.pytestplugin import *

tests/sqlalchemy/conftest.py

@@ -0,0 +1,2 @@
+from sqlalchemy.testing.suite import *
+