feat: adding drop table if exists and alembic render import

sede-open · Oct 4, 2024 · 60c408a · 60c408a
1 parent c97fbc9
commit 60c408a
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/src/databricks/sqlalchemy/_ddl.py b/src/databricks/sqlalchemy/_ddl.py
@@ -101,6 +101,16 @@ def visit_create_table(self, create, **kw):
 
         return text
 
+    def liquid_cluster_on_table(self, liquid_cluster_columns):
+        columns = liquid_cluster_columns
+
+        return """CLUSTER BY ({cols})""".format(cols=', '.join(columns))
+
+    def visit_drop_table(self, drop, **kw):
+        text = "\nDROP TABLE IF EXISTS "
+
+        return text + self.preparer.format_table(drop.element)
+
     def visit_unique_constraint(self, constraint, **kw):
         logger.warning("Databricks does not support unique constraints")
         pass

diff --git a/src/databricks/sqlalchemy/base.py b/src/databricks/sqlalchemy/base.py
@@ -32,6 +32,7 @@
     pass
 else:
     from alembic.ddl import DefaultImpl
+    from alembic.autogenerate import render
 
     class DatabricksImpl(DefaultImpl):
         __dialect__ = "databricks"