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>

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