Add support for column comments

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

Author: cbornet

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__)
@@ -40,16 +40,22 @@ def visit_identity_column(self, identity, **kw):
         return text
 
     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

@@ -188,6 +188,26 @@ 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"))
+
+    SampleTable = Table(
+        table_name,
+        metadata_obj,
+        Column("name", String(255), comment="some comment")
+    )
+
+    metadata_obj.create_all(db_engine)
+
+    columns = db_engine.dialect.get_columns(
+        connection=db_engine.connect(), table_name=table_name
+    )
+
+    assert columns[0].get('comment') == "some comment"
+
+    metadata_obj.drop_all(db_engine)
+
+
 def test_bulk_insert_with_core(db_engine, metadata_obj, session):
     import random