Add missing_privileges and execute_sql_script to migrate

pgdevkit/migrate.py

@@ -303,6 +303,14 @@ def record_applied(conninfo: str, tracking_table: str, filename: str) -> bool:
             return False
 
 
+def created_table_names(sql: str) -> list[str]:
+    """Table names any CREATE TABLE statement in this raw SQL script targets. Public
+    wrapper around the same detection `apply_migration` uses internally, for callers that
+    run a script directly (e.g. via `execute_sql_script`) instead of through a tracked
+    migration file, and still want to know what tables -- if any -- it created."""
+    return _created_table_names(_split_sql(sql))
+
+
 def verify_created_tables(conninfo: str, stmts: list[str]) -> list[str]:
     """Table names from this migration's CREATE TABLE statements that do NOT exist in the
     database. Empty means everything landed."""
@@ -318,6 +326,39 @@ def verify_created_tables(conninfo: str, stmts: list[str]) -> list[str]:
     return missing
 
 
+def missing_privileges(conninfo: str, role: str, tables: list[str], privilege: str = "select") -> list[str]:
+    """Table names from `tables` that `role` cannot currently exercise `privilege` on
+    (checked via Postgres's own `has_table_privilege`). Empty means the role can access
+    all of them. For guarding against the classic "migration creates a table, nobody
+    grants it to the app's runtime role" gap: check the tables a migration just created
+    against the role that will actually query them at runtime."""
+    if not tables:
+        return []
+    missing = []
+    with psycopg.connect(conninfo) as con:
+        for tbl in tables:
+            row = con.execute("select has_table_privilege(%s, %s, %s)", (role, tbl, privilege)).fetchone()
+            if not (row and row[0]):
+                missing.append(tbl)
+    return missing
+
+
+def _execute_stmts(conninfo: str, stmts: list[str]) -> None:
+    with psycopg.connect(conninfo) as con:
+        for stmt in stmts:
+            con.execute(cast(LiteralString, stmt))
+        con.commit()
+
+
+def execute_sql_script(conninfo: str, sql: str) -> None:
+    """Run a raw SQL script as one committed transaction, split into statements the same
+    statement-boundary-safe way apply_migration is (dollar-quoted blocks, string literals,
+    and line comments never get split mid-statement). Unlike apply_migration, this does
+    no tracking-table bookkeeping and isn't forward-only -- for scripts meant to re-run
+    every time, like an idempotent `GRANT ... ON ALL TABLES IN SCHEMA` privilege sync."""
+    _execute_stmts(conninfo, _split_sql(sql))
+
+
 @dataclass
 class ApplyResult:
     filename: str
@@ -337,11 +378,7 @@ def apply_migration(
     stmts = _split_sql(sql)
 
     if not already_done:
-        # DDL in its own committed transaction.
-        with psycopg.connect(conninfo) as con:
-            for stmt in stmts:
-                con.execute(cast(LiteralString, stmt))
-            con.commit()
+        _execute_stmts(conninfo, stmts)
 
     # Tracking insert is a separate connection/transaction so a missing tracking table
     # never rolls back the DDL that was just applied.

pyproject.toml

@@ -11,7 +11,7 @@ packages = ["pgdevkit"]
 
 [project]
 name = "pgdevkit"
-version = "0.8.0"
+version = "0.9.0"
 description = "A helper for developing with Postgres"
 readme = "README.md"
 requires-python = ">=3.14"

tests/test_migrate_grants.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import psycopg
+import pytest
+
+from pgdevkit.migrate import created_table_names, execute_sql_script, missing_privileges
+
+ROLE = "pgdevkit_test_grants_role"
+
+
+@pytest.fixture
+def app_role(clean_db: str):
+    """A throwaway, no-login role to check privileges against -- dropped afterwards since
+    roles are cluster-wide, not scoped to the (per-test) database like clean_db's schemas."""
+    with psycopg.connect(clean_db, autocommit=True) as con:
+        con.execute(f"DROP ROLE IF EXISTS {ROLE}")
+        con.execute(f"CREATE ROLE {ROLE} NOLOGIN")
+    yield ROLE
+    with psycopg.connect(clean_db, autocommit=True) as con:
+        # DROP OWNED BY revokes any grants left on tables the role can still see --
+        # otherwise DROP ROLE fails with "cannot be dropped because some objects depend
+        # on it" for as long as the granted-on table (e.g. widgets) still exists.
+        con.execute(f"DROP OWNED BY {ROLE}")
+        con.execute(f"DROP ROLE IF EXISTS {ROLE}")
+
+
+def test_created_table_names_finds_tables_in_a_raw_script():
+    sql = "GRANT SELECT ON foo TO some_role; CREATE TABLE public.widgets (id int primary key);"
+    assert created_table_names(sql) == ["public.widgets"]
+
+
+def test_created_table_names_empty_for_a_script_with_no_create_table():
+    assert created_table_names("GRANT SELECT ON ALL TABLES IN SCHEMA public TO some_role;") == []
+
+
+def test_execute_sql_script_runs_multiple_statements_in_one_transaction(clean_db: str):
+    execute_sql_script(
+        clean_db,
+        """
+        -- a leading comment shouldn't confuse statement splitting
+        CREATE TABLE public.widgets (id int primary key);
+        INSERT INTO public.widgets VALUES (1);
+        """,
+    )
+    with psycopg.connect(clean_db) as con:
+        row = con.execute("select count(*) from public.widgets").fetchone()
+        assert row == (1,)
+
+
+def test_execute_sql_script_rolls_back_all_statements_on_any_failure(clean_db: str):
+    with pytest.raises(psycopg.Error):
+        execute_sql_script(
+            clean_db,
+            "CREATE TABLE public.widgets (id int primary key); NOT VALID SQL HERE;",
+        )
+    with psycopg.connect(clean_db) as con:
+        row = con.execute("select to_regclass('public.widgets')").fetchone()
+        assert row is not None and row[0] is None
+
+
+def test_missing_privileges_reports_table_without_grant(clean_db: str, app_role: str):
+    with psycopg.connect(clean_db, autocommit=True) as con:
+        con.execute("CREATE TABLE public.widgets (id int primary key)")
+    assert missing_privileges(clean_db, app_role, ["public.widgets"]) == ["public.widgets"]
+
+
+def test_missing_privileges_empty_once_granted(clean_db: str, app_role: str):
+    with psycopg.connect(clean_db, autocommit=True) as con:
+        con.execute("CREATE TABLE public.widgets (id int primary key)")
+        con.execute(f"GRANT SELECT ON public.widgets TO {app_role}")
+    assert missing_privileges(clean_db, app_role, ["public.widgets"]) == []
+
+
+def test_missing_privileges_empty_for_no_tables(clean_db: str, app_role: str):
+    assert missing_privileges(clean_db, app_role, []) == []
+
+
+def test_execute_sql_script_grant_on_all_tables_covers_tables_created_after_the_fact(
+    clean_db: str, app_role: str
+):
+    """The actual use case: a re-run of `GRANT ... ON ALL TABLES IN SCHEMA` after a new
+    table appears picks it up automatically -- unlike ALTER DEFAULT PRIVILEGES, which only
+    covers tables created later by the exact role that ran the ALTER DEFAULT PRIVILEGES
+    statement itself."""
+    with psycopg.connect(clean_db, autocommit=True) as con:
+        con.execute("CREATE TABLE public.widgets (id int primary key)")
+    assert missing_privileges(clean_db, app_role, ["public.widgets"]) == ["public.widgets"]
+
+    execute_sql_script(clean_db, f"GRANT SELECT ON ALL TABLES IN SCHEMA public TO {app_role}")
+
+    assert missing_privileges(clean_db, app_role, ["public.widgets"]) == []