DRAFT: This scaffolds in the code for running parallel dialect compli…

…ance

tests. Setting this up isn't documented anywhere and the implementation here
doesn't actually work because the test runner keeps trying to create
tables within the catalog.schema in the connection string.

Need to research more about why this happens. But wanted to stash this code
for now as we'll want to address this for faster compliance test runs
in the future.

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

src/databricks/sqlalchemy/__init__.py

@@ -282,6 +282,13 @@ def get_schema_names(self, connection, **kw):
 
         # TODO: replace with call to cursor.schemas() once its performance matches raw SQL
         return [row[0] for row in connection.execute("SHOW SCHEMAS")]
+
+    @classmethod
+    def load_provisioning(cls):
+        try:
+            __import__("databricks.sqlalchemy.provision")
+        except ImportError:
+            pass
 
 
 @event.listens_for(Engine, "do_connect")

src/databricks/sqlalchemy/provision.py

@@ -0,0 +1,12 @@
+from sqlalchemy.testing.provision import create_db, drop_db
+
+@create_db.for_db("databricks")
+def _databricks_create_db(cfg, eng, ident):
+    with eng.begin() as conn:
+        create_string = "CREATE SCHEMA `main`.`%s`" % ident
+        conn.exec_driver_sql(create_string)
+
+@drop_db.for_db("databricks")
+def _databricks_drop_db(cfg, eng, ident):
+    with eng.begin() as conn:
+        conn.exec_driver_sql("DROP SCHEMA `main`.`%s`" % ident)

src/databricks/sqlalchemy/setup.cfg

@@ -2,3 +2,6 @@
 [sqla_testing]
 requirement_cls=databricks.sqlalchemy.requirements:Requirements
 profile_file=profiles.txt
+
+[db]
+databricks=<enter valid connection string here>