Add support for table comments

src/databricks/sqlalchemy/_ddl.py

@@ -16,7 +16,13 @@ def __init__(self, dialect):
 
 class DatabricksDDLCompiler(compiler.DDLCompiler):
     def post_create_table(self, table):
-        return " USING DELTA"
+        post = " USING DELTA"
+        if table.comment:
+            comment = self.sql_compiler.render_literal_value(
+                table.comment, sqltypes.String()
+            )
+            post += " COMMENT " + comment
+        return post
 
     def visit_unique_constraint(self, constraint, **kw):
         logger.warning("Databricks does not support unique constraints")
@@ -61,7 +67,7 @@ def get_column_specification(self, column, **kwargs):
         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(
+            logger.warning(
                 "Databricks dialect ignores SQLAlchemy's autoincrement semantics. Use explicit Identity() instead."
             )
 

src/databricks/sqlalchemy/base.py

@@ -1,5 +1,4 @@
-import re
-from typing import Any, List, Optional, Dict, Union, Collection, Iterable, Tuple
+from typing import Any, List, Optional, Dict, Union
 
 import databricks.sqlalchemy._ddl as dialect_ddl_impl
 import databricks.sqlalchemy._types as dialect_type_impl
@@ -17,13 +16,13 @@
 import sqlalchemy
 from sqlalchemy import DDL, event
 from sqlalchemy.engine import Connection, Engine, default, reflection
-from sqlalchemy.engine.reflection import ObjectKind
 from sqlalchemy.engine.interfaces import (
     ReflectedForeignKeyConstraint,
     ReflectedPrimaryKeyConstraint,
     ReflectedColumn,
-    TableKey,
+    ReflectedTableComment,
 )
+from sqlalchemy.engine.reflection import ReflectionDefaults
 from sqlalchemy.exc import DatabaseError, SQLAlchemyError
 
 try:
@@ -285,7 +284,7 @@ def get_table_names(self, connection: Connection, schema=None, **kwargs):
         views_result = self.get_view_names(connection=connection, schema=schema)
 
         # In Databricks, SHOW TABLES FROM <schema> returns both tables and views.
-        # Potential optimisation: rewrite this to instead query informtation_schema
+        # Potential optimisation: rewrite this to instead query information_schema
         tables_minus_views = [
             row.tableName for row in tables_result if row.tableName not in views_result
         ]
@@ -328,7 +327,7 @@ def get_materialized_view_names(
     def get_temp_view_names(
         self, connection: Connection, schema: Optional[str] = None, **kw: Any
     ) -> List[str]:
-        """A wrapper around get_view_names taht fetches only the names of temporary views"""
+        """A wrapper around get_view_names that fetches only the names of temporary views"""
         return self.get_view_names(connection, schema, only_temp=True)
 
     def do_rollback(self, dbapi_connection):
@@ -375,6 +374,30 @@ def get_schema_names(self, connection, **kw):
         schema_list = [row[0] for row in result]
         return schema_list
 
+    @reflection.cache
+    def get_table_comment(
+        self,
+        connection: Connection,
+        table_name: str,
+        schema: Optional[str] = None,
+        **kw: Any,
+    ) -> ReflectedTableComment:
+        result = self._describe_table_extended(
+            connection=connection,
+            table_name=table_name,
+            schema_name=schema,
+        )
+        # Type ignore is because mypy knows that self._describe_table_extended *can*
+        # return None (even though it never will since expect_result defaults to True)
+        comment_row: Dict[str, str] = next(
+            filter(lambda r: r["col_name"] == "Comment", result), None
+        )  # type: ignore
+        return (
+            {"text": comment_row["data_type"]}
+            if comment_row
+            else ReflectionDefaults.table_comment()
+        )
+
 
 @event.listens_for(Engine, "do_connect")
 def receive_do_connect(dialect, conn_rec, cargs, cparams):

src/databricks/sqlalchemy/requirements.py

@@ -159,6 +159,11 @@ def table_reflection(self):
         """target database has general support for table reflection"""
         return sqlalchemy.testing.exclusions.open()
 
+    @property
+    def comment_reflection(self):
+        """Indicates if the database support table comment reflection"""
+        return sqlalchemy.testing.exclusions.open()
+
     @property
     def temp_table_reflection(self):
         """ComponentReflection test is intricate and simply cannot function without this exclusion being defined here.

src/databricks/sqlalchemy/test/_extra.py

@@ -3,6 +3,9 @@
 
 import datetime
 
+from sqlalchemy import Integer, String, schema, inspect
+from sqlalchemy.testing import util
+from sqlalchemy.testing.config import requirements
 from sqlalchemy.testing.suite.test_types import (
     _LiteralRoundTripFixture,
     fixtures,
@@ -63,8 +66,39 @@ class DateTimeTZTestCustom(_DateFixture, fixtures.TablesTest):
 
     @testing.requires.datetime_implicit_bound
     def test_select_direct(self, connection):
-
         # We need to pass the TIMESTAMP type to the literal function
         # so that the value is processed correctly.
         result = connection.scalar(select(literal(self.data, TIMESTAMP)))
         eq_(result, self.data)
+
+
+class TableDDLTestCustom(fixtures.TestBase):
+    """This test confirms that a table comment can be dropped.
+    The difference with TableDDLTest is that the comment value is '' and not None after
+    being dropped.
+    """
+
+    __backend__ = True
+
+    def _simple_fixture(self, schema=None):
+        return Table(
+            "test_table",
+            self.metadata,
+            Column("id", Integer, primary_key=True, autoincrement=False),
+            Column("data", String(50)),
+            schema=schema,
+        )
+
+    @requirements.comment_reflection
+    @util.provide_metadata
+    def test_drop_table_comment(self, connection):
+        table = self._simple_fixture()
+        table.create(connection, checkfirst=False)
+        table.comment = "a comment"
+        connection.execute(schema.SetTableComment(table))
+        connection.execute(schema.DropTableComment(table))
+        eq_(inspect(connection).get_table_comment("test_table"), {"text": ""})
+
+
+class FutureTableDDLTestCustom(fixtures.FutureEngineMixin, TableDDLTestCustom):
+    pass

src/databricks/sqlalchemy/test/_future.py

@@ -13,9 +13,7 @@
     ComponentReflectionTest,
     ComponentReflectionTestExtra,
     CTETest,
-    FutureTableDDLTest,
     InsertBehaviorTest,
-    TableDDLTest,
 )
 from sqlalchemy.testing.suite import (
     ArrayTest,
@@ -251,36 +249,7 @@ class FutureWeCanSetDefaultSchemaWEventsTest(FutureWeCanSetDefaultSchemaWEventsT
     pass
 
 
-class FutureTableDDLTest(FutureTableDDLTest):
-    @pytest.mark.skip(reason=render_future_feature(FutureFeature.TBL_COMMENTS))
-    def test_add_table_comment(self):
-        """We could use requirements.comment_reflection here to disable this but prefer a more meaningful skip message"""
-        pass
-
-    @pytest.mark.skip(reason=render_future_feature(FutureFeature.TBL_COMMENTS))
-    def test_drop_table_comment(self):
-        """We could use requirements.comment_reflection here to disable this but prefer a more meaningful skip message"""
-        pass
-
-
-class TableDDLTest(TableDDLTest):
-    @pytest.mark.skip(reason=render_future_feature(FutureFeature.TBL_COMMENTS))
-    def test_add_table_comment(self, connection):
-        """We could use requirements.comment_reflection here to disable this but prefer a more meaningful skip message"""
-        pass
-
-    @pytest.mark.skip(reason=render_future_feature(FutureFeature.TBL_COMMENTS))
-    def test_drop_table_comment(self, connection):
-        """We could use requirements.comment_reflection here to disable this but prefer a more meaningful skip message"""
-        pass
-
-
 class ComponentReflectionTest(ComponentReflectionTest):
-    @pytest.mark.skip(reason=render_future_feature(FutureFeature.TBL_COMMENTS))
-    def test_get_multi_table_comment(self):
-        """There are 84 permutations of this test that are skipped."""
-        pass
-
     @pytest.mark.skip(reason=render_future_feature(FutureFeature.TBL_OPTS, True))
     def test_multi_get_table_options_tables(self):
         """It's not clear what the expected ouput from this method would even _be_. Requires research."""
@@ -302,14 +271,6 @@ def test_get_multi_pk_constraint(self):
     def test_get_multi_check_constraints(self):
         pass
 
-    @pytest.mark.skip(reason=render_future_feature(FutureFeature.TBL_COMMENTS))
-    def test_get_comments(self):
-        pass
-
-    @pytest.mark.skip(reason=render_future_feature(FutureFeature.TBL_COMMENTS))
-    def test_get_comments_with_schema(self):
-        pass
-
     @pytest.mark.skip(reason=render_future_feature(FutureFeature.TBL_COMMENTS))
     def test_comments_unicode(self):
         pass

src/databricks/sqlalchemy/test/_unsupported.py

@@ -55,6 +55,7 @@ class SkipReason(Enum):
     TIMEZONE_OPT = "timezone-optional TIMESTAMP fields"
     TRANSACTIONS = "transactions"
     UNIQUE = "UNIQUE constraints"
+    DROP_TBL = "drop table comment"
 
 
 def render_skip_reason(rsn: SkipReason, setup_error=False, extra=False) -> str:
@@ -222,6 +223,13 @@ def test_uuid_returning(self):
 
 
 class FutureTableDDLTest(FutureTableDDLTest):
+    @pytest.mark.skip(reason=render_skip_reason(SkipReason.DROP_TBL))
+    def test_drop_table_comment(self, connection):
+        """The DropTableComment statement is supported but it sets the comment to ''
+        instead of None so this test can't pass.
+        """
+        pass
+
     @pytest.mark.skip(render_skip_reason(SkipReason.INDEXES))
     def test_create_index_if_not_exists(self):
         """We could use requirements.index_reflection and requirements.index_ddl_if_exists
@@ -238,6 +246,13 @@ def test_drop_index_if_exists(self):
 
 
 class TableDDLTest(TableDDLTest):
+    @pytest.mark.skip(reason=render_skip_reason(SkipReason.DROP_TBL))
+    def test_drop_table_comment(self, connection):
+        """The DropTableComment statement is supported but it sets the comment to ''
+        instead of None so this test can't pass.
+        """
+        pass
+
     @pytest.mark.skip(reason=render_skip_reason(SkipReason.INDEXES))
     def test_create_index_if_not_exists(self, connection):
         """We could use requirements.index_reflection and requirements.index_ddl_if_exists

src/databricks/sqlalchemy/test/test_suite.py

@@ -10,4 +10,9 @@
 from databricks.sqlalchemy.test._regression import *
 from databricks.sqlalchemy.test._unsupported import *
 from databricks.sqlalchemy.test._future import *
-from databricks.sqlalchemy.test._extra import TinyIntegerTest, DateTimeTZTestCustom
+from databricks.sqlalchemy.test._extra import (
+    TinyIntegerTest,
+    DateTimeTZTestCustom,
+    TableDDLTestCustom,
+    FutureTableDDLTestCustom
+)

src/databricks/sqlalchemy/test_local/test_ddl.py

@@ -1,16 +1,24 @@
 import pytest
 from sqlalchemy import Column, MetaData, String, Table, create_engine
-from sqlalchemy.schema import CreateTable, DropColumnComment, SetColumnComment
+from sqlalchemy.schema import (
+    CreateTable,
+    DropColumnComment,
+    DropTableComment,
+    SetColumnComment,
+    SetTableComment,
+)
 
 
-class TestTableCommentDDL:
+class DDLTestBase:
     engine = create_engine(
         "databricks://token:****@****?http_path=****&catalog=****&schema=****"
     )
 
     def compile(self, stmt):
         return str(stmt.compile(bind=self.engine))
 
+
+class TestColumnCommentDDL(DDLTestBase):
     @pytest.fixture
     def metadata(self) -> MetaData:
         """Assemble a metadata object with one table containing one column."""
@@ -45,3 +53,43 @@ def test_alter_table_drop_column_comment(self, column):
         stmt = DropColumnComment(column)
         output = self.compile(stmt)
         assert output == "ALTER TABLE foobar ALTER COLUMN foo COMMENT ''"
+
+
+class TestTableCommentDDL(DDLTestBase):
+    @pytest.fixture
+    def metadata(self) -> MetaData:
+        """Assemble a metadata object with one table containing one column."""
+        metadata = MetaData()
+
+        col1 = Column("foo", String)
+        col2 = Column("foo", String)
+        tbl_w_comment = Table("martin", metadata, col1, comment="foobar")
+        tbl_wo_comment = Table("prs", metadata, col2)
+
+        return metadata
+
+    @pytest.fixture
+    def table_with_comment(self, metadata) -> Table:
+        return metadata.tables.get("martin")
+
+    @pytest.fixture
+    def table_without_comment(self, metadata) -> Table:
+        return metadata.tables.get("prs")
+
+    def test_create_table_with_comment(self, table_with_comment):
+        stmt = CreateTable(table_with_comment)
+        output = self.compile(stmt)
+        assert "USING DELTA COMMENT 'foobar'" in output
+
+    def test_alter_table_add_comment(self, table_without_comment: Table):
+        table_without_comment.comment = "wireless mechanical keyboard"
+        stmt = SetTableComment(table_without_comment)
+        output = self.compile(stmt)
+
+        assert output == "COMMENT ON TABLE prs IS 'wireless mechanical keyboard'"
+
+    def test_alter_table_drop_comment(self, table_with_comment):
+        """The syntax for COMMENT ON is here: https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-comment.html"""
+        stmt = DropTableComment(table_with_comment)
+        output = self.compile(stmt)
+        assert output == "COMMENT ON TABLE martin IS NULL"