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) <noreply@anthropic.com>

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"]

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

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
+        with open(_EMPTY_PRECLEAN_SQL, encoding="utf-8") as preclean:
-        # executes as its own top-level statement: a single DO-block runs every DROP in one
+            drop_sql = preclean.read()
-        # 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
-"""
         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,
         )

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()

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()