From 5f3ee0a6693a0843d964df8f6426d9325af746cc Mon Sep 17 00:00:00 2001 From: Kevin Veen-Birkenbach Date: Mon, 17 Aug 2026 08:15:58 +0200 Subject: [PATCH] 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) --- src/baudolo/restore/db/cluster.py | 138 +++++++++++++++++- src/baudolo/restore/db/cluster_preclean.sql | 3 + .../e2e/test_e2e_postgres_cluster_restore.py | 52 ++++++- .../restore/test_cluster_preclean_scope.py | 126 ++++++++++++++++ tests/unit/restore/test_cluster_replay.py | 4 +- 5 files changed, 312 insertions(+), 11 deletions(-) create mode 100644 tests/unit/restore/test_cluster_preclean_scope.py diff --git a/src/baudolo/restore/db/cluster.py b/src/baudolo/restore/db/cluster.py index cb6bc7e..bf1ace8 100644 --- a/src/baudolo/restore/db/cluster.py +++ b/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, ) diff --git a/src/baudolo/restore/db/cluster_preclean.sql b/src/baudolo/restore/db/cluster_preclean.sql index caf862c..51868d2 100644 --- a/src/baudolo/restore/db/cluster_preclean.sql +++ b/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, diff --git a/tests/e2e/test_e2e_postgres_cluster_restore.py b/tests/e2e/test_e2e_postgres_cluster_restore.py index 113c3ef..c011cac 100644 --- a/tests/e2e/test_e2e_postgres_cluster_restore.py +++ b/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", diff --git a/tests/unit/restore/test_cluster_preclean_scope.py b/tests/unit/restore/test_cluster_preclean_scope.py new file mode 100644 index 0000000..bb38a11 --- /dev/null +++ b/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() diff --git a/tests/unit/restore/test_cluster_replay.py b/tests/unit/restore/test_cluster_replay.py index b91e66a..0b49b7a 100644 --- a/tests/unit/restore/test_cluster_replay.py +++ b/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)