Expanded Comment support

schemainspect/misc.py

@@ -53,10 +53,12 @@ def unquoted_identifier(identifier, *, schema=None, identity_arguments=None):
     return s
 
 
-def quoted_identifier(identifier, schema=None, identity_arguments=None):
+def quoted_identifier(identifier, schema=None, identity_arguments=None, table=None):
     if identifier is None and schema is not None:
         return '"{}"'.format(schema.replace('"', '""'))
     s = '"{}"'.format(identifier.replace('"', '""'))
+    if table:
+        s = '"{}".{}'.format(table.replace('"', '""'), s)
     if schema:
         s = '"{}".{}'.format(schema.replace('"', '""'), s)
     if identity_arguments is not None:

schemainspect/pg/obj.py

@@ -34,6 +34,7 @@
 COLLATIONS_QUERY = resource_text("sql/collations.sql")
 COLLATIONS_QUERY_9 = resource_text("sql/collations9.sql")
 RLSPOLICIES_QUERY = resource_text("sql/rlspolicies.sql")
+COMMENTS_QUERY = resource_text("sql/comments.sql")
 
 
 class InspectedSelectable(BaseInspectedSelectable):
@@ -994,6 +995,34 @@ def key(self):
         return self.object_type, self.quoted_full_name, self.target_user, self.privilege
 
 
+class InspectedComment(Inspected):
+    def __init__(self, object_type, identifier, comment):
+        self.identifier = identifier
+        self.object_type = object_type
+        self.comment = comment
+
+    @property
+    def drop_statement(self):
+        return "comment on {} {} is null;".format(self.object_type, self.identifier)
+
+    @property
+    def create_statement(self):
+        return "comment on {} {} is $cmt${}$cmt$;".format(
+            self.object_type, self.identifier, self.comment
+        )
+
+    @property
+    def key(self):
+        return "{} {}".format(self.object_type, self.identifier)
+
+    def __eq__(self, other):
+        return (
+            self.object_type == other.object_type
+            and self.identifier == other.identifier
+            and self.comment == other.comment
+        )
+
+
 RLS_POLICY_CREATE = """create policy {name}
 on {table_name}
 as {permissiveness}
@@ -1134,6 +1163,7 @@ def processed(q):
         self.SCHEMAS_QUERY = processed(SCHEMAS_QUERY)
         self.PRIVILEGES_QUERY = processed(PRIVILEGES_QUERY)
         self.TRIGGERS_QUERY = processed(TRIGGERS_QUERY)
+        self.COMMENTS_QUERY = processed(COMMENTS_QUERY)
 
         super(PostgreSQL, self).__init__(c, include_internal)
 
@@ -1160,6 +1190,7 @@ def load_all(self):
         self.load_rlspolicies()
         self.load_types()
         self.load_domains()
+        self.load_comments()
 
         self.load_deps()
         self.load_deps_all()
@@ -1663,6 +1694,18 @@ def col(defn):
         ]  # type: list[InspectedType]
         self.domains = od((t.signature, t) for t in domains)
 
+    def load_comments(self):
+        q = self.execute(self.COMMENTS_QUERY)
+        comments = [
+            InspectedComment(
+                i.object_type,
+                i.identifier,
+                i.comment,
+            )
+            for i in q
+        ]  # type: list[InspectedComment]
+        self.comments = od((t.key, t) for t in comments)
+
     def filter_schema(self, schema=None, exclude_schema=None):
         if schema and exclude_schema:
             raise ValueError("Can only have schema or exclude schema, not both")
@@ -1765,4 +1808,5 @@ def __eq__(self, other):
             and self.triggers == other.triggers
             and self.collations == other.collations
             and self.rlspolicies == other.rlspolicies
+            and self.comments == other.comments
         )

schemainspect/pg/sql/comments.sql

@@ -0,0 +1,23 @@
+SELECT
+  CASE (d.iden).type
+    WHEN 'domain constraint' THEN 'constraint'
+    WHEN 'table column' THEN 'column'
+    WHEN 'table constraint' THEN 'constraint'
+    ELSE (d.iden).type::TEXT
+  END AS object_type,
+  (d.iden).identity AS identifier,
+  d.description AS comment
+FROM (
+  SELECT
+    pg_identify_object(classoid, objoid, objsubid) AS iden,
+    DESCRIPTION
+  FROM pg_description
+) d
+WHERE
+    (
+        (d.iden).schema IS NULL
+        AND (d.iden).type = 'trigger'
+    ) OR (
+        (d.iden).schema <> 'pg_catalog'
+        AND (d.iden).schema <> 'information_schema'
+    );

tests/test_all.py

@@ -279,7 +279,9 @@ def setup_pg_schema(s):
             language sql;
         """
     )
-    s.execute("comment on function films_f(date, text, date) is 'films_f comment'")
+    s.execute(
+        "comment on function public.films_f(date, text, date) is 'films_f comment'"
+    )
     s.execute(
         """
         CREATE OR REPLACE FUNCTION inc_f(integer) RETURNS integer AS $$
@@ -495,6 +497,19 @@ def asserts_pg(i, has_timescale=False):
     with raises(ValueError):
         tid.change_string_to_enum_statement("t")
 
+    # comments
+    assert len(i.comments) == 2
+    assert (
+        i.comments[
+            "function public.films_f(pg_catalog.date,pg_catalog.text,pg_catalog.date)"
+        ].create_statement
+        == "comment on function public.films_f(pg_catalog.date,pg_catalog.text,pg_catalog.date) is $cmt$films_f comment$cmt$;"
+    )
+    assert (
+        i.comments["table public.emptytable"].create_statement
+        == "comment on table public.emptytable is $cmt$emptytable comment$cmt$;"
+    )
+
 
 def test_weird_names(db):
     with S(db) as s: