Merge pull request #6 from sede-open/task/add-column-comment-code

PeterDKay · web-flow · commit 9656f9506267 · 2023-08-04T18:29:10.000+01:00
First test with new compiler function
diff --git a/src/databricks/sqlalchemy/dialect/__init__.py b/src/databricks/sqlalchemy/dialect/__init__.py
@@ -245,6 +245,7 @@ def get_indexes(self, connection, table_name, schema=None, **kw):
         return []
 
     def get_table_names(self, connection, schema=None, **kwargs):
+        # TODO: Change this to ignore views!!!!!!!
         TABLE_NAME = 1
         catalog = "`" + self.catalog + "`"
         schema = ("`" + schema + "`") or ("`" + self.schema + "`")
diff --git a/src/databricks/sqlalchemy/dialect/base.py b/src/databricks/sqlalchemy/dialect/base.py
@@ -1,6 +1,8 @@
 import re
+from alembic.ddl.base import ColumnComment
 from sqlalchemy import util, exc
 from sqlalchemy.sql import compiler, sqltypes, ColumnElement
+from sqlalchemy.ext.compiler import compiles
 
 
 class DatabricksIdentifierPreparer(compiler.IdentifierPreparer):
@@ -123,3 +125,25 @@ def visit_drop_table(self, drop, **kw):
         text = "\nDROP TABLE IF EXISTS "
 
         return text + self.preparer.format_table(drop.element)
+
+
+@compiles(ColumnComment, "databricks")
+def visit_column_comment(
+    element: ColumnComment, compiler: DatabricksDDLCompiler, **kw
+) -> str:
+    ddl = "ALTER TABLE `{schema}`.{table_name} ALTER COLUMN {column_name} COMMENT {comment}"
+    comment = (
+        compiler.sql_compiler.render_literal_value(
+            element.comment, sqltypes.String()
+        )
+        if element.comment is not None
+        else "NULL"
+    )
+
+    return ddl.format(
+        schema=element.schema,
+        table_name=element.table_name,
+        column_name=element.column_name,
+        comment=comment,
+    )
+
