Basic metadata operation now supported. This requires a non HMS catalog.

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

Author: Jesse Whitehouse

src/databricks/sqlalchemy/dialect/__init__.py

@@ -1,16 +1,25 @@
+"""This module's layout loosely follows example of SQLAlchemy's postgres dialect
+"""
+
 from sqlalchemy.engine import default
 from sqlalchemy.exc import DatabaseError
 
 from databricks import sql
 
 
+from databricks.sqlalchemy.dialect.base import DatabricksIdentifierPreparer
+
 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 = "thrift"
     default_schema_name: str = "default"
 
+    preparer = DatabricksIdentifierPreparer
+
     @classmethod
     def dbapi(cls):
         return sql
@@ -23,6 +32,8 @@ def create_connect_args(self, url):
             "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")
         }
 
         return [], kwargs

src/databricks/sqlalchemy/dialect/base.py

@@ -0,0 +1,12 @@
+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="`")

tests/e2e/sqlalchemy/test_queries.py

@@ -1,9 +1,9 @@
 import os, datetime
 import pytest
 from unittest import skipIf
-from sqlalchemy import create_engine, select, Column
+from sqlalchemy import create_engine, select, insert, Column, MetaData, Table
 from sqlalchemy.orm import declarative_base, Session
-from sqlalchemy.types import SMALLINT, Integer, BigInteger, Float, DECIMAL, BOOLEAN, String
+from sqlalchemy.types import SMALLINT, Integer, BOOLEAN, String
 
 
 @pytest.fixture
@@ -12,8 +12,10 @@ def db_engine():
     HOST = os.environ.get("host")
     HTTP_PATH = os.environ.get("http_path")
     ACCESS_TOKEN = os.environ.get("access_token")
+    CATALOG = os.environ.get("catalog")
+    SCHEMA = os.environ.get("schema")
 
-    engine = create_engine(f"databricks+thrift://token:{ACCESS_TOKEN}@{HOST}?http_path={HTTP_PATH}")
+    engine = create_engine(f"databricks+thrift://token:{ACCESS_TOKEN}@{HOST}?http_path={HTTP_PATH}&catalog={CATALOG}&schema={SCHEMA}")
     return engine
 
 
@@ -26,23 +28,57 @@ def base(db_engine):
 def session(db_engine):
     return Session(bind=db_engine)
 
+@pytest.fixture()
+def metadata_obj(db_engine):
+    return MetaData(bind=db_engine)
+
 
 def test_can_connect(db_engine):
     simple_query = "SELECT 1"
     result = db_engine.execute(simple_query).fetchall()
     assert len(result) == 1
 
-@skipIf(True, 'metadata operations not yet supported')
-def test_create_insert_drop_table(base, session: Session):
-    """Make sure we can automatically create and drop a table defined with SQLAlchemy's MetaData object
-    while exercising all supported types.
+
+def test_create_insert_drop_table_core(base, db_engine, metadata_obj: MetaData):
+    """
+    """
+
+    SampleTable = Table(
+         "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s")),
+         metadata_obj,
+         Column("name", String(255)),
+         Column("episodes", Integer),
+         Column("some_bool", BOOLEAN)
+    )
+
+    metadata_obj.create_all()
+
+    insert_stmt = insert(SampleTable).values(name="Bim Adwunmi", episodes=6, some_bool=True)
+
+    with db_engine.connect() as conn:
+        conn.execute(insert_stmt)
+
+    select_stmt = select(SampleTable)
+    resp = db_engine.execute(select_stmt)
+
+    result = resp.fetchall()
+
+    assert len(result) == 1
+
+    metadata_obj.drop_all()
+
+
+@skipIf(True, 'Unity catalog must be supported')
+def test_create_insert_drop_table_orm(base, session: Session):
+    """ORM classes built on the declarative base class must have a primary key. 
+    This is restricted to Unity Catalog.
     """
 
     class SampleObject(base):
 
         __tablename__ = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
 
-        name = Column(String, primary_key=True)
+        name = Column(String(255), primary_key=True)
         episodes = Column(Integer),
         some_bool = Column(BOOLEAN)