diff --git a/src/baudolo/restore/__main__.py b/src/baudolo/restore/__main__.py index c68f324..b444812 100644 --- a/src/baudolo/restore/__main__.py +++ b/src/baudolo/restore/__main__.py @@ -3,10 +3,11 @@ from __future__ import annotations import argparse import sys -from .paths import BackupPaths -from .files import restore_volume_files -from .db.postgres import restore_postgres_sql +from .db.cluster import restore_cluster_sql from .db.mariadb import restore_mariadb_sql +from .db.postgres import restore_postgres_sql +from .files import restore_volume_files +from .paths import BackupPaths def _add_common_backup_args(p: argparse.ArgumentParser) -> None: @@ -59,6 +60,27 @@ def main(argv: list[str] | None = None) -> int: p_pg.add_argument("--db-password", required=True) p_pg.add_argument("--empty", action="store_true") + # ------------------------------------------------------------------ + # cluster + # ------------------------------------------------------------------ + p_cluster = sub.add_parser( + "cluster", help="Restore a full PostgreSQL cluster dump (pg_dumpall)" + ) + _add_common_backup_args(p_cluster) + p_cluster.add_argument("--container", required=True) + p_cluster.add_argument( + "--instance", + required=True, + help="Instance the dump was taken from; names .cluster.backup.sql", + ) + p_cluster.add_argument( + "--db-user", + required=True, + help="Superuser of the instance; the dump creates roles and databases", + ) + p_cluster.add_argument("--db-password", required=True) + p_cluster.add_argument("--empty", action="store_true") + # ------------------------------------------------------------------ # mariadb # ------------------------------------------------------------------ @@ -111,6 +133,22 @@ def main(argv: list[str] | None = None) -> int: ) return 0 + if args.cmd == "cluster": + restore_cluster_sql( + container=args.container, + user=args.db_user, + password=args.db_password, + sql_path=BackupPaths( + args.volume_name, + args.backup_hash, + args.version, + repo_name=args.repo_name, + backups_dir=args.backups_dir, + ).cluster_file(args.instance), + empty=args.empty, + ) + return 0 + if args.cmd == "mariadb": user = args.db_user or args.db_name restore_mariadb_sql( diff --git a/src/baudolo/restore/db/cluster.py b/src/baudolo/restore/db/cluster.py new file mode 100644 index 0000000..b3c7f59 --- /dev/null +++ b/src/baudolo/restore/db/cluster.py @@ -0,0 +1,103 @@ +"""Replay a full PostgreSQL cluster dump produced by ``pg_dumpall``. + +The backup side writes one when a databases.csv row asks for every database of +an instance (``database = '*'``, see ``backup/db.py``). Until now nothing read +it back, so that dump was stored and unrestorable - a format whose producer has +no consumer. + +A cluster stream differs from a single-database one in three ways that decide +the implementation: + +* it recreates roles and databases, so it must be replayed against the control + database rather than into a target database; +* ``CREATE DATABASE`` cannot run inside a transaction block, so unlike + :mod:`baudolo.restore.db.postgres` the replay must not be wrapped in + ``--single-transaction``; +* it is replayed as a superuser, so the superuser-only statements that the + single-database path filters out are exactly the ones that have to survive. +""" + +from __future__ import annotations + +import os +import re +import tempfile +from collections.abc import Iterable, Iterator + +from ..run import docker_exec + +CONTROL_DB = "postgres" +_CLUSTER_PRECLEAN_SQL = os.path.join(os.path.dirname(__file__), "cluster_preclean.sql") +_CREATE_ROLE = re.compile(rb'^CREATE ROLE "?([^";]+)"?;\s*$') + + +def _psql(user: str) -> list[str]: + """The replay client: no --single-transaction, CREATE DATABASE forbids it.""" + return ["psql", "-v", "ON_ERROR_STOP=1", "-U", user, "-d", CONTROL_DB] + + +def filter_own_role_creation(lines: Iterable[bytes], user: str) -> Iterator[bytes]: + """Drop the ``CREATE ROLE`` of the role holding this session. + + A pg_dumpall stream recreates every role of the cluster, the bootstrap + superuser included, and the pre-clean cannot drop the one it is connected + as - so that single statement always collides. Its ``ALTER ROLE`` is kept: + that is what re-applies the attributes and the password the dump captured. + + Args: + lines: dump lines including their trailing newlines. + user: the connecting role. + + Yields: + Every line except that one CREATE. + """ + for line in lines: + found = _CREATE_ROLE.match(line) + if found and found.group(1).decode() == user: + continue + yield line + + +def restore_cluster_sql( + *, + container: str, + user: str, + password: str, + sql_path: str, + empty: bool, +) -> None: + """Replay a pg_dumpall stream into a running instance. + + Args: + container: the running engine to replay into. + user: a superuser of that instance; the dump creates roles and + databases, which an application role may not do. + password: its password, handed to psql through the container's env. + sql_path: the ``.cluster.backup.sql`` of a generation. + empty: drop the cluster's databases and roles first. Without it the + replay stops at the first object that already exists, which is the + honest outcome: recreating a cluster over a populated one is a + decision, not a default. + """ + if not os.path.isfile(sql_path): + raise FileNotFoundError(sql_path) + + docker_env = {"PGPASSWORD": password} + + if empty: + with open(_CLUSTER_PRECLEAN_SQL, encoding="utf-8") as preclean: + drop_sql = preclean.read() + docker_exec( + container, + _psql(user), + stdin=drop_sql.encode(), + docker_env=docker_env, + ) + + with open(sql_path, "rb") as src, tempfile.TemporaryFile() as filtered: + for line in filter_own_role_creation(src, user): + filtered.write(line) + filtered.seek(0) + docker_exec(container, _psql(user), stdin=filtered, docker_env=docker_env) + + print(f"PostgreSQL cluster restore complete from '{os.path.basename(sql_path)}'.") diff --git a/src/baudolo/restore/db/cluster_preclean.sql b/src/baudolo/restore/db/cluster_preclean.sql new file mode 100644 index 0000000..caf862c --- /dev/null +++ b/src/baudolo/restore/db/cluster_preclean.sql @@ -0,0 +1,30 @@ +-- Pre-clean for `restore cluster --empty`. A pg_dumpall stream recreates roles +-- and databases, so replaying it into a populated cluster dies on the first +-- CREATE ROLE. Emitted as one DROP per row and run via \gexec so each executes +-- as its own top-level statement: DROP DATABASE cannot run inside a +-- transaction block, which rules out a single DO block. +-- The phase column pins the order: databases must be gone before their owners +-- can be dropped, and DROP OWNED BY releases what a role still holds in the +-- 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. +SELECT statement + FROM ( + SELECT 1 AS phase, + format('DROP DATABASE IF EXISTS %I', datname) AS statement + FROM pg_database + WHERE NOT datistemplate + AND datname <> current_database() + UNION ALL + SELECT 2, format('DROP OWNED BY %I', rolname) + FROM pg_roles + WHERE NOT starts_with(rolname, 'pg_') + AND rolname <> current_user + UNION ALL + SELECT 3, format('DROP ROLE IF EXISTS %I', rolname) + FROM pg_roles + WHERE NOT starts_with(rolname, 'pg_') + AND rolname <> current_user + ) drops + ORDER BY phase +\gexec diff --git a/src/baudolo/restore/paths.py b/src/baudolo/restore/paths.py index a037321..536b11e 100644 --- a/src/baudolo/restore/paths.py +++ b/src/baudolo/restore/paths.py @@ -27,3 +27,7 @@ class BackupPaths: def sql_file(self, db_name: str) -> str: return os.path.join(self.root(), "sql", f"{db_name}.backup.sql") + + def cluster_file(self, instance: str) -> str: + """The pg_dumpall stream a `database = '*'` row produces.""" + return os.path.join(self.root(), "sql", f"{instance}.cluster.backup.sql") diff --git a/tests/e2e/test_e2e_postgres_cluster_restore.py b/tests/e2e/test_e2e_postgres_cluster_restore.py new file mode 100644 index 0000000..5d3a904 --- /dev/null +++ b/tests/e2e/test_e2e_postgres_cluster_restore.py @@ -0,0 +1,188 @@ +import unittest + +from .helpers import ( + POSTGRES_DATA_DIR, + POSTGRES_IMAGE, + backup_path, + backup_run, + cleanup_docker, + create_minimal_compose_dir, + ensure_empty_dir, + latest_version_dir, + require_docker, + run, + unique, + wait_for_postgres, + write_databases_csv, +) + +# A `database = '*'` row makes the backup side write one pg_dumpall stream for +# the whole instance instead of a dump per database - the shape an application +# with several databases in one engine produces. This proves the stream is +# replayable: two databases and their owning role are dropped outright, and the +# cluster restore has to bring all three back. Before the cluster subcommand +# existed the dump was stored and unreadable. +# Each statement runs on its own: psql wraps a multi-statement -c in one +# transaction, and CREATE DATABASE is forbidden inside one - the same rule that +# keeps the cluster replay out of --single-transaction. +SEED_SQL = ( + "CREATE ROLE app LOGIN PASSWORD 'apppw'", + "CREATE DATABASE first OWNER app", + "CREATE DATABASE second OWNER app", +) +DROP_SQL = ( + "DROP DATABASE first", + "DROP DATABASE second", + "DROP ROLE app", +) +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');" + + +class TestE2EPostgresClusterRestore(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + require_docker() + cls.prefix = unique("baudolo-e2e-pg-cluster") + 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.containers = [cls.pg_container] + cls.volumes = [cls.pg_volume] + + run(["docker", "volume", "create", cls.pg_volume]) + run( + [ + "docker", + "run", + "-d", + "--name", + cls.pg_container, + "-e", + "POSTGRES_PASSWORD=pgpw", + "-v", + f"{cls.pg_volume}:{POSTGRES_DATA_DIR}", + POSTGRES_IMAGE, + ] + ) + wait_for_postgres(cls.pg_container, user="postgres") + + for statement in SEED_SQL: + cls._psql("postgres", statement) + cls._psql("first", FIRST_SQL) + cls._psql("second", SECOND_SQL) + + cls.databases_csv = f"/tmp/{cls.prefix}/databases.csv" + write_databases_csv( + cls.databases_csv, [(cls.pg_container, "*", "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) + cls.dump = ( + backup_path(cls.backups_dir, cls.repo_name, cls.version, cls.pg_volume) + / "sql" + / f"{cls.pg_container}.cluster.backup.sql" + ) + + # The disaster: both databases and the role that owns them are gone. + for statement in DROP_SQL: + cls._psql("postgres", statement) + + run( + [ + "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) + + @classmethod + def _psql(cls, database: str, sql: str) -> str: + p = run( + [ + "docker", + "exec", + cls.pg_container, + "sh", + "-lc", + f'psql -U postgres -d {database} -t -A -c "{sql}"', + ] + ) + return (p.stdout or "").strip() + + 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_both_databases_are_back(self) -> None: + listed = self._psql( + "postgres", + "SELECT datname FROM pg_database WHERE datname IN ('first','second') ORDER BY 1", + ) + self.assertEqual(listed.split(), ["first", "second"]) + + def test_each_database_carries_its_own_payload(self) -> None: + self.assertEqual(self._psql("first", "SELECT v FROM t"), "first-payload") + self.assertEqual(self._psql("second", "SELECT v FROM t"), "second-payload") + + def test_the_superusers_own_create_was_filtered(self) -> None: + # The dump recreates every role including the one the replay connects + # as; only its ALTER may survive, or the stream dies on the first + # statement with ON_ERROR_STOP. + self.assertEqual( + self._psql( + "postgres", "SELECT rolsuper FROM pg_roles WHERE rolname = 'postgres'" + ), + "t", + ) + + def test_the_owning_role_is_back(self) -> None: + self.assertEqual( + self._psql( + "postgres", "SELECT rolname FROM pg_roles WHERE rolname = 'app'" + ), + "app", + ) + + def test_ownership_survived(self) -> None: + self.assertEqual( + self._psql( + "postgres", + "SELECT pg_get_userbyid(datdba) FROM pg_database WHERE datname = 'first'", + ), + "app", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/restore/test_cluster_replay.py b/tests/unit/restore/test_cluster_replay.py new file mode 100644 index 0000000..64cee35 --- /dev/null +++ b/tests/unit/restore/test_cluster_replay.py @@ -0,0 +1,124 @@ +import tempfile +import unittest +from unittest.mock import MagicMock, patch + +from baudolo.restore.db import cluster as cluster_mod +from baudolo.restore.paths import BackupPaths + + +class TestClusterReplay(unittest.TestCase): + def _replay(self, *, empty: bool): + calls = [] + + def _capture(container, argv, **kwargs): + calls.append((argv, kwargs.get("stdin"))) + return MagicMock() + + with tempfile.NamedTemporaryFile(suffix=".sql") as sql: + sql.write(b"CREATE ROLE app;\nCREATE DATABASE app OWNER app;\n") + sql.flush() + with patch.object(cluster_mod, "docker_exec", side_effect=_capture): + cluster_mod.restore_cluster_sql( + container="db", + user="postgres", + password="pw", + sql_path=sql.name, + empty=empty, + ) + return calls + + def test_the_replay_is_not_wrapped_in_a_transaction(self) -> None: + argv, _ = self._replay(empty=False)[0] + self.assertNotIn( + "--single-transaction", + argv, + "CREATE DATABASE cannot run inside a transaction block, so unlike the " + "single-database replay this stream must not be wrapped in one", + ) + self.assertIn("ON_ERROR_STOP=1", argv) + + def test_the_replay_targets_the_control_database(self) -> None: + argv, _ = self._replay(empty=False)[0] + self.assertEqual(argv[argv.index("-d") + 1], cluster_mod.CONTROL_DB) + self.assertEqual(argv[argv.index("-U") + 1], "postgres") + + def test_without_empty_nothing_is_dropped_first(self) -> None: + self.assertEqual(len(self._replay(empty=False)), 1) + + def test_empty_drops_databases_before_their_owners(self) -> None: + calls = self._replay(empty=True) + self.assertEqual(len(calls), 2, f"expected pre-clean + replay: {calls}") + preclean = calls[0][1].decode() + self.assertLess( + preclean.index("DROP DATABASE"), + preclean.index("DROP ROLE"), + "a role cannot be dropped while it still owns a database", + ) + 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: + preclean = self._replay(empty=True)[0][1].decode() + self.assertIn("NOT datistemplate", preclean) + self.assertIn("datname <> current_database()", preclean) + self.assertIn("starts_with(rolname, 'pg_')", preclean) + self.assertIn("rolname <> current_user", preclean) + + def test_only_the_connecting_role_loses_its_create(self) -> None: + # Captured from pg_dumpall 17: the bootstrap superuser is recreated like + # any other role, and the pre-clean cannot drop the one holding the + # session - so that single CREATE always collides while its ALTER, which + # carries the attributes and the password, must survive. + dump = [ + b"CREATE ROLE app;\n", + b"ALTER ROLE app WITH NOSUPERUSER INHERIT LOGIN PASSWORD 'SCRAM-SHA-256$...';\n", + b"CREATE ROLE postgres;\n", + b"ALTER ROLE postgres WITH SUPERUSER INHERIT LOGIN PASSWORD 'SCRAM-SHA-256$...';\n", + b'CREATE ROLE "odd-name";\n', + ] + kept = list(cluster_mod.filter_own_role_creation(dump, "postgres")) + self.assertNotIn(b"CREATE ROLE postgres;\n", kept) + self.assertIn(b"CREATE ROLE app;\n", kept) + self.assertIn(b'CREATE ROLE "odd-name";\n', kept) + self.assertEqual( + sum(1 for line in kept if line.startswith(b"ALTER ROLE postgres")), + 1, + "the ALTER re-applies the superuser's attributes and password", + ) + + def test_a_quoted_connecting_role_is_matched_too(self) -> None: + kept = list( + cluster_mod.filter_own_role_creation( + [b'CREATE ROLE "odd-name";\n'], "odd-name" + ) + ) + self.assertEqual(kept, []) + + def test_a_role_whose_name_merely_starts_the_same_is_kept(self) -> None: + kept = list( + cluster_mod.filter_own_role_creation( + [b"CREATE ROLE postgresql;\n"], "postgres" + ) + ) + self.assertEqual(kept, [b"CREATE ROLE postgresql;\n"]) + + def test_a_missing_dump_is_reported_as_such(self) -> None: + with self.assertRaises(FileNotFoundError): + cluster_mod.restore_cluster_sql( + container="db", + user="postgres", + password="pw", + sql_path="/nonexistent/x.cluster.backup.sql", + empty=False, + ) + + def test_the_path_helper_names_the_dumpall_file(self) -> None: + paths = BackupPaths("vol", "hash", "v1", repo_name="repo", backups_dir="/B") + self.assertEqual( + paths.cluster_file("bigbluebutton"), + "/B/hash/repo/v1/vol/sql/bigbluebutton.cluster.backup.sql", + ) + + +if __name__ == "__main__": + unittest.main()