feat(sftkit): improve pg introspection, prevent deletion of pg_ functions

Author: mikonse

sftkit/pyproject.toml

@@ -34,6 +34,7 @@ source = ["sftkit"]
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "session"
 minversion = "6.0"
 testpaths = ["tests"]
 

sftkit/sftkit/database/_migrations.py

@@ -6,14 +6,17 @@
 
 import asyncpg
 
+from sftkit.database import Connection
+from sftkit.database.introspection import list_views, list_triggers, list_functions, list_constraints
+
 logger = logging.getLogger(__name__)
 
 MIGRATION_VERSION_RE = re.compile(r"^-- migration: (?P<version>\w+)$")
 MIGRATION_REQURES_RE = re.compile(r"^-- requires: (?P<version>\w+)$")
 MIGRATION_TABLE = "schema_revision"
 
 
-async def _run_postgres_code(conn: asyncpg.Connection, code: str, file_name: Path):
+async def _run_postgres_code(conn: Connection, code: str, file_name: Path):
     if all(line.startswith("--") for line in code.splitlines()):
         return
     try:
@@ -32,33 +35,23 @@ async def _run_postgres_code(conn: asyncpg.Connection, code: str, file_name: Pat
         raise ValueError(f"Syntax or Access error when executing SQL code ({file_name!s}): {message!r}") from exc
 
 
-async def _drop_all_views(conn: asyncpg.Connection, schema: str):
+async def _drop_all_views(conn: Connection, schema: str):
     # TODO: we might have to find out the dependency order of the views if drop cascade does not work
-    result = await conn.fetch(
-        "select table_name from information_schema.views where table_schema = $1 and table_name !~ '^pg_';",
-        schema,
-    )
-    views = [row["table_name"] for row in result]
+    views = await list_views(conn, schema)
     if len(views) == 0:
         return
 
     # we use drop if exists here as the cascade dropping might lead the view to being already dropped
     # due to being a dependency of another view
-    drop_statements = "\n".join([f"drop view if exists {view} cascade;" for view in views])
+    drop_statements = "\n".join([f"drop view if exists {view.table_name} cascade;" for view in views])
     await conn.execute(drop_statements)
 
 
-async def _drop_all_triggers(conn: asyncpg.Connection, schema: str):
-    result = await conn.fetch(
-        "select distinct on (trigger_name, event_object_table) trigger_name, event_object_table "
-        "from information_schema.triggers where trigger_schema = $1",
-        schema,
-    )
+async def _drop_all_triggers(conn: Connection, schema: str):
+    triggers = await list_triggers(conn, schema)
     statements = []
-    for row in result:
-        trigger_name = row["trigger_name"]
-        table = row["event_object_table"]
-        statements.append(f"drop trigger {trigger_name} on {table};")
+    for trigger in triggers:
+        statements.append(f'drop trigger "{trigger.trigger_name}" on "{trigger.event_object_table}";')
 
     if len(statements) == 0:
         return
@@ -67,27 +60,20 @@ async def _drop_all_triggers(conn: asyncpg.Connection, schema: str):
     await conn.execute(drop_statements)
 
 
-async def _drop_all_functions(conn: asyncpg.Connection, schema: str):
-    result = await conn.fetch(
-        "select proname, pg_get_function_identity_arguments(oid) as signature, prokind from pg_proc "
-        "where pronamespace = $1::regnamespace;",
-        schema,
-    )
+async def _drop_all_functions(conn: Connection, schema: str):
+    funcs = await list_functions(conn, schema)
     drop_statements = []
-    for row in result:
-        kind = row["prokind"].decode("utf-8")
-        name = row["proname"]
-        signature = row["signature"]
-        if kind in ("f", "w"):
+    for func in funcs:
+        if func.prokind in ("f", "w"):
             drop_type = "function"
-        elif kind == "a":
+        elif func.prokind == "a":
             drop_type = "aggregate"
-        elif kind == "p":
+        elif func.prokind == "p":
             drop_type = "procedure"
         else:
-            raise RuntimeError(f'Unknown postgres function type "{kind}"')
+            raise RuntimeError(f'Unknown postgres function type "{func.prokind}"')
 
-        drop_statements.append(f"drop {drop_type} {name}({signature}) cascade;")
+        drop_statements.append(f'drop {drop_type} "{func.proname}"({func.signature}) cascade;')
 
     if len(drop_statements) == 0:
         return
@@ -96,37 +82,31 @@ async def _drop_all_functions(conn: asyncpg.Connection, schema: str):
     await conn.execute(drop_code)
 
 
-async def _drop_all_constraints(conn: asyncpg.Connection, schema: str):
+async def _drop_all_constraints(conn: Connection, schema: str):
     """drop all constraints in the given schema which are not unique, primary or foreign key constraints"""
-    result = await conn.fetch(
-        "select con.conname as constraint_name, rel.relname as table_name, con.contype as constraint_type "
-        "from pg_catalog.pg_constraint con "
-        "   join pg_catalog.pg_namespace nsp on nsp.oid = con.connamespace "
-        "   left join pg_catalog.pg_class rel on rel.oid = con.conrelid "
-        "where nsp.nspname = $1 and con.conname !~ '^pg_' "
-        "   and con.contype != 'p' and con.contype != 'f' and con.contype != 'u';",
-        schema,
-    )
-    constraints = []
-    for row in result:
-        constraint_name = row["constraint_name"]
-        constraint_type = row["constraint_type"].decode("utf-8")
-        table_name = row["table_name"]
+    constraints = await list_constraints(conn, schema)
+    drop_statements = []
+    for constraint in constraints:
+        constraint_name = constraint.conname
+        constraint_type = constraint.contype
+        table_name = constraint.relname
+        if constraint_type in ("p", "f", "u"):
+            continue
         if constraint_type == "c":
-            constraints.append(f"alter table {table_name} drop constraint {constraint_name};")
+            drop_statements.append(f'alter table "{table_name}" drop constraint "{constraint_name}";')
         elif constraint_type == "t":
-            constraints.append(f"drop constraint trigger {constraint_name};")
+            drop_statements.append(f"drop constraint trigger {constraint_name};")
         else:
             raise RuntimeError(f'Unknown constraint type "{constraint_type}" for constraint "{constraint_name}"')
 
-    if len(constraints) == 0:
+    if len(drop_statements) == 0:
         return
 
-    drop_statements = "\n".join(constraints)
-    await conn.execute(drop_statements)
+    drop_cmd = "\n".join(drop_statements)
+    await conn.execute(drop_cmd)
 
 
-async def _drop_db_code(conn: asyncpg.Connection, schema: str):
+async def _drop_db_code(conn: Connection, schema: str):
     await _drop_all_triggers(conn, schema=schema)
     await _drop_all_functions(conn, schema=schema)
     await _drop_all_views(conn, schema=schema)

sftkit/sftkit/database/introspection/__init__.py

@@ -0,0 +1,89 @@
+from pydantic import BaseModel
+
+from sftkit.database import Connection
+
+
+class PgFunctionDef(BaseModel):
+    proname: str
+    pronamespace: int  # oid
+    proowner: int  # oid
+    prolang: int  # oid
+    procost: float
+    prorows: int
+    provariadic: int  # oid
+    prosupport: str
+    prokind: str
+    prosecdef: bool
+    proleakproof: bool
+    proisstrict: bool
+    proretset: bool
+    provolatile: str
+    proparallel: str
+    pronargs: int
+    pronargdefaults: int
+    prorettype: int  # oid
+    proargtypes: list[int]  # oid
+    proallargtypes: list[int] | None  # oid
+    proargmodes: list[str] | None
+    proargnames: list[str] | None
+    # proargdefaults: pg_node_tree | None
+    protrftypes: list[str] | None
+    prosrc: str
+    probin: str | None
+    # prosqlbody: pg_node_tree | None
+    proconfig: list[str] | None
+    proacl: list[str] | None
+    signature: str
+
+
+async def list_functions(conn: Connection, schema: str) -> list[PgFunctionDef]:
+    return await conn.fetch_many(
+        PgFunctionDef,
+        "select pg_proc.*, pg_get_function_identity_arguments(oid) as signature from pg_proc "
+        "where pronamespace = $1::regnamespace and pg_proc.proname !~ '^pg_';",
+        schema,
+    )
+
+
+class PgViewDef(BaseModel):
+    table_name: str
+
+
+async def list_views(conn: Connection, schema: str) -> list[PgViewDef]:
+    return await conn.fetch_many(
+        PgViewDef,
+        "select table_name from information_schema.views where table_schema = $1 and table_name !~ '^pg_';",
+        schema,
+    )
+
+
+class PgTriggerDef(BaseModel):
+    trigger_name: str
+    event_object_table: str
+
+
+async def list_triggers(conn: Connection, schema: str) -> list[PgTriggerDef]:
+    return await conn.fetch_many(
+        PgTriggerDef,
+        "select distinct on (trigger_name, event_object_table) trigger_name, event_object_table "
+        "from information_schema.triggers where trigger_schema = $1",
+        schema,
+    )
+
+
+class PgConstraintDef(BaseModel):
+    conname: str
+    relname: str
+    contype: str
+
+
+async def list_constraints(conn: Connection, schema: str) -> list[PgConstraintDef]:
+    return await conn.fetch_many(
+        PgConstraintDef,
+        "select con.conname, rel.relname, con.contype "
+        "from pg_catalog.pg_constraint con "
+        "   join pg_catalog.pg_namespace nsp on nsp.oid = con.connamespace "
+        "   left join pg_catalog.pg_class rel on rel.oid = con.conrelid "
+        "where nsp.nspname = $1 and con.conname !~ '^pg_';",
+        schema,
+    )

sftkit/tests/assets/minimal_db/code/constraints.sql

@@ -0,0 +1 @@
+alter table "user" add constraint username_allowlist check (name != 'exclusion');

sftkit/tests/assets/minimal_db/code/functions.sql

@@ -0,0 +1,13 @@
+create or replace function test_func(
+    arg1 bigint,
+    arg2 text
+) returns boolean as
+$$
+<<locals>> declare
+    tmp_var double precision;
+begin
+    tmp_var = arg1 > 0 and arg2 != 'bla';
+    return tmp_var;
+end;
+$$ language plpgsql
+    set search_path = "$user", public;

sftkit/tests/assets/minimal_db/code/triggers.sql

@@ -0,0 +1,14 @@
+create or replace function user_trigger() returns trigger as
+$$
+begin
+    return NEW;
+end
+$$ language plpgsql
+    stable
+    set search_path = "$user", public;
+
+create trigger create_user_trigger
+    before insert
+    on "user"
+    for each row
+execute function user_trigger();

sftkit/tests/assets/minimal_db/code/views.sql

@@ -0,0 +1,6 @@
+create view user_with_post_count as
+    select u.*, author_counts.count
+    from "user" as u
+    join (
+        select p.author_id, count(*) as count from post as p group by p.author_id
+    ) as author_counts on u.id = author_counts.author_id;

sftkit/tests/assets/minimal_db/migrations/0001-initial.sql

@@ -0,0 +1,15 @@
+-- migration: YbcL1B3z3yKj0TrH
+-- requires: null
+
+create table "user" (
+    id bigint primary key generated always as identity,
+    name text not null unique,
+    is_registered bool not null default false,
+    comment text
+);
+
+create table post (
+    id bigint primary key generated always as identity,
+    title text not null unique,
+    author_id bigint not null references "user"(id)
+);

sftkit/tests/conftest.py

@@ -0,0 +1,67 @@
+import os
+import random
+import string
+from pathlib import Path
+
+import pytest
+import pytest_asyncio
+from pytest_asyncio import is_async_test
+
+from sftkit.database import DatabaseConfig, Pool, create_db_pool, Connection, Database
+
+ASSETS_DIR = Path(__file__).parent / "assets"
+
+
+@pytest.fixture(scope="session")
+def db_config() -> DatabaseConfig:
+    return DatabaseConfig(
+        host=os.environ.get("SFTKIT_TEST_DB_HOST"),
+        port=os.environ.get("SFTKIT_TEST_DB_PORT", 5432),
+        user=os.environ.get("SFTKIT_TEST_DB_USER"),
+        password=os.environ.get("SFTKIT_TEST_DB_PASSWORD"),
+        dbname=os.environ.get("SFTKIT_TEST_DB_DBNAME", "sftkit_test"),
+    )
+
+
+@pytest_asyncio.fixture(loop_scope="session", scope="session")
+async def static_test_db_pool(db_config: DatabaseConfig) -> Pool:
+    pool = await create_db_pool(cfg=db_config, n_connections=10)
+    yield pool
+    await pool.close()
+
+
+@pytest_asyncio.fixture(loop_scope="session", scope="function")
+async def test_db(db_config: DatabaseConfig, static_test_db_pool: Pool) -> Database:
+    dbname = "".join(random.choices(string.ascii_lowercase, k=20))
+    cfg = db_config.model_copy()
+    cfg.dbname = dbname
+    await static_test_db_pool.execute(f'create database "{dbname}"')
+    if db_config.user:
+        await static_test_db_pool.execute(f'alter database "{dbname}" owner to "{db_config.user}"')
+    mininal_db_assets = ASSETS_DIR / "minimal_db"
+    database = Database(
+        config=cfg, migrations_dir=mininal_db_assets / "migrations", code_dir=mininal_db_assets / "code"
+    )
+    await database.apply_migrations()
+    yield database
+    await static_test_db_pool.execute(f'drop database "{dbname}"')
+
+
+@pytest_asyncio.fixture(loop_scope="session", scope="function")
+async def test_db_pool(test_db: Database) -> Pool:
+    pool = await test_db.create_pool(n_connections=2)
+    yield pool
+    await pool.close()
+
+
+@pytest_asyncio.fixture(loop_scope="session", scope="function")
+async def test_db_conn(test_db_pool: Pool) -> Connection:
+    async with test_db_pool.acquire() as conn:
+        yield conn
+
+
+def pytest_collection_modifyitems(items):
+    pytest_asyncio_tests = (item for item in items if is_async_test(item))
+    session_scope_marker = pytest.mark.asyncio(loop_scope="session")
+    for async_test in pytest_asyncio_tests:
+        async_test.add_marker(session_scope_marker, append=False)

sftkit/tests/test_dummy.py

@@ -0,0 +1,22 @@
+from sftkit.database import Connection
+from sftkit.database.introspection import list_functions, list_views, list_triggers, list_constraints
+
+
+async def test_introspection_functions(test_db_conn: Connection):
+    funcs = await list_functions(test_db_conn, "public")
+    assert len([x for x in funcs if x.proname == "test_func"]) > 0
+
+
+async def test_introspection_view(test_db_conn: Connection):
+    views = await list_views(test_db_conn, "public")
+    assert len([x for x in views if x.table_name == "user_with_post_count"]) > 0
+
+
+async def test_introspection_triggers(test_db_conn: Connection):
+    triggers = await list_triggers(test_db_conn, "public")
+    assert len([x for x in triggers if x.trigger_name == "create_user_trigger"]) > 0
+
+
+async def test_introspection_constraints(test_db_conn: Connection):
+    constraints = await list_constraints(test_db_conn, "public")
+    assert len([x for x in constraints if x.conname == "username_allowlist"]) > 0