From 8409843ff9f0a7cbbf4f072533baeed5d9bfc63b Mon Sep 17 00:00:00 2001 From: Kevin Veen-Birkenbach Date: Mon, 20 Jul 2026 19:59:56 +0200 Subject: [PATCH] fix(restore): wrap the postgres dump replay in a single transaction The dump replay ran statement-by-statement with autocommit. When restore --empty runs against a LIVE database, the pre-clean drops a table, the replay recreates it and autocommits, and a concurrent writer (discourse's mini_scheduler upserting scheduler_stats(id=1)) inserts the same primary key into the empty table before the dump's COPY loads it. The COPY then aborts with a duplicate-key violation under ON_ERROR_STOP and the whole restore fails. Running the replay with --single-transaction keeps the recreated table invisible to other sessions until commit, so the writer can never insert the racing row. The --empty pre-clean stays multi-statement (\gexec, one DROP per statement): running every DROP in one transaction exhausts max_locks_per_transaction on large schemas (e.g. gitlab). Extract the pre-clean SQL from the inline string into restore/db/empty_preclean.sql (loaded via dirname(__file__)) and declare it as package-data so it ships in the wheel. Add a unit test guarding the single-transaction/multi-statement split and an e2e that reproduces the live-writer race. Co-Authored-By: Claude Opus 4.8 (1M context) --- pyproject.toml | 3 + src/baudolo/restore/db/empty_preclean.sql | 64 ++++++ src/baudolo/restore/db/postgres.py | 79 ++------ ...postgres_single_transaction_live_writer.py | 185 ++++++++++++++++++ .../test_postgres_single_transaction.py | 44 +++++ 5 files changed, 309 insertions(+), 66 deletions(-) create mode 100644 src/baudolo/restore/db/empty_preclean.sql create mode 100644 tests/e2e/test_e2e_postgres_single_transaction_live_writer.py create mode 100644 tests/unit/restore/test_postgres_single_transaction.py diff --git a/pyproject.toml b/pyproject.toml index 5adfa2e..8957378 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,3 +27,6 @@ package-dir = { "" = "src" } [tool.setuptools.packages.find] where = ["src"] exclude = ["tests*"] + +[tool.setuptools.package-data] +"baudolo.restore.db" = ["*.sql"] diff --git a/src/baudolo/restore/db/empty_preclean.sql b/src/baudolo/restore/db/empty_preclean.sql new file mode 100644 index 0000000..81c39ee --- /dev/null +++ b/src/baudolo/restore/db/empty_preclean.sql @@ -0,0 +1,64 @@ +-- Owner-filtered pre-clean for `restore --empty`. Emitted as one DROP per row and +-- run via \gexec so each executes as its own top-level statement: a single DO-block +-- would run every DROP in one transaction and exhaust max_locks_per_transaction on +-- large schemas (e.g. gitlab). Also drops user-owned non-public schemas so a dump +-- that CREATE SCHEMAs (e.g. discourse's discourse_functions) does not fail on an +-- already-existing schema. Extension members (pg_trgm's set_limit) are +-- superuser-owned; IF EXISTS absorbs the CASCADE fallout. +SELECT format('DROP %s IF EXISTS public.%s CASCADE', obj.type, obj.name) + FROM ( + SELECT format('%I', c.relname) AS name, + CASE c.relkind + WHEN 'v' THEN 'VIEW' + WHEN 'm' THEN 'MATERIALIZED VIEW' + WHEN 'f' THEN 'FOREIGN TABLE' + ELSE 'TABLE' + END AS type + FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = 'public' AND c.relkind IN ('r', 'p', 'v', 'm', 'f') + AND pg_get_userbyid(c.relowner) = current_user + UNION ALL + -- Overloaded functions share a proname; DROP needs the identity + -- signature or psql aborts with "function name is not unique". + SELECT format('%I(%s)', p.proname, pg_get_function_identity_arguments(p.oid)) AS name, + CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END AS type + FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace + WHERE n.nspname = 'public' AND p.prokind IN ('f', 'p', 'w') + AND pg_get_userbyid(p.proowner) = current_user + UNION ALL + SELECT format('%I', c.relname) AS name, 'SEQUENCE' AS type + FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = 'public' AND c.relkind = 'S' + AND pg_get_userbyid(c.relowner) = current_user + UNION ALL + SELECT format('%I', t.typname) AS name, 'TYPE' AS type + FROM pg_type t JOIN pg_namespace n ON n.oid = t.typnamespace + WHERE n.nspname = 'public' + AND pg_get_userbyid(t.typowner) = current_user + AND (t.typtype IN ('e', 'd') + OR (t.typtype = 'c' AND EXISTS ( + SELECT 1 FROM pg_class c2 + WHERE c2.oid = t.typrelid AND c2.relkind = 'c'))) + UNION ALL + SELECT format('%I', col.collname) AS name, 'COLLATION' AS type + FROM pg_collation col JOIN pg_namespace n ON n.oid = col.collnamespace + WHERE n.nspname = 'public' + AND pg_get_userbyid(col.collowner) = current_user + UNION ALL + SELECT format('%I', ts.cfgname) AS name, 'TEXT SEARCH CONFIGURATION' AS type + FROM pg_ts_config ts JOIN pg_namespace n ON n.oid = ts.cfgnamespace + WHERE n.nspname = 'public' + AND pg_get_userbyid(ts.cfgowner) = current_user + UNION ALL + SELECT format('%I', d.dictname) AS name, 'TEXT SEARCH DICTIONARY' AS type + FROM pg_ts_dict d JOIN pg_namespace n ON n.oid = d.dictnamespace + WHERE n.nspname = 'public' + AND pg_get_userbyid(d.dictowner) = current_user + ) obj +UNION ALL +SELECT format('DROP SCHEMA IF EXISTS %I CASCADE', n.nspname) + FROM pg_namespace n + WHERE NOT starts_with(n.nspname, 'pg_') + AND n.nspname NOT IN ('public', 'information_schema') + AND pg_get_userbyid(n.nspowner) = current_user +\gexec diff --git a/src/baudolo/restore/db/postgres.py b/src/baudolo/restore/db/postgres.py index e186dcf..b08553b 100644 --- a/src/baudolo/restore/db/postgres.py +++ b/src/baudolo/restore/db/postgres.py @@ -7,6 +7,7 @@ from collections.abc import Iterable, Iterator from ..run import docker_exec _SUPERUSER_ONLY_PREFIXES = (b"COMMENT ON EXTENSION", b"ALTER DEFAULT PRIVILEGES") +_EMPTY_PRECLEAN_SQL = os.path.join(os.path.dirname(__file__), "empty_preclean.sql") def filter_superuser_only_lines(lines: Iterable[bytes]) -> Iterator[bytes]: @@ -53,71 +54,8 @@ def restore_postgres_sql( docker_env = {"PGPASSWORD": password} if empty: - # Owner-filtered pre-clean emitted as one DROP per row and run via \gexec so each - # executes as its own top-level statement: a single DO-block runs every DROP in one - # transaction and exhausts max_locks_per_transaction on large schemas (e.g. gitlab). - # Also drop user-owned non-public schemas so a dump that CREATE SCHEMAs (e.g. - # discourse's discourse_functions) does not fail on an already-existing schema. - # Extension members (pg_trgm's set_limit) are superuser-owned; IF EXISTS absorbs CASCADE fallout. - drop_sql = r""" -SELECT format('DROP %s IF EXISTS public.%s CASCADE', obj.type, obj.name) - FROM ( - SELECT format('%I', c.relname) AS name, - CASE c.relkind - WHEN 'v' THEN 'VIEW' - WHEN 'm' THEN 'MATERIALIZED VIEW' - WHEN 'f' THEN 'FOREIGN TABLE' - ELSE 'TABLE' - END AS type - FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace - WHERE n.nspname = 'public' AND c.relkind IN ('r', 'p', 'v', 'm', 'f') - AND pg_get_userbyid(c.relowner) = current_user - UNION ALL - -- Overloaded functions share a proname; DROP needs the identity - -- signature or psql aborts with "function name is not unique". - SELECT format('%I(%s)', p.proname, pg_get_function_identity_arguments(p.oid)) AS name, - CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END AS type - FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace - WHERE n.nspname = 'public' AND p.prokind IN ('f', 'p', 'w') - AND pg_get_userbyid(p.proowner) = current_user - UNION ALL - SELECT format('%I', c.relname) AS name, 'SEQUENCE' AS type - FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace - WHERE n.nspname = 'public' AND c.relkind = 'S' - AND pg_get_userbyid(c.relowner) = current_user - UNION ALL - SELECT format('%I', t.typname) AS name, 'TYPE' AS type - FROM pg_type t JOIN pg_namespace n ON n.oid = t.typnamespace - WHERE n.nspname = 'public' - AND pg_get_userbyid(t.typowner) = current_user - AND (t.typtype IN ('e', 'd') - OR (t.typtype = 'c' AND EXISTS ( - SELECT 1 FROM pg_class c2 - WHERE c2.oid = t.typrelid AND c2.relkind = 'c'))) - UNION ALL - SELECT format('%I', col.collname) AS name, 'COLLATION' AS type - FROM pg_collation col JOIN pg_namespace n ON n.oid = col.collnamespace - WHERE n.nspname = 'public' - AND pg_get_userbyid(col.collowner) = current_user - UNION ALL - SELECT format('%I', ts.cfgname) AS name, 'TEXT SEARCH CONFIGURATION' AS type - FROM pg_ts_config ts JOIN pg_namespace n ON n.oid = ts.cfgnamespace - WHERE n.nspname = 'public' - AND pg_get_userbyid(ts.cfgowner) = current_user - UNION ALL - SELECT format('%I', d.dictname) AS name, 'TEXT SEARCH DICTIONARY' AS type - FROM pg_ts_dict d JOIN pg_namespace n ON n.oid = d.dictnamespace - WHERE n.nspname = 'public' - AND pg_get_userbyid(d.dictowner) = current_user - ) obj -UNION ALL -SELECT format('DROP SCHEMA IF EXISTS %I CASCADE', n.nspname) - FROM pg_namespace n - WHERE NOT starts_with(n.nspname, 'pg_') - AND n.nspname NOT IN ('public', 'information_schema') - AND pg_get_userbyid(n.nspowner) = current_user -\gexec -""" + with open(_EMPTY_PRECLEAN_SQL, encoding="utf-8") as preclean: + drop_sql = preclean.read() docker_exec( container, ["psql", "-v", "ON_ERROR_STOP=1", "-U", user, "-d", db_name], @@ -134,7 +72,16 @@ SELECT format('DROP SCHEMA IF EXISTS %I CASCADE', n.nspname) filtered.seek(0) docker_exec( container, - ["psql", "-v", "ON_ERROR_STOP=1", "-U", user, "-d", db_name], + [ + "psql", + "--single-transaction", + "-v", + "ON_ERROR_STOP=1", + "-U", + user, + "-d", + db_name, + ], stdin=filtered, docker_env=docker_env, ) diff --git a/tests/e2e/test_e2e_postgres_single_transaction_live_writer.py b/tests/e2e/test_e2e_postgres_single_transaction_live_writer.py new file mode 100644 index 0000000..549aa24 --- /dev/null +++ b/tests/e2e/test_e2e_postgres_single_transaction_live_writer.py @@ -0,0 +1,185 @@ +# tests/e2e/test_e2e_postgres_single_transaction_live_writer.py +import unittest + +from .helpers import ( + POSTGRES_IMAGE, + POSTGRES_DATA_DIR, + backup_run, + cleanup_docker, + create_minimal_compose_dir, + ensure_empty_dir, + latest_version_dir, + require_docker, + run, + unique, + wait_for_postgres, + write_databases_csv, +) + +# The discourse restore-drill race: `restore --empty` replays the dump into a +# LIVE database while a background writer keeps touching a primary-key row +# (discourse's mini_scheduler upserts scheduler_stats(id=1)). Without a +# single-transaction replay, the pre-clean drops the table, the replay recreates +# it and auto-commits, the writer wins the gap and inserts id=1, and the dump's +# COPY of the same id then aborts with a duplicate-key violation under +# ON_ERROR_STOP -> the whole restore fails. The --single-transaction replay keeps +# the recreated table invisible until commit, so the writer can never insert the +# racing row and the restore completes. A wide filler table makes the COPY slow +# enough that the non-transactional variant loses the race deterministically. +SEED_SQL = ( + "CREATE TABLE public.scheduler_stats (id int primary key, v text);" + "INSERT INTO public.scheduler_stats VALUES (1, 'from-dump');" + "CREATE TABLE public.filler (id serial primary key, blob text);" + "INSERT INTO public.filler (blob)" + " SELECT repeat('x', 512) FROM generate_series(1, 100000);" +) + +WRITER_LOOP = ( + "while true; do " + "psql -h 127.0.0.1 -U postgres -d appdb " + "-c \"INSERT INTO public.scheduler_stats(id, v) VALUES (1, 'live') " + 'ON CONFLICT (id) DO NOTHING;" >/dev/null 2>&1; ' + "done" +) + + +class TestE2EPostgresSingleTransactionLiveWriter(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + require_docker() + cls.prefix = unique("baudolo-e2e-pg-single-txn") + cls.backups_dir = f"/tmp/{cls.prefix}/Backups" + ensure_empty_dir(cls.backups_dir) + cls.compose_dir = create_minimal_compose_dir(f"/tmp/{cls.prefix}") + cls.repo_name = cls.prefix + + cls.pg_container = f"{cls.prefix}-pg" + cls.pg_volume = f"{cls.prefix}-pg-vol" + cls.writer = f"{cls.prefix}-writer" + cls.containers = [cls.pg_container, cls.writer] + cls.volumes = [cls.pg_volume] + + run(["docker", "volume", "create", cls.pg_volume]) + run( + [ + "docker", + "run", + "-d", + "--name", + cls.pg_container, + "-e", + "POSTGRES_PASSWORD=pgpw", + "-e", + "POSTGRES_DB=appdb", + "-e", + "POSTGRES_USER=postgres", + "-v", + f"{cls.pg_volume}:{POSTGRES_DATA_DIR}", + POSTGRES_IMAGE, + ] + ) + wait_for_postgres(cls.pg_container, user="postgres", timeout_s=90) + + run( + [ + "docker", + "exec", + cls.pg_container, + "sh", + "-lc", + f'psql -U postgres -d appdb -v ON_ERROR_STOP=1 -c "{SEED_SQL}"', + ] + ) + + cls.databases_csv = f"/tmp/{cls.prefix}/databases.csv" + write_databases_csv( + cls.databases_csv, [(cls.pg_container, "appdb", "postgres", "pgpw")] + ) + backup_run( + backups_dir=cls.backups_dir, + repo_name=cls.repo_name, + compose_dir=cls.compose_dir, + databases_csv=cls.databases_csv, + database_containers=[cls.pg_container], + images_no_stop_required=[POSTGRES_IMAGE], + ) + cls.hash, cls.version = latest_version_dir(cls.backups_dir, cls.repo_name) + + run( + [ + "docker", + "run", + "-d", + "--name", + cls.writer, + "--network", + f"container:{cls.pg_container}", + "-e", + "PGPASSWORD=pgpw", + POSTGRES_IMAGE, + "sh", + "-lc", + WRITER_LOOP, + ] + ) + + cls.restore = run( + [ + "baudolo-restore", + "postgres", + cls.pg_volume, + cls.hash, + cls.version, + "--backups-dir", + cls.backups_dir, + "--repo-name", + cls.repo_name, + "--container", + cls.pg_container, + "--db-name", + "appdb", + "--db-user", + "postgres", + "--db-password", + "pgpw", + "--empty", + ], + capture=True, + check=False, + ) + + run(["docker", "rm", "-f", cls.writer], capture=True, check=False) + + @classmethod + def tearDownClass(cls) -> None: + cleanup_docker(containers=cls.containers, volumes=cls.volumes) + + def _scalar(self, sql: str) -> str: + p = run( + [ + "docker", + "exec", + self.pg_container, + "sh", + "-lc", + f'psql -U postgres -d appdb -t -A -c "{sql}"', + ] + ) + return (p.stdout or "").strip() + + def test_restore_survived_the_live_writer(self) -> None: + self.assertEqual( + self.restore.returncode, + 0, + f"restore aborted (duplicate-key race not contained):\n{self.restore.stderr}", + ) + + def test_primary_key_row_restored(self) -> None: + self.assertEqual( + self._scalar("SELECT count(*) FROM public.scheduler_stats WHERE id=1;"), + "1", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/restore/test_postgres_single_transaction.py b/tests/unit/restore/test_postgres_single_transaction.py new file mode 100644 index 0000000..5353309 --- /dev/null +++ b/tests/unit/restore/test_postgres_single_transaction.py @@ -0,0 +1,44 @@ +import tempfile +import unittest +from unittest.mock import MagicMock, patch + +from baudolo.restore.db import postgres as pg_mod + + +class TestPostgresSingleTransaction(unittest.TestCase): + def test_replay_is_single_transaction_but_preclean_is_not(self) -> None: + calls = [] + + def _capture(container, argv, **kwargs): + calls.append(argv) + return MagicMock() + + with tempfile.NamedTemporaryFile(suffix=".sql") as sql: + sql.write(b"CREATE TABLE t (id int);\nINSERT INTO t VALUES (1);\n") + sql.flush() + with patch.object(pg_mod, "docker_exec", side_effect=_capture): + pg_mod.restore_postgres_sql( + container="db", + db_name="discourse", + user="discourse", + password="pw", + sql_path=sql.name, + empty=True, + ) + + self.assertEqual(len(calls), 2, f"expected pre-clean + replay: {calls}") + preclean, replay = calls[0], calls[1] + self.assertNotIn( + "--single-transaction", + preclean, + "pre-clean must stay multi-statement or it exhausts max_locks on large schemas", + ) + self.assertIn( + "--single-transaction", + replay, + "dump replay must be atomic so a live concurrent writer cannot trip a duplicate-key abort", + ) + + +if __name__ == "__main__": + unittest.main()