import argparse import hashlib import json import subprocess from pathlib import Path import pytest from ops import restore_gcp_generated_postgres_snapshot as restore SERVICE_STATE = { "ActiveState": "active", "SubState": "running", "MainPID": "148735", "NRestarts": "0", "ExecMainStartTimestamp": "Thu 2026-07-09 07:00:04 UTC", } def manifest_rows(database: str, *, target: bool, row_count: int = 1) -> list[dict]: identity = { "kind": "identity", "database": database, "server_version_num": 160014, "transaction_read_only": "on", "server_address": "10.61.0.3" if target else None, "server_port": 5432 if target else None, "ssl": target, "ssl_version": "TLSv1.3" if target else None, "database_collation": "en_US.UTF8" if target else "en_US.utf8", "database_ctype": "en_US.UTF8" if target else "en_US.utf8", "server_encoding": "UTF8", } rows = [ identity, {"kind": "schemas", "items": ["kb_stage", "public"]}, {"kind": "extensions", "items": [{"name": "pgcrypto", "version": "1.3"}]}, { "kind": "application_roles", "items": [ { "name": "kb_apply", "can_login": True, "superuser": False, "create_db": False, "create_role": False, "inherit": True, "replication": False, "bypass_rls": False, } ], }, ] rows.extend( {"kind": kind, "items": []} for kind in ( "columns", "constraints", "indexes", "sequences", "views", "functions", "triggers", "types", "policies", ) ) rows.extend( [ { "kind": "table", "schema": "public", "table": "claims", "row_count": row_count, "rowset_md5": "same" if row_count == 1 else "different", }, { "kind": "performance", "query": "count_claims", "elapsed_ms": 1.0, "observed_rows": row_count, }, ] ) return rows def manifest_text(database: str, *, target: bool, row_count: int = 1) -> str: return "".join(json.dumps(row) + "\n" for row in manifest_rows(database, target=target, row_count=row_count)) def args_for(tmp_path: Path, **overrides) -> argparse.Namespace: dump = tmp_path / "teleo-canonical.dump" dump.write_bytes(b"PGDMPfixture") source_manifest = tmp_path / "source-manifest.jsonl" source_manifest.write_text(manifest_text("teleo", target=False), encoding="utf-8") values = { "operation": "restore", "execute": True, "target_db": "teleo_clone_dbfirst_test", "run_id": "gcp-restore-test12345", "host": restore.DEFAULT_HOST, "project": restore.DEFAULT_PROJECT, "password_secret": restore.DEFAULT_SECRET, "service": restore.DEFAULT_SERVICE, "command_timeout": 2.0, "psql_bin": "psql", "run_root": tmp_path / "runs", "dump": dump, "expected_dump_sha256": hashlib.sha256(dump.read_bytes()).hexdigest(), "source_manifest": source_manifest, "manifest_sql": Path("ops/postgres_parity_manifest.sql").resolve(), "pg_restore_bin": "pg_restore", "restore_timeout": 5.0, "manifest_timeout": 5.0, "max_target_query_ms": 2000.0, "max_target_source_ratio": 20.0, } values.update(overrides) return argparse.Namespace(**values) def install_restore_fakes( monkeypatch: pytest.MonkeyPatch, *, target_manifest: str, ) -> dict[str, bool]: state = {"exists": False, "dropped": False} monkeypatch.setattr(restore.os, "geteuid", lambda: 0) monkeypatch.setattr(restore, "service_state", lambda *_args, **_kwargs: dict(SERVICE_STATE)) monkeypatch.setattr(restore, "resolve_password", lambda *_args, **_kwargs: "private-password") monkeypatch.setattr( restore, "validate_pg_restore_version", lambda *_args, **_kwargs: { "pg_restore_bin": "pg_restore", "pg_restore_version": "pg_restore (PostgreSQL) 16.14", "pg_restore_major": 16, "source_postgres_major": 16, "compatible": True, }, ) monkeypatch.setattr(restore, "database_exists", lambda *_args, **_kwargs: state["exists"]) def fake_create(*_args, **_kwargs): state["exists"] = True def fake_drop(*_args, **_kwargs): existed = state["exists"] state["exists"] = False state["dropped"] = True return {"existed_before": existed, "clone_database_remaining": 0} monkeypatch.setattr(restore, "create_database", fake_create) monkeypatch.setattr(restore, "restore_dump", lambda *_args, **_kwargs: None) monkeypatch.setattr(restore, "capture_target_manifest", lambda *_args, **_kwargs: target_manifest) monkeypatch.setattr( restore, "rollback_database_state", lambda *_args, **_kwargs: {"exists": True, "datallowconn": False, "connections": 0}, ) monkeypatch.setattr(restore, "drop_clone", fake_drop) return state def test_restore_success_retains_only_exact_parity_clone( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: args = args_for(tmp_path) state = install_restore_fakes( monkeypatch, target_manifest=manifest_text(args.target_db, target=True), ) result = restore.run_restore(args) assert result["status"] == "pass" assert result["phase"] == "retained_for_no_send_replay" assert result["parity"]["problems"] == [] assert result["parity"]["details"]["source_total_rows"] == 1 assert result["target"]["connectivity"]["private_connectivity"] is True assert result["target"]["connectivity"]["ssl"] is True assert result["live_service"]["unchanged"] is True assert result["cleanup"]["clone_database_remaining"] == 1 assert state == {"exists": True, "dropped": False} run_dir = args.run_root / args.run_id assert (run_dir / "target-manifest.jsonl").stat().st_mode & 0o777 == 0o600 assert (run_dir / "restore-receipt.json").stat().st_mode & 0o777 == 0o600 def test_parity_failure_drops_generated_clone_and_receipts_both_failures( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: args = args_for(tmp_path) state = install_restore_fakes( monkeypatch, target_manifest=manifest_text(args.target_db, target=True, row_count=2), ) result = restore.run_restore(args) assert result["status"] == "fail" assert result["error"]["phase"] == "parity_compare" assert result["parity"]["problems"] == ["table_content_mismatches=1"] assert result["cleanup"]["attempted"] is True assert result["cleanup"]["clone_database_remaining"] == 0 assert state == {"exists": False, "dropped": True} def test_create_command_failure_still_drops_a_committed_clone( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: args = args_for(tmp_path) state = install_restore_fakes( monkeypatch, target_manifest=manifest_text(args.target_db, target=True), ) def committed_then_failed(*_args, **_kwargs): state["exists"] = True raise restore.RestoreError("connection dropped after create committed") monkeypatch.setattr(restore, "create_database", committed_then_failed) result = restore.run_restore(args) assert result["status"] == "fail" assert result["error"]["phase"] == "database_create" assert result["cleanup"]["attempted"] is True assert result["cleanup"]["clone_database_remaining"] == 0 assert state == {"exists": False, "dropped": True} def test_changed_live_service_turns_successful_restore_into_cleanup( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: args = args_for(tmp_path) state = install_restore_fakes( monkeypatch, target_manifest=manifest_text(args.target_db, target=True), ) calls = 0 def changing_service(*_args, **_kwargs): nonlocal calls calls += 1 result = dict(SERVICE_STATE) if calls > 1: result["NRestarts"] = "1" return result monkeypatch.setattr(restore, "service_state", changing_service) result = restore.run_restore(args) assert result["status"] == "fail" assert result["error"]["phase"] == "service_and_rollback_readback" assert result["live_service"]["unchanged"] is False assert result["cleanup"]["clone_database_remaining"] == 0 assert state["dropped"] is True def test_cleanup_is_bound_to_passing_restore_receipt( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: args = args_for(tmp_path) state = install_restore_fakes( monkeypatch, target_manifest=manifest_text(args.target_db, target=True), ) assert restore.run_restore(args)["status"] == "pass" cleanup_args = argparse.Namespace( **{ key: getattr(args, key) for key in ( "target_db", "run_id", "host", "project", "password_secret", "service", "command_timeout", "psql_bin", "run_root", ) }, operation="cleanup", execute=True, ) result = restore.run_cleanup(cleanup_args) assert result["status"] == "pass" assert result["database"] == {"existed_before": True, "clone_database_remaining": 0} assert result["rollback_database"] == {"exists": True, "datallowconn": False, "connections": 0} assert result["live_service"]["unchanged"] is True assert state == {"exists": False, "dropped": True} def test_restore_dump_uses_tls_no_owner_no_acl_and_keeps_secret_out_of_argv( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: args = args_for(tmp_path) captured = {} def fake_run(command, *, timeout, env=None, input_text=None): captured.update(command=command, timeout=timeout, env=env, input_text=input_text) return subprocess.CompletedProcess(command, 0, "", "") monkeypatch.setattr(restore, "_run", fake_run) restore.restore_dump(args, "private-password") assert captured["command"][0] == "pg_restore" assert captured["command"][1] == "-d" assert "sslmode=require" in captured["command"][2] assert "--no-owner" in captured["command"] assert "--no-acl" in captured["command"] assert "--exit-on-error" in captured["command"] assert all("private-password" not in argument for argument in captured["command"]) assert captured["env"]["PGPASSWORD"] == "private-password" def test_pg_restore_preflight_rejects_client_older_than_snapshot( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: manifest_path = tmp_path / "source.jsonl" manifest_path.write_text(manifest_text("teleo", target=False), encoding="utf-8") source_manifest = restore.load_manifest(manifest_path) def fake_run(command, *, timeout, env=None, input_text=None): assert command == ["pg_restore", "--version"] assert timeout > 0 assert env is None assert input_text is None return subprocess.CompletedProcess(command, 0, "pg_restore (PostgreSQL) 15.18\n", "") monkeypatch.setattr(restore, "_run", fake_run) with pytest.raises(restore.RestoreError, match="major 15 is older than source PostgreSQL major 16"): restore.validate_pg_restore_version("pg_restore", source_manifest, timeout=2.0) @pytest.mark.parametrize("target", ["teleo_canonical", "teleo_clone_UPPER", "other_clone_test"]) def test_parse_args_rejects_non_disposable_target(target: str) -> None: with pytest.raises(SystemExit): restore.parse_args( [ "cleanup", "--execute", "--target-db", target, "--run-id", "gcp-restore-test12345", ] ) def test_parse_args_requires_explicit_execute() -> None: with pytest.raises(SystemExit): restore.parse_args( [ "cleanup", "--target-db", "teleo_clone_test", "--run-id", "gcp-restore-test12345", ] )