SNOW-140347 get tabe comment views

Author: sfc-gh-stakeda

snowdialect.py

@@ -572,18 +572,43 @@ def get_schema_names(self, connection, **kw):
 
         return [self.normalize_name(row[1]) for row in cursor]
 
-    def get_table_comment(self, connection, table_name, schema=None, **kw):
+    def _get_table_comment(self, connection, table_name, schema=None, **kw):
         """
-        Returns comment of table in a dictionary as described by SQLAlchemy spec
+        Returns comment of table in a dictionary as described by SQLAlchemy spec.
         """
-        sql_command = "SHOW /* sqlalchemy:get_table_comment */ " \
-                      "TABLES LIKE '{0}'{1}".format(
+        sql_command = "SHOW /* sqlalchemy:_get_table_comment */ " \
+                      "TABLES LIKE '{}'{}".format(
                                         table_name,
-                                        (' IN SCHEMA {0}'.format(self.normalize_name(schema))) if schema else ''
+                                        (' IN SCHEMA {}'.format(self.normalize_name(schema))) if schema else ''
                                     )
         cursor = connection.execute(sql_command)
-        ans = cursor.fetchone()
-        return {'text': ans['comment'] if ans['comment'] else None}
+        return cursor.fetchone()
+
+    def _get_view_comment(self, connection, table_name, schema=None, **kw):
+        """
+        Returns comment of view in a dictionary as described by SQLAlchemy spec.
+        """
+        sql_command = "SHOW /* sqlalchemy:_get_view_comment */ " \
+            "VIEWS LIKE '{}'{}".format(
+                table_name,
+                (' IN SCHEMA {}'.format(self.normalize_name(schema))) if schema else ''
+            )
+        cursor = connection.execute(sql_command)
+        return cursor.fetchone()
+
+    def get_table_comment(self, connection, table_name, schema=None, **kw):
+        """
+        Returns comment associated with a table (or view) in a dictionary as
+        SQLAlchemy expects. Note that since SQLAlchemy may not (in fact,
+        typically does not) know if this is a table or a view, we have to
+        handle both cases here.
+        """
+        result = self._get_table_comment(connection, table_name, schema)
+        if result is None:
+            # the "table" being reflected is actually a view
+            result = self._get_view_comment(connection, table_name, schema)
+
+        return {'text': result['comment'] if result and result['comment'] else None}
 
 
 @sa_vnt.listens_for(Table, 'before_create')

test/test_core.py

@@ -546,10 +546,47 @@ def test_view_definition(engine_testaccount, db_parameters):
                sql.strip()
         assert inspector.get_view_names() == [test_view_name]
     finally:
-        engine_testaccount.execute(
-            "DROP TABLE IF EXISTS {0}".format(test_table_name))
-        engine_testaccount.execute(
-            "DROP VIEW IF EXISTS {0}".format(test_view_name))
+        engine_testaccount.execute(text(
+            "DROP TABLE IF EXISTS {0}".format(test_table_name)))
+        engine_testaccount.execute(text(
+            "DROP VIEW IF EXISTS {0}".format(test_view_name)))
+
+
+def test_view_comment_reading(engine_testaccount, db_parameters):
+    """
+    Tests reading a comment from a view once it's defined
+    """
+    test_table_name = "test_table_sqlalchemy"
+    test_view_name = "testview_sqlalchemy"
+    engine_testaccount.execute("""
+CREATE OR REPLACE TABLE {} (
+    id INTEGER,
+    name STRING
+)
+""".format(test_table_name))
+    sql = """
+CREATE OR REPLACE VIEW {} AS
+SELECT * FROM {} WHERE id > 10""".format(
+        test_view_name, test_table_name)
+    engine_testaccount.execute(text(sql).execution_options(
+        autocommit=True))
+    comment_text = "hello my viewing friends"
+    sql = "COMMENT ON VIEW {} IS '{}';".format(
+        test_view_name, comment_text)
+    engine_testaccount.execute(text(sql).execution_options(
+        autocommit=True))
+    try:
+        inspector = inspect(engine_testaccount)
+        # NOTE: sqlalchemy doesn't have a way to get view comments specifically,
+        # but the code to get table comments should work for views too
+        assert inspector.get_table_comment(test_view_name) == {'text': comment_text}
+        assert inspector.get_table_comment(test_table_name) == {'text': None}
+        assert str(inspector.get_columns(test_table_name)) == str(inspector.get_columns(test_view_name))
+    finally:
+        engine_testaccount.execute(text(
+            "DROP TABLE IF EXISTS {0}".format(test_table_name)))
+        engine_testaccount.execute(text(
+            "DROP VIEW IF EXISTS {0}".format(test_view_name)))
 
 
 @pytest.mark.skip("Temp table cannot be viewed for some reason")