Merge pull request #1 from PeterDKay/main

Releasing dialect changes from personal forked repo to sede-open

Author: PeterDKay

src/databricks/sqlalchemy/dialect/__init__.py

@@ -80,6 +80,7 @@ class DatabricksDialect(default.DefaultDialect):
     supports_multivalues_insert: bool = True
     supports_native_decimal: bool = True
     supports_sane_rowcount: bool = False
+    supports_comments: bool = True
 
     @classmethod
     def dbapi(cls):
@@ -174,6 +175,7 @@ def get_columns(self, connection, table_name, schema=None, **kwargs):
                 "nullable": bool(col.NULLABLE),
                 "default": col.COLUMN_DEF,
                 "autoincrement": False if col.IS_AUTO_INCREMENT == "NO" else True,
+                'comment': col.REMARKS,
             }
             columns.append(this_column)
 
@@ -244,11 +246,14 @@ def get_indexes(self, connection, table_name, schema=None, **kw):
 
     def get_table_names(self, connection, schema=None, **kwargs):
         TABLE_NAME = 1
+        catalog = "`" + self.catalog + "`"
+        schema = ("`" + schema + "`") or ("`" + self.schema + "`")
+
         with self.get_driver_connection(
             connection
         )._dbapi_connection.dbapi_connection.cursor() as cur:
             sql_str = "SHOW TABLES FROM {}".format(
-                ".".join([self.catalog, schema or self.schema])
+                ".".join([catalog, schema])
             )
             data = cur.execute(sql_str).fetchall()
             _tables = [i[TABLE_NAME] for i in data]

src/databricks/sqlalchemy/dialect/base.py

@@ -1,5 +1,5 @@
 import re
-from sqlalchemy.sql import compiler
+from sqlalchemy.sql import compiler, sqltypes, ColumnElement
 
 
 class DatabricksIdentifierPreparer(compiler.IdentifierPreparer):
@@ -15,3 +15,52 @@ def __init__(self, dialect):
 class DatabricksDDLCompiler(compiler.DDLCompiler):
     def post_create_table(self, table):
         return " USING DELTA"
+
+    def visit_set_column_comment(self, create, **kw):
+        """
+        Example syntax for adding column comment:
+        "ALTER TABLE schema.table_name CHANGE COLUMN COLUMN_NAME COMMENT 'Comment to be added to column';"
+
+        """
+        return """ALTER TABLE {0} CHANGE COLUMN {1} COMMENT {2}""".format(
+            self._format_table_from_column(
+                 create, use_schema=True
+            ),
+            self.preparer.format_column(
+                create.element, use_table=False
+            ),
+            self.sql_compiler.render_literal_value(
+                create.element.comment, sqltypes.String()
+            ),
+        )
+
+    def visit_drop_column_comment(self, drop, **kw):
+        """
+        Example syntax for dropping column comment:
+        "ALTER TABLE schema.table_name CHANGE COLUMN COLUMN_NAME COMMENT '';"
+
+        Note: There is no syntactical 'DROP' statement in this case, the comment must be replaced with an empty string
+        """
+        return "ALTER TABLE {0} CHANGE COLUMN {1} COMMENT '';".format(
+            self._format_table_from_column(
+                 drop, use_schema=True
+            ),
+            self.preparer.format_column(
+                drop.element, use_table=False
+            )
+        )
+
+    def _format_table_from_column(self, column_object, use_schema=False):
+        """
+        Prepare a quoted table name from the column object (including schema if specified)
+        """
+        schema_table_column = self.preparer.format_column(
+                column_object.element, use_table=True, use_schema=True
+            )
+
+        name = schema_table_column.split(".")[1]
+
+        if use_schema:
+            name = schema_table_column.split(".")[0] + '.' + name
+
+        return name