Track extension enums

.circleci/config.yml

@@ -63,7 +63,10 @@ jobs:
       - run:
           name: Wait for db
           command: dockerize -wait tcp://localhost:5432 -timeout 1m
-      - run: sudo apt-get install -y postgresql-client
+      - run: |
+          # avoid trying to install outdated and unavailable distro-info-data package
+          sudo apt-mark hold distro-info-data
+          sudo apt-get install -y --no-install-recommends postgresql-client
       - run:
           name: create postgres user
           command: psql postgresql://@localhost/circleci -c 'create role postgres'

schemainspect/pg/obj.py

@@ -555,11 +555,12 @@ def __eq__(self, other):
 
 
 class InspectedEnum(Inspected):
-    def __init__(self, name, schema, elements, pg_version=None):
+    def __init__(self, name, schema, elements, pg_version=None, is_extension=None):
         self.name = name
         self.schema = schema
         self.elements = elements
         self.pg_version = pg_version
+        self.is_extension = is_extension
         self.dependents = []
         self.dependent_on = []
 
@@ -1169,13 +1170,17 @@ def load_deps(self):
                         r.dependent_on.append(e_sig)
                         c.enum.dependents.append(k)
 
+                    elif e_sig in self.extension_enums:
+                        r.dependent_on.append(e_sig)
+                        c.enum.dependents.append(k)
+
             if r.parent_table:
                 pt = self.relations[r.parent_table]
                 r.dependent_on.append(r.parent_table)
                 pt.dependents.append(r.signature)
 
     def get_dependency_by_signature(self, signature):
-        things = [self.selectables, self.enums, self.triggers]
+        things = [self.selectables, self.enums, self.extension_enums, self.triggers]
 
         for thing in things:
             try:
@@ -1282,12 +1287,16 @@ def load_all_relations(self):
                 schema=i.schema,
                 elements=i.elements,
                 pg_version=self.pg_version,
+                is_extension=i.is_extension,
             )
             for i in q
         ]
-        self.enums = od((i.quoted_full_name, i) for i in enumlist)
-        q = self.c.execute(self.ALL_RELATIONS_QUERY)
+        self.enums = od((i.quoted_full_name, i) for i in enumlist if not i.is_extension)
+        self.extension_enums = od(
+            (i.quoted_full_name, i) for i in enumlist if i.is_extension
+        )
 
+        q = self.c.execute(self.ALL_RELATIONS_QUERY)
         for _, g in groupby(q, lambda x: (x.relationtype, x.schema, x.name)):
             clist = list(g)
             f = clist[0]
@@ -1299,6 +1308,8 @@ def get_enum(name, schema):
                 quoted_full_name = "{}.{}".format(
                     quoted_identifier(schema), quoted_identifier(name)
                 )
+                if quoted_full_name in self.extension_enums:
+                    return self.extension_enums[quoted_full_name]
                 return self.enums[quoted_full_name]
 
             columns = [
@@ -1618,6 +1629,7 @@ def __eq__(self, other):
             and self.relations == other.relations
             and self.sequences == other.sequences
             and self.enums == other.enums
+            and self.extension_enums == other.extension_enums
             and self.constraints == other.constraints
             and self.extensions == other.extensions
             and self.functions == other.functions

schemainspect/pg/sql/enums.sql

@@ -15,14 +15,14 @@ SELECT
       FROM pg_catalog.pg_enum e
       WHERE e.enumtypid = t.oid
       ORDER BY e.enumsortorder
-  ) as elements
+  ) as elements,
+  e.objid is not null as is_extension
 FROM pg_catalog.pg_type t
      LEFT JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace
      left outer join extension_oids e
        on t.oid = e.objid
 WHERE
   t.typcategory = 'E'
-  and e.objid is null
   -- SKIP_INTERNAL and n.nspname not in ('pg_internal', 'pg_catalog', 'information_schema', 'pg_toast')
   -- SKIP_INTERNAL and n.nspname not like 'pg_temp_%' and n.nspname not like 'pg_toast_temp_%'
 ORDER BY 1, 2;

