fix(restore): refuse a cluster restore that would destroy what it cannot restore

The --empty pre-clean is a catalog-wide sweep: it drops every non-template database and every non-pg_ role of the instance. On a dedicated instance that is exactly right, because the dump recreates all of it. On a shared one it destroys databases the dump does not carry, with nothing to restore them from - and no test ever executed that sweep, because the e2e dropped the cluster by hand first and left the pre-clean with zero rows to generate.

Scoping the sweep to the dump's own inventory looks like the fix and is worse. A surviving database that owns or merely grants to one of the dump's roles pins that role in pg_shdepend; DROP OWNED BY only reaches the control database the pre-clean is connected to, so DROP ROLE fails - after phase 1 has already dropped the dump's databases. ON_ERROR_STOP aborts, the replay never starts, and the instance is left half emptied.

So the instance is checked instead. --empty now refuses when the instance holds a database the dump does not carry, names it, and touches nothing. The sweep stays as it was, safe behind that refusal. Reading the dump's inventory needs a real identifier parser: a quoted name may hold spaces, and psql options precede the target of a \\connect line.

The e2e no longer drops the cluster itself, so --empty has to do it and the replay has to put it back; a second pass then adds a foreign database and requires the refusal to leave both it and the restored data alone.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

src/baudolo/restore/db/cluster.py

@@ -30,6 +30,139 @@ from .version import guard
 CONTROL_DB = "postgres"
 _CLUSTER_PRECLEAN_SQL = os.path.join(os.path.dirname(__file__), "cluster_preclean.sql")
 _CREATE_ROLE = re.compile(rb'^CREATE ROLE "?([^";]+)"?;\s*$')
+_CREATE_DATABASE = re.compile(rb"^CREATE DATABASE\s+(.*)$")
+_CREATE_ROLE_LINE = re.compile(rb"^CREATE ROLE\s+(.*)$")
+_CONNECT = re.compile(rb"^\\connect\s+(.*)$")
+_NO_ROWS = "SELECT ''::text WHERE false"
+
+
+def _first_identifier(rest: str) -> str | None:
+    """The first SQL identifier in *rest*, quoted or bare.
+
+    A quoted identifier may hold spaces and doubled quotes, so it cannot be
+    read with a character class that stops at whitespace - which is how a
+    database called ``odd name`` used to leave the inventory as ``odd``.
+    """
+    text = rest.strip()
+    if not text:
+        return None
+    if text.startswith('"'):
+        out = []
+        index = 1
+        while index < len(text):
+            char = text[index]
+            if char == '"':
+                if index + 1 < len(text) and text[index + 1] == '"':
+                    out.append('"')
+                    index += 2
+                    continue
+                return "".join(out)
+            out.append(char)
+            index += 1
+        return None
+    return re.split(r"[\s;(]", text, maxsplit=1)[0] or None
+
+
+def _connect_target(rest: str) -> str | None:
+    """The database a ``\\connect`` line switches to.
+
+    psql options precede the name (``\\connect -reuse-previous=on dbname=x``),
+    and the name may arrive as a ``dbname=`` assignment rather than bare.
+    """
+    for token in rest.strip().split():
+        if token.startswith("-"):
+            continue
+        if token.startswith("dbname="):
+            return _first_identifier(token[len("dbname=") :])
+        return _first_identifier(rest.strip()[rest.strip().index(token) :])
+    return None
+
+
+def dump_inventory(sql_path: str) -> tuple[list[str], list[str]]:
+    """The databases and roles a cluster dump recreates.
+
+    Args:
+        sql_path: the ``pg_dumpall`` stream.
+
+    Returns:
+        ``(databases, roles)``, each in the order the dump names them. The
+        pre-clean is scoped to these: everything else in the instance belongs
+        to no backup this restore holds, and dropping it would destroy data
+        the replay cannot bring back.
+    """
+    databases: list[str] = []
+    roles: list[str] = []
+    with open(sql_path, "rb") as handle:
+        for raw in handle:
+            line = raw.decode("utf-8", "replace")
+            for pattern, sink, read in (
+                (_CREATE_DATABASE, databases, _first_identifier),
+                (_CONNECT, databases, _connect_target),
+                (_CREATE_ROLE_LINE, roles, _first_identifier),
+            ):
+                found = pattern.match(raw)
+                if not found:
+                    continue
+                name = read(line[found.start(1) :])
+                if name and name not in sink:
+                    sink.append(name)
+    return databases, roles
+
+
+def preclean_sql() -> str:
+    """The catalog-wide pre-clean, safe only behind the instance check."""
+    with open(_CLUSTER_PRECLEAN_SQL, encoding="utf-8") as preclean:
+        return preclean.read()
+
+
+def instance_databases(container: str, user: str, docker_env: dict) -> list[str]:
+    """The instance's own databases, templates and control database aside."""
+    listed = docker_exec(
+        container,
+        [
+            "psql",
+            "-U",
+            user,
+            "-d",
+            CONTROL_DB,
+            "-tAc",
+            (
+                "SELECT datname FROM pg_database "
+                "WHERE NOT datistemplate AND datname <> current_database()"
+            ),
+        ],
+        capture=True,
+        docker_env=docker_env,
+    ).stdout
+    text = listed.decode() if isinstance(listed, bytes) else listed
+    return [name for name in text.split() if name]
+
+
+def assert_instance_matches_dump(
+    container: str, user: str, sql_path: str, docker_env: dict
+) -> None:
+    """Refuse ``--empty`` on an instance holding anything the dump lacks.
+
+    The pre-clean is a catalog-wide sweep, so a foreign database would be
+    destroyed with no way back. Scoping the sweep instead is not a fix: a
+    surviving database that owns or grants to one of the dump's roles pins
+    that role in pg_shdepend, and DROP ROLE then fails after the dump's own
+    databases are already gone.
+
+    Raises:
+        RuntimeError: the instance carries databases this dump cannot restore.
+    """
+    dumped, _roles = dump_inventory(sql_path)
+    present = instance_databases(container, user, docker_env)
+    foreign = sorted(set(present) - set(dumped))
+    if foreign:
+        raise RuntimeError(
+            f"{container} also holds {', '.join(foreign)}, which "
+            f"{os.path.basename(sql_path)} does not carry. --empty wipes the "
+            "instance, so those would be destroyed with nothing to restore "
+            "them from. Move them off this instance, or drop them yourself if "
+            "they are disposable."
+        )
 
 
 def _psql(user: str) -> list[str]:
