Merge pull request #114 from Dan-Knott/dk/pg-exclude-extension-objects

postgresql: exclude extension-owned objects from introspection

Author: DanCardin

.github/workflows/test.yml

@@ -29,9 +29,6 @@ jobs:
           - postgres-drivername: postgresql+psycopg
             sqlalchemy-version: 1.4.0
 
-    env:
-      PMR_POSTGRES_IMAGE: postgres:13
-
     steps:
       - uses: actions/checkout@v3
 

src/sqlalchemy_declarative_extensions/dialects/postgresql/query.py

@@ -23,6 +23,7 @@
 from sqlalchemy_declarative_extensions.dialects.postgresql.schema import (
     databases_query,
     default_acl_query,
+    extensions_query,
     functions_query,
     object_acl_query,
     objects_query,
@@ -47,13 +48,33 @@
 from sqlalchemy_declarative_extensions.procedure import Procedure as BaseProcedure
 from sqlalchemy_declarative_extensions.sql import qualify_name
 
+EXTENSION_SCHEMAS = {
+    "postgis_topology": ["topology"],
+    "postgis_tiger_geocoder": ["tiger", "tiger_data"],
+    "pg_partman": ["partman"],
+    "timescaledb": ["timescaledb_internal"],
+    "pgmq": ["pgmq"],
+    "pgrouting": ["pgrouting"],
+    "orafce": ["orafce"],
+    "pgagent": ["pgagent"],
+}
+EXTENSION_TRIGGERS = {
+    "postgis_topology": {"layer_integrity_checks": "topology.layer"},
+}
+
 
 def get_schemas_postgresql(connection: Connection):
     from sqlalchemy_declarative_extensions.schema.base import Schema
 
+    extension_schemas = []
+    for name, _ in connection.execute(extensions_query).fetchall():
+        if name in EXTENSION_SCHEMAS:
+            extension_schemas.extend(EXTENSION_SCHEMAS[name])
+
     return {
         schema: Schema(schema)
         for schema, *_ in connection.execute(schemas_query).fetchall()
+        if schema not in extension_schemas
     }
 
 
@@ -220,8 +241,18 @@ def get_functions_postgresql(connection: Connection) -> Sequence[BaseFunction]:
 
 def get_triggers_postgresql(connection: Connection):
     triggers = []
+
+    extension_triggers = {}
+    for name, _ in connection.execute(extensions_query).fetchall():
+        if name in EXTENSION_TRIGGERS:
+            extension_triggers.update(EXTENSION_TRIGGERS[name])
+
     for t in connection.execute(triggers_query).fetchall():
         on = t.on_name if t.on_schema == "public" else f"{t.on_schema}.{t.on_name}"
+
+        if extension_triggers.get(t.name) == on:
+            continue
+
         execute = (
             t.execute_name
             if t.execute_schema == "public"

src/sqlalchemy_declarative_extensions/dialects/postgresql/schema.py

@@ -3,13 +3,14 @@
     and_,
     bindparam,
     column,
+    exists,
     func,
     literal,
     table,
     text,
     union,
 )
-from sqlalchemy.dialects.postgresql import ARRAY, CHAR
+from sqlalchemy.dialects.postgresql import ARRAY, CHAR, REGCLASS
 
 from sqlalchemy_declarative_extensions.sqlalchemy import select
 
@@ -39,6 +40,21 @@
     column("datname"),
 )
 
+pg_depend = table(
+    "pg_depend",
+    column("objid"),
+    column("classid"),
+    column("deptype"),
+)
+
+pg_extension = table(
+    "pg_extension",
+    column("oid"),
+    column("extname"),
+    column("extversion"),
+    column("extconfig"),
+)
+
 pg_namespace = table(
     "pg_namespace",
     column("oid"),
@@ -146,11 +162,24 @@ def _schema_not_pg(column=pg_namespace.c.nspname):
     )
 
 
