Add support for table 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__)
@@ -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")
@@ -39,13 +45,26 @@ def visit_identity_column(self, identity, **kw):
         )
         return text
 
+    def visit_set_table_comment(self, create, **kw):
+        return "ALTER TABLE %s COMMENT %s" % (
+            self.preparer.format_table(create.element),
+            self.sql_compiler.render_literal_value(
+                create.element.comment, sqltypes.String()
+            ),
+        )
+
+    def visit_drop_table_comment(self, create, **kw):
+        return "ALTER TABLE %s COMMENT ''" % (
+            self.preparer.format_table(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.
         """
         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/_information_schema.py

@@ -0,0 +1,23 @@
+from sqlalchemy import MetaData, Table, Column, String, DateTime
+
+metadata = MetaData()
+
+tables = Table(
+    "tables",
+    metadata,
+    Column("table_catalog", String),
+    Column("table_schema", String),
+    Column("table_name", String),
+    Column("table_type", String),
+    Column("is_insertable_into", String),
+    Column("commit_action", String),
+    Column("table_owner", String),
+    Column("comment", String),
+    Column("created", DateTime),
+    Column("created_by", String),
+    Column("last_altered", DateTime),
+    Column("last_altered_by", String),
+    Column("data_source_format", String),
+    Column("storage_sub_directory", String),
+    schema="information_schema",
+)

src/databricks/sqlalchemy/base.py

@@ -1,9 +1,9 @@
-import re
-from typing import Any, List, Optional, Dict, Union, Collection, Iterable, Tuple
+from typing import Any, List, Optional, Dict, Union, Iterable, Tuple
 
 import databricks.sqlalchemy._ddl as dialect_ddl_impl
 import databricks.sqlalchemy._types as dialect_type_impl
 from databricks import sql
+from databricks.sqlalchemy._information_schema import tables
 from databricks.sqlalchemy._parse import (
     _describe_table_extended_result_to_dict_list,
     _match_table_not_found_string,
@@ -15,15 +15,16 @@
 )
 
 import sqlalchemy
-from sqlalchemy import DDL, event
+from sqlalchemy import DDL, event, select, bindparam, exc
 from sqlalchemy.engine import Connection, Engine, default, reflection
-from sqlalchemy.engine.reflection import ObjectKind
 from sqlalchemy.engine.interfaces import (
     ReflectedForeignKeyConstraint,
     ReflectedPrimaryKeyConstraint,
     ReflectedColumn,
+    ReflectedTableComment,
     TableKey,
 )
+from sqlalchemy.engine.reflection import ReflectionDefaults, ObjectKind, ObjectScope
 from sqlalchemy.exc import DatabaseError, SQLAlchemyError
 
 try:
@@ -285,7 +286,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 +329,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 +376,87 @@ def get_schema_names(self, connection, **kw):
         schema_list = [row[0] for row in result]
         return schema_list
 
+    def get_multi_table_comment(
+        self,
+        connection,
+        schema=None,
+        filter_names=None,
+        scope=ObjectScope.ANY,
+        kind=ObjectKind.ANY,
+        **kw,
+    ) -> Iterable[Tuple[TableKey, ReflectedTableComment]]:
+        result = []
+        _schema = schema or self.schema
+        if ObjectScope.DEFAULT in scope:
+            query = (
+                select(tables.c.table_name, tables.c.comment)
+                .select_from(tables)
+                .where(
+                    tables.c.table_catalog == self.catalog,
+                    tables.c.table_schema == _schema,
+                )
+            )
+
+            if ObjectKind.ANY not in kind:
+                where_in = set()
+                if ObjectKind.TABLE in kind:
+                    where_in.update(
+                        ["BASE TABLE", "MANAGED", "EXTERNAL", "STREAMING_TABLE"]
+                    )
+                if ObjectKind.VIEW in kind:
+                    where_in.update(["VIEW"])
+                if ObjectKind.MATERIALIZED_VIEW in kind:
+                    where_in.update(["MATERIALIZED_VIEW"])
+                query = query.where(tables.c.table_type.in_(where_in))
+
+            if filter_names:
+                query = query.where(tables.c.table_name.in_(bindparam("filter_names")))
+                result = connection.execute(
+                    query, {"filter_names": [f.lower() for f in filter_names]}
+                )
+            else:
+                result = connection.execute(query)
+
+        if ObjectScope.TEMPORARY in scope and ObjectKind.VIEW in kind:
+            result = list(result)
+            temp_views = self.get_view_names(connection, schema, only_temp=True)
+            if filter_names:
+                temp_views = set(temp_views).intersection(
+                    [f.lower() for f in filter_names]
+                )
+            result.extend(zip(temp_views, [None] * len(temp_views)))
+
+        # TODO: figure out how to return sqlalchemy.interfaces in a way that mypy respects
+        return (
+            (
+                (schema, table),
+                {"text": comment}
+                if comment is not None
+                else ReflectionDefaults.table_comment(),
+            )
+            for table, comment in result
+        )  # type: ignore
+
+    def get_table_comment(
+        self,
+        connection: Connection,
+        table_name: str,
+        schema: Optional[str] = None,
+        **kw: Any,
+    ) -> ReflectedTableComment:
+        data = self.get_multi_table_comment(
+            connection,
+            schema,
+            [table_name],
+            **kw,
+        )
+        try:
+            return dict(data)[(schema, table_name.lower())]
+        except KeyError:
+            raise exc.NoSuchTableError(
+                f"{schema}.{table_name}" if schema else table_name
+            ) from None
+
 
 @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/_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