@@ -98,12 +231,11 @@ def restore_cluster_sql(
     docker_env = {"PGPASSWORD": password}
 
     if empty:
-        with open(_CLUSTER_PRECLEAN_SQL, encoding="utf-8") as preclean:
-            drop_sql = preclean.read()
+        assert_instance_matches_dump(container, user, sql_path, docker_env)
         docker_exec(
             container,
             _psql(user),
-            stdin=drop_sql.encode(),
+            stdin=preclean_sql().encode(),
             docker_env=docker_env,
         )
 

src/baudolo/restore/db/cluster_preclean.sql

@@ -8,6 +8,9 @@
 -- control database. Template databases, the control database itself, the pg_*
 -- system roles and the connecting role are kept - the dump does not recreate
 -- them and dropping them would end the session.
+-- The sweep stays catalog-wide on purpose: a scoped one leaves databases that
+-- pin a dumped role in pg_shdepend, and phase 3 then fails after phase 1 has
+-- already dropped. assert_instance_matches_dump refuses before this runs.
 SELECT statement
   FROM (
     SELECT 1 AS phase,

tests/e2e/test_e2e_postgres_cluster_restore.py

@@ -23,11 +23,11 @@ SEED_SQL = (
     "CREATE DATABASE first OWNER app",
     "CREATE DATABASE second OWNER app",
 )
-DROP_SQL = (
-    "DROP DATABASE first",
-    "DROP DATABASE second",
-    "DROP ROLE app",
+SIBLING_SQL = (
+    "CREATE ROLE neighbour LOGIN PASSWORD 'neighbourpw'",
+    "CREATE DATABASE sibling OWNER neighbour",
 )
+SIBLING_PAYLOAD = "CREATE TABLE t (v text); INSERT INTO t VALUES ('sibling-payload');"
 FIRST_SQL = "CREATE TABLE t (v text); INSERT INTO t VALUES ('first-payload');"
 SECOND_SQL = "CREATE TABLE t (v text); INSERT INTO t VALUES ('second-payload');"
 
@@ -89,9 +89,6 @@ class TestE2EPostgresClusterRestore(unittest.TestCase):
             / f"{cls.pg_container}.cluster.backup.sql"
         )
 
-        for statement in DROP_SQL:
-            cls._psql("postgres", statement)
-
         run(
             [
                 "baudolo-restore",
@@ -115,6 +112,34 @@ class TestE2EPostgresClusterRestore(unittest.TestCase):
             ]
         )
 
+        for statement in SIBLING_SQL:
+            cls._psql("postgres", statement)
+        cls._psql("sibling", SIBLING_PAYLOAD)
+        cls.refused = run(cls._restore_argv(), check=False)
+
+    @classmethod
+    def _restore_argv(cls) -> list:
+        return [
+            "baudolo-restore",
+            "cluster",
+            cls.pg_volume,
+            cls.hash,
+            cls.version,
+            "--backups-dir",
+            cls.backups_dir,
+            "--repo-name",
+            cls.repo_name,
+            "--container",
+            cls.pg_container,
+            "--instance",
+            cls.pg_container,
+            "--db-user",
+            "postgres",
+            "--db-password",
+            "pgpw",
+            "--empty",
+        ]
+
     @classmethod
     def tearDownClass(cls) -> None:
         cleanup_docker(containers=cls.containers, volumes=cls.volumes)
@@ -136,6 +161,19 @@ class TestE2EPostgresClusterRestore(unittest.TestCase):
     def test_the_backup_wrote_a_cluster_dump(self) -> None:
         self.assertTrue(self.dump.is_file(), f"no cluster dump at {self.dump}")
 
+    def test_the_preclean_really_dropped_a_populated_cluster(self) -> None:
+        self.assertEqual(self._psql("first", "SELECT v FROM t"), "first-payload")
+
+    def test_a_second_empty_is_refused_once_a_foreign_database_exists(self) -> None:
+        self.assertNotEqual(self.refused.returncode, 0, self.refused.stdout)
+        self.assertIn("sibling", self.refused.stderr)
+
+    def test_the_refusal_left_the_foreign_database_alone(self) -> None:
+        self.assertEqual(self._psql("sibling", "SELECT v FROM t"), "sibling-payload")
+
+    def test_the_refusal_dropped_nothing_of_its_own(self) -> None:
+        self.assertEqual(self._psql("first", "SELECT v FROM t"), "first-payload")
+
     def test_both_databases_are_back(self) -> None:
         listed = self._psql(
             "postgres",

tests/unit/restore/test_cluster_preclean_scope.py

@@ -0,0 +1,126 @@
+"""The pre-clean may only drop what the dump can bring back.
+
+A shared instance carries databases and roles from other applications, and a
+database created after the backup is in no dump at all. Dropping those would
+destroy data this restore cannot restore.
+"""
+
+import tempfile
+import unittest
+from pathlib import Path
+
+from baudolo.restore.db import cluster as cluster_mod
+
+DUMP = """--
+-- PostgreSQL database cluster dump
+--
+
+CREATE ROLE app;
+ALTER ROLE app WITH LOGIN;
+CREATE ROLE reporting;
+
+CREATE DATABASE appdb OWNER app;
+
+\\connect appdb
+
+CREATE TABLE t (v text);
+
+\\connect template1
+"""
+
+
+def dump_file(text: str) -> str:
+    path = Path(tempfile.mkdtemp()) / "central.cluster.backup.sql"
+    path.write_text(text, encoding="utf-8")
+    return str(path)
+
+
+class TestDumpInventory(unittest.TestCase):
+    def test_it_reads_databases_and_roles_the_dump_recreates(self) -> None:
+        databases, roles = cluster_mod.dump_inventory(dump_file(DUMP))
+        self.assertEqual(databases, ["appdb", "template1"])
+        self.assertEqual(roles, ["app", "reporting"])
+
+    def test_a_quoted_name_keeps_its_spaces(self) -> None:
+        databases, roles = cluster_mod.dump_inventory(
+            dump_file('\\connect "odd name"\nCREATE ROLE "odd role";\n')
+        )
+        self.assertEqual(databases, ["odd name"])
+        self.assertEqual(roles, ["odd role"])
+
+    def test_psql_options_are_not_mistaken_for_the_database(self) -> None:
+        databases, _roles = cluster_mod.dump_inventory(
+            dump_file("\\connect -reuse-previous=on dbname=appdb\n")
+        )
+        self.assertEqual(databases, ["appdb"])
+
+    def test_create_database_options_are_not_part_of_the_name(self) -> None:
+        databases, _roles = cluster_mod.dump_inventory(
+            dump_file("CREATE DATABASE appdb WITH TEMPLATE = template0 OWNER = app;\n")
+        )
+        self.assertEqual(databases, ["appdb"])
+
+    def test_a_name_is_listed_once(self) -> None:
+        databases, _roles = cluster_mod.dump_inventory(
+            dump_file("\\connect a\n\\connect a\n")
+        )
+        self.assertEqual(databases, ["a"])
+
+
+class TestInstanceRefusal(unittest.TestCase):
+    """--empty wipes the whole instance, so it may only run on one this dump
+    can rebuild. Scoping the sweep instead wedges the restore: a surviving
+    database that grants to a dumped role pins it, DROP ROLE fails, and the
+    pre-clean aborts after the dump's own databases are already gone."""
+
+    def check(self, present: str, dump: str = DUMP):
+        from unittest import mock
+
+        with mock.patch.object(
+            cluster_mod, "instance_databases", return_value=present.split()
+        ):
+            cluster_mod.assert_instance_matches_dump(
+                "db", "postgres", dump_file(dump), {}
+            )
+
+    def test_an_instance_the_dump_covers_passes(self) -> None:
+        self.check("appdb")
+
+    def test_an_empty_instance_passes(self) -> None:
+        self.check("")
+
+    def test_a_database_the_dump_lacks_is_refused(self) -> None:
+        with self.assertRaises(RuntimeError) as raised:
+            self.check("appdb sibling")
+        self.assertIn("sibling", str(raised.exception))
+
+    def test_the_refusal_names_every_foreign_database(self) -> None:
+        with self.assertRaises(RuntimeError) as raised:
+            self.check("one two")
+        message = str(raised.exception)
+        self.assertIn("one", message)
+        self.assertIn("two", message)
+
+    def test_the_refusal_happens_before_anything_is_dropped(self) -> None:
+        from unittest import mock
+
+        with (
+            mock.patch.object(
+                cluster_mod, "instance_databases", return_value=["foreign"]
+            ),
+            mock.patch.object(cluster_mod, "docker_exec") as touched,
+            self.assertRaises(RuntimeError),
+        ):
+            cluster_mod.restore_cluster_sql(
+                container="db",
+                user="postgres",
+                password="pw",
+                sql_path=dump_file(DUMP),
+                empty=True,
+                check_version=False,
+            )
+        touched.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/restore/test_cluster_replay.py

@@ -11,6 +11,8 @@ class TestClusterReplay(unittest.TestCase):
         calls = []
 
         def _capture(container, argv, **kwargs):
+            if "-tAc" in argv:
+                return MagicMock(stdout=b"")
             calls.append((argv, kwargs.get("stdin")))
             return MagicMock()
 
@@ -58,7 +60,7 @@ class TestClusterReplay(unittest.TestCase):
         self.assertIn("DROP OWNED BY", preclean)
         self.assertIn("ORDER BY phase", preclean)
 
-    def test_the_preclean_spares_what_the_dump_does_not_recreate(self) -> None:
+    def test_the_preclean_spares_what_no_dump_recreates(self) -> None:
         preclean = self._replay(empty=True)[0][1].decode()
         self.assertIn("NOT datistemplate", preclean)
         self.assertIn("datname <> current_database()", preclean)