+def _schema_not_from_extension(namespace_oid_column=pg_namespace.c.oid):
+    return _not_from_extension(namespace_oid_column, "pg_namespace")
+
+
+def _not_from_extension(obj_id_col, class_name):
+    return ~exists().where(pg_depend.c.objid == obj_id_col).where(
+        pg_depend.c.classid == literal(class_name).cast(REGCLASS)
+    ).where(pg_depend.c.deptype == literal("e"))
+
+
 _schema_not_public = pg_namespace.c.nspname != "public"
 _table_not_pg = pg_class.c.relname.notlike("pg_%")
 
 schemas_query = (
-    select(pg_namespace.c.nspname).where(_schema_not_pg()).where(_schema_not_public)
+    select(pg_namespace.c.nspname)
+    .where(_schema_not_pg())
+    .where(_schema_not_public)
+    .where(_schema_not_from_extension())
 )
 
 
@@ -160,6 +189,10 @@ def _schema_not_pg(column=pg_namespace.c.nspname):
     .where(_schema_not_public)
 )
 
+extensions_query = select(
+    pg_extension.c.extname.label("name"), pg_extension.c.extversion.label("version")
+)
+
 schema_exists_query = text(
     "SELECT schema_name FROM information_schema.schemata WHERE schema_name = :schema"
 )
@@ -238,26 +271,25 @@ def _schema_not_pg(column=pg_namespace.c.nspname):
     )
     .where(_table_not_pg)
     .where(_schema_not_pg())
+    .where(_schema_not_from_extension())
 )
 
-views_query = union(
+views_query = (
     select(
-        pg_views.c.schemaname.label("schema"),
-        pg_views.c.viewname.label("name"),
-        pg_views.c.definition.label("definition"),
-        literal(False).label("materialized"),
+        pg_namespace.c.nspname.label("schema"),
+        pg_class.c.relname.label("name"),
+        func.pg_get_viewdef(pg_class.c.oid).label("definition"),
+        (pg_class.c.relkind == literal("m")).label("materialized"),
     )
-    .where(_schema_not_pg(pg_views.c.schemaname))
-    .where(pg_views.c.viewname.notin_(["pg_stat_statements"])),
-    select(
-        pg_matviews.c.schemaname.label("schema"),
-        pg_matviews.c.matviewname.label("name"),
-        pg_matviews.c.definition.label("definition"),
-        literal(True).label("materialized"),
-    ).where(_schema_not_pg(pg_matviews.c.schemaname)),
+    .select_from(
+        pg_class.join(pg_namespace, pg_class.c.relnamespace == pg_namespace.c.oid)
+    )
+    .where(pg_class.c.relkind.in_(["v", "m"]))
+    .where(_schema_not_pg(pg_namespace.c.nspname))
+    .where(_schema_not_from_extension())
+    .where(_not_from_extension(pg_class.c.oid, "pg_class"))
 )
 
-
 views_subquery = views_query.cte()
 view_query = (
     select(views_subquery)
@@ -283,6 +315,8 @@ def _schema_not_pg(column=pg_namespace.c.nspname):
     )
     .where(pg_namespace.c.nspname.notin_(["pg_catalog", "information_schema"]))
     .where(pg_proc.c.prokind == "p")
+    .where(_schema_not_from_extension())
+    .where(_not_from_extension(pg_proc.c.oid, "pg_proc"))
 )
 
 functions_query = (
@@ -302,6 +336,7 @@ def _schema_not_pg(column=pg_namespace.c.nspname):
     )
     .where(pg_namespace.c.nspname.notin_(["pg_catalog", "information_schema"]))
     .where(pg_proc.c.prokind != "p")
+    .where(_not_from_extension(pg_proc.c.oid, "pg_proc"))
 )
 
 
@@ -335,4 +370,6 @@ def _schema_not_pg(column=pg_namespace.c.nspname):
         .join(proc_nsp, pg_proc.c.pronamespace == proc_nsp.c.oid)
     )
     .where(pg_trigger.c.tgisinternal.is_(False))
+    .where(_schema_not_from_extension())
+    .where(_not_from_extension(pg_trigger.c.oid, "pg_trigger"))
 )

tests/conftest.py

@@ -36,7 +36,7 @@ def pmr_postgres_container(pytestconfig, pmr_postgres_config: PostgresConfig):
 
 @pytest.fixture
 def pmr_postgres_config():
-    return PostgresConfig(port=None, ci_port=None)
+    return PostgresConfig(image="postgres:13", port=None, ci_port=None)
 
 
 @pytest.fixture

tests/dialect/postgresql/test_exclude_extension_objects.py

@@ -0,0 +1,54 @@
+import pytest
+from pytest_mock_resources import PostgresConfig, create_postgres_fixture
+from sqlalchemy import func, select, text
+
+from sqlalchemy_declarative_extensions import Functions, Schemas, Triggers, Views
+from sqlalchemy_declarative_extensions.function.compare import compare_functions
+from sqlalchemy_declarative_extensions.schema.compare import compare_schemas
+from sqlalchemy_declarative_extensions.trigger.compare import compare_triggers
+from sqlalchemy_declarative_extensions.view.compare import compare_views
+
+pg = create_postgres_fixture(scope="function", engine_kwargs={"echo": True})
+
+
+@pytest.fixture
+def pmr_postgres_config():
+    return PostgresConfig(image="postgis/postgis:13-3.5", port=None, ci_port=None)
+
+
+@pytest.fixture(autouse=True)
+def postgis_extension(pg):
+    with pg.connect() as connection:
+        connection.execute(text("CREATE EXTENSION postgis"))
+        connection.execute(text("CREATE EXTENSION postgis_topology"))
+        connection.commit()
+
+
+def test_functions(pg):
+    with pg.connect() as connection:
+        diff = compare_functions(connection, Functions())
+        result = connection.execute(select(func.PostGis_Version())).scalar()
+        assert result
+    assert diff == []
+
+
+def test_schemas(pg):
+    with pg.connect() as connection:
+        diff = compare_schemas(connection, Schemas())
+    assert diff == []
+
+
+def test_triggers(pg):
+    with pg.connect() as connection:
+        diff = compare_triggers(connection, Triggers())
+    assert diff == []
+
+
+def test_views(pg):
+    with pg.connect() as connection:
+        diff = compare_views(connection, Views())
+        result = connection.execute(
+            text("SELECT * FROM geometry_columns")
+        ).one_or_none()
+        assert not result
+    assert diff == []

tests/examples/test_create_role/conftest.py

@@ -6,4 +6,4 @@
 
 @pytest.fixture(scope="session")
 def pmr_postgres_config():
-    return PostgresConfig(port=None, ci_port=None)
+    return PostgresConfig(image="postgres:13", port=None, ci_port=None)

tests/examples/test_create_rows/conftest.py

@@ -6,4 +6,4 @@
 
 @pytest.fixture(scope="session")
 def pmr_postgres_config():
-    return PostgresConfig(port=None, ci_port=None)
+    return PostgresConfig(image="postgres:13", port=None, ci_port=None)

tests/examples/test_delete_rows/conftest.py

@@ -6,4 +6,4 @@
 
 @pytest.fixture(scope="session")
 def pmr_postgres_config():
-    return PostgresConfig(port=None, ci_port=None)
+    return PostgresConfig(image="postgres:13", port=None, ci_port=None)

tests/examples/test_delete_rows_ignore_unspecified/conftest.py

@@ -6,4 +6,4 @@
 
 @pytest.fixture(scope="session")
 def pmr_postgres_config():
-    return PostgresConfig(port=None, ci_port=None)
+    return PostgresConfig(image="postgres:13", port=None, ci_port=None)

tests/examples/test_delete_unspecified_grants/conftest.py

@@ -12,4 +12,4 @@ def alembic_engine(_alembic_engine):
 
 @pytest.fixture(scope="session")
 def pmr_postgres_config():
-    return PostgresConfig(port=None, ci_port=None)
+    return PostgresConfig(image="postgres:13", port=None, ci_port=None)