tests/pgxn/dummycolor--0.0.1.sql

@@ -0,0 +1,7 @@
+SET client_min_messages = warning;
+
+CREATE TYPE color AS ENUM (
+    'blue',
+    'aqua',
+    'eggshell',
+    'seashell');

tests/pgxn/dummycolor.control

@@ -0,0 +1,4 @@
+# dummy extension
+comment = 'color enum dummy extension'
+default_version = '0.0.1'
+relocatable = true

tests/test_all.py

@@ -618,6 +618,61 @@ def test_sequences(db):
         assert owned.quoted_table_and_column_name == '"public"."t"."id"'
 
 
+def test_enums(db):
+    with S(db) as s:
+        s.execute(
+            """
+        create type color as enum ('blue', 'teal', 'eggshell', 'spamshell');
+        """
+        )
+
+        i = get_inspector(s)
+
+        assert len(i.extension_enums) == 0
+        enums = list(i.enums)
+        assert enums == ['"public"."color"']
+
+        color = i.enums['"public"."color"']
+        assert color.name == "color"
+        assert color.schema == "public"
+        assert color.elements == ["blue", "teal", "eggshell", "spamshell"]
+        assert color.pg_version == i.pg_version
+        assert not color.is_extension
+
+
+def test_enums_from_extensions(db):
+    # This test requires an extension defining an enum to be installed.  To run
+    # this test, copy the dummycolor extension files in test/pgxn/ into the
+    # postgres server extension directory.
+    #
+    # :; pgsharedir=$(pg_config --sharedir) && cp tests/pgxn/dummycolor* $pgsharedir/extension/
+    with S(db) as s:
+        try:
+            s.execute(
+                """
+            create schema color;
+            create extension dummycolor schema color;
+            """
+            )
+        except sqlalchemy.exc.OperationalError as e:
+            if "could not open extension control file" in str(e):
+                pytest.skip("color enum dummy extension is not available")
+            raise
+
+        i = get_inspector(s)
+
+        assert len(i.enums) == 0
+        enums = list(i.extension_enums)
+        assert enums == ['"color"."color"']
+
+        color = i.extension_enums['"color"."color"']
+        assert color.name == "color"
+        assert color.schema == "color"
+        assert color.elements == ["blue", "aqua", "eggshell", "seashell"]
+        assert color.pg_version == i.pg_version
+        assert color.is_extension
+
+
 def test_postgres_inspect(db, pytestconfig):
     if pytestconfig.getoption("timescale"):
         pytest.skip("--timescale was specified")

tests/test_deps.py

@@ -1,3 +1,5 @@
+import pytest
+import sqlalchemy
 from sqlbag import S
 
 from schemainspect import get_inspector
@@ -97,6 +99,38 @@ def test_enum_deps(db):
         assert e in i.selectables[v].dependent_on
 
 
+def test_extension_enum_deps(db):
+    ENUM_DEP_SAMPLE = """\
+create schema color;
+create extension dummycolor schema color;
+
+create table t(id integer primary key, color color.color);
+
+create view v as select * from t;
+
+"""
+    with S(db) as s:
+        try:
+            s.execute(ENUM_DEP_SAMPLE)
+        except sqlalchemy.exc.OperationalError as e:
+            if "could not open extension control file" in str(e):
+                pytest.skip("color enum dummy extension is not available")
+            raise
+
+        i = get_inspector(s)
+
+        e = '"color"."color"'
+        t = '"public"."t"'
+        v = '"public"."v"'
+
+        assert e not in i.enums
+        assert e in i.extension_enums
+
+        assert i.extension_enums[e].dependents == [t, v]
+        assert e in i.selectables[t].dependent_on
+        assert e in i.selectables[v].dependent_on
+
+
 def test_relationships(db):
     # commented-out dependencies are the dependencies that aren't tracked directly by postgres
     with S(db) as s: