Updates to dialect to handle comments in alembic

PeterDKay · PeterDKay · commit e324ff90f724 · 2023-05-11T14:21:39.000+01:00
diff --git a/src/databricks/sqlalchemy/dialect/__init__.py b/src/databricks/sqlalchemy/dialect/__init__.py
@@ -80,6 +80,7 @@ class DatabricksDialect(default.DefaultDialect):
     supports_multivalues_insert: bool = True
     supports_native_decimal: bool = True
     supports_sane_rowcount: bool = False
+    supports_comments: bool = True
 
     @classmethod
     def dbapi(cls):
@@ -174,6 +175,7 @@ def get_columns(self, connection, table_name, schema=None, **kwargs):
                 "nullable": bool(col.NULLABLE),
                 "default": col.COLUMN_DEF,
                 "autoincrement": False if col.IS_AUTO_INCREMENT == "NO" else True,
+                'comment': col.REMARKS,
             }
             columns.append(this_column)
 
diff --git a/src/databricks/sqlalchemy/dialect/base.py b/src/databricks/sqlalchemy/dialect/base.py
@@ -1,5 +1,5 @@
 import re
-from sqlalchemy.sql import compiler
+from sqlalchemy.sql import compiler, sqltypes, ColumnElement
 
 
 class DatabricksIdentifierPreparer(compiler.IdentifierPreparer):
@@ -15,3 +15,52 @@ def __init__(self, dialect):
 class DatabricksDDLCompiler(compiler.DDLCompiler):
     def post_create_table(self, table):
         return " USING DELTA"
+    def visit_set_column_comment(self, create, **kw):
+        """
+        Example syntax for adding column comment:
+        "ALTER TABLE schema.table_name CHANGE COLUMN COLUMN_NAME COMMENT 'Comment to be added to column';"
+
+        """
+        return """ALTER TABLE {0} CHANGE COLUMN {1} COMMENT {2}""".format(
+            'testing.alembic_test1',
+            self._format_table_from_column(
+                 create, use_schema=True
+            ),
+            self.preparer.format_column(
+                create.element, use_table=False
+            ),
+            self.sql_compiler.render_literal_value(
+                create.element.comment, sqltypes.String()
+            ),
+        )
+
+    def visit_drop_column_comment(self, drop, **kw):
+        """
+        Example syntax for dropping column comment:
+        "ALTER TABLE schema.table_name CHANGE COLUMN COLUMN_NAME COMMENT '';"
+
+        Note: There is no syntactical 'DROP' statement in this case, the comment must be replaced with an empty string
+        """
+        return "ALTER TABLE {0} CHANGE COLUMN {1} COMMENT '';".format(
+            self._format_table_from_column(
+                 drop, use_schema=True
+            ),
+            self.preparer.format_column(
+                drop.element, use_table=False
+            )
+        )
+
+    def _format_table_from_column(self, column_object, use_schema=False):
+        """
+        Prepare a quoted table name from the column object (including schema if specified)
+        """
+        schema_table_column = self.preparer.format_column(
+                column_object.element, use_table=True, use_schema=True
+            )
+
+        name = schema_table_column.split(".")[1]
+
+        if use_schema:
+            name = schema_table_column.split(".")[0] + '.' + name
+
+        return name
