Add support for column comments

src/databricks/sqlalchemy/_ddl.py

@@ -1,5 +1,5 @@
 import re
-from sqlalchemy.sql import compiler
+from sqlalchemy.sql import compiler, sqltypes
 import logging
 
 logger = logging.getLogger(__name__)
@@ -39,17 +39,38 @@ 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):
-        """Currently we override this method only to 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 feature in the future, similar to the Microsoft SQL Server dialect.
-        """
+        # 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
+        # feature in the future, similar to the Microsoft SQL Server dialect.
         if column is column.table._autoincrement_column or column.autoincrement is True:
             logger.warn(
                 "Databricks dialect ignores SQLAlchemy's autoincrement semantics. Use explicit Identity() instead."
             )
 
-        return super().get_column_specification(column, **kwargs)
+        colspec = super().get_column_specification(column, **kwargs)
+        if column.comment is not None:
+            literal = self.sql_compiler.render_literal_value(
+                column.comment, sqltypes.STRINGTYPE
+            )
+            colspec += " COMMENT " + literal
+
+        return colspec
 
 
 class DatabricksStatementCompiler(compiler.SQLCompiler):

src/databricks/sqlalchemy/_parse.py

@@ -354,6 +354,7 @@ def parse_column_info_from_tgetcolumnsresponse(thrift_resp_row) -> ReflectedColu
         "type": final_col_type,
         "nullable": bool(thrift_resp_row.NULLABLE),
         "default": thrift_resp_row.COLUMN_DEF,
+        "comment": thrift_resp_row.REMARKS,
     }
 
     # TODO: figure out how to return sqlalchemy.interfaces in a way that mypy respects

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:
@@ -188,6 +189,41 @@ def test_create_table_not_null(db_engine, metadata_obj: MetaData):
     metadata_obj.drop_all(db_engine)
 
 
+def test_column_comment(db_engine, metadata_obj: MetaData):
+    table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
+
+    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=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") == ""
+
+    metadata_obj.drop_all(db_engine)
+
+
 def test_bulk_insert_with_core(db_engine, metadata_obj, session):
     import random