Add SetColumnComment and DropColumnComment DDL elements

Signed-off-by: Christophe Bornet <[email protected]>

Author: cbornet

src/databricks/sqlalchemy/_ddl.py

@@ -39,6 +39,21 @@ def visit_identity_column(self, identity, **kw):
         )
         return text
 
+    def visit_set_column_comment(self, create, **kw):
+        return "ALTER TABLE %s ALTER COLUMN %s COMMENT %s" % (
+            self.preparer.format_table(create.element.table),
+            self.preparer.format_column(create.element),
+            self.sql_compiler.render_literal_value(
+                create.element.comment, sqltypes.String()
+            ),
+        )
+
+    def visit_drop_column_comment(self, create, **kw):
+        return "ALTER TABLE %s ALTER COLUMN %s COMMENT ''" % (
+            self.preparer.format_table(create.element.table),
+            self.preparer.format_column(create.element),
+        )
+
     def get_column_specification(self, column, **kwargs):
         # Emit a log message if a user attempts to set autoincrement=True on a column.
         # See comments in test_suite.py. We may implement implicit IDENTITY using this

src/databricks/sqlalchemy/test_local/e2e/test_basic.py

@@ -18,6 +18,7 @@
 from sqlalchemy.engine import Engine
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column
+from sqlalchemy.schema import DropColumnComment, SetColumnComment
 from sqlalchemy.types import BOOLEAN, DECIMAL, Date, DateTime, Integer, String
 
 try:
@@ -191,19 +192,34 @@ def test_create_table_not_null(db_engine, metadata_obj: MetaData):
 def test_column_comment(db_engine, metadata_obj: MetaData):
     table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
 
-    SampleTable = Table(
-        table_name,
-        metadata_obj,
-        Column("name", String(255), comment="some comment")
-    )
+    column = Column("name", String(255), comment="some comment")
+    SampleTable = Table(table_name, metadata_obj, column)
 
     metadata_obj.create_all(db_engine)
+    connection = db_engine.connect()
 
     columns = db_engine.dialect.get_columns(
-        connection=db_engine.connect(), table_name=table_name
+        connection=connection, table_name=table_name
+    )
+
+    assert columns[0].get("comment") == "some comment"
+
+    column.comment = "other comment"
+    connection.execute(SetColumnComment(column))
+
+    columns = db_engine.dialect.get_columns(
+        connection=connection, table_name=table_name
+    )
+
+    assert columns[0].get("comment") == "other comment"
+
+    connection.execute(DropColumnComment(column))
+
+    columns = db_engine.dialect.get_columns(
+        connection=connection, table_name=table_name
     )
 
-    assert columns[0].get('comment') == "some comment"
+    assert columns[0].get("comment") == ""
 
     metadata_obj.drop_all(db_engine)