diff --git a/pgdevkit/migrate.py b/pgdevkit/migrate.py index ffcd160..31df1ef 100644 --- a/pgdevkit/migrate.py +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index 4444bc1..6fab0c5 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/test_migrate_grants.py b/tests/test_migrate_grants.py new file mode 100644 index 0000000..f166451 --- /dev/null +++ b/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"]) == [] diff --git a/uv.lock b/uv.lock index 698a146..644aad1 100644 --- a/uv.lock +++ b/uv.lock @@ -313,7 +313,7 @@ wheels = [ [[package]] name = "pgdevkit" -version = "0.8.0" +version = "0.9.0" source = { editable = "." } dependencies = [ { name = "docker" },