import argparse import hashlib import json import subprocess from pathlib import Path import pytest from ops import restore_gcp_generated_postgres_snapshot as restore from ops.capture_vps_canonical_postgres_snapshot import build_provenance_binding SERVICE_STATE = { "ActiveState": "active", "SubState": "running", "MainPID": "148735", "NRestarts": "0", "ExecMainStartTimestamp": "Thu 2026-07-09 07:00:04 UTC", } AUTHORIZATION_REF = "codex-delegation:019f47af-f2cc-7c10-b928-a9283afa2621" 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") manifest_sql = Path("ops/postgres_parity_manifest.sql").resolve() source_context_sha256 = "4" * 64 snapshot = { "exported_snapshot_id": "00000003-000001A2-1", "txid_snapshot": "100:120:105,111", "wal_lsn": "0/16B6C50", "system_identifier": "7649789040005668902", "captured_at_utc": "2026-07-15T01:02:03.456789+00:00", } source = { "ssh_target": "root@77.42.65.182", "container": "teleo-pg", "database": "teleo", "service": "leoclean-gateway.service", } dump_sha256 = hashlib.sha256(dump.read_bytes()).hexdigest() source_manifest_sha256 = hashlib.sha256(source_manifest.read_bytes()).hexdigest() manifest_sql_sha256 = hashlib.sha256(manifest_sql.read_bytes()).hexdigest() source_receipt = tmp_path / "source-receipt.json" source_receipt.write_text( json.dumps( { "artifact": "vps_canonical_postgres_exported_snapshot", "schema": restore.SOURCE_SNAPSHOT_RECEIPT_SCHEMA, "receipt_version": 2, "generated_at_utc": "2026-07-15T01:02:04+00:00", "status": "pass", "run_id": "canonical-20260715", "capture_authorization_ref": AUTHORIZATION_REF, "source": source, "snapshot_exported": True, "snapshot": snapshot, "dump": {"path": "/private/source.dump", "bytes": dump.stat().st_size, "sha256": dump_sha256}, "manifest": { "path": "/private/source-manifest.jsonl", "sha256": source_manifest_sha256, "manifest_sql_sha256": manifest_sql_sha256, "table_count": 1, "total_rows": 1, }, "source_context": {"path": "/private/source-context.txt", "sha256": source_context_sha256}, "source_service": {"before": SERVICE_STATE, "after": SERVICE_STATE, "unchanged": True}, "provenance_binding": build_provenance_binding( run_id="canonical-20260715", authorization_ref=AUTHORIZATION_REF, source=source, snapshot=snapshot, dump_sha256=dump_sha256, manifest_sql_sha256=manifest_sql_sha256, source_manifest_sha256=source_manifest_sha256, source_context_sha256=source_context_sha256, ), "service_unchanged": True, "production_db_mutated": False, "telegram_send_attempted": 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, "cloudsql_instance": restore.DEFAULT_INSTANCE, "expected_private_network": restore.DEFAULT_PRIVATE_NETWORK, "expected_compute_instance": restore.DEFAULT_COMPUTE_INSTANCE, "expected_compute_zone": restore.DEFAULT_COMPUTE_ZONE, "expected_source_ssh_target": restore.DEFAULT_SOURCE_SSH_TARGET, "expected_source_container": restore.DEFAULT_SOURCE_CONTAINER, "expected_source_database": restore.DEFAULT_SOURCE_DATABASE, "expected_source_service": restore.DEFAULT_SOURCE_SERVICE, "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": dump_sha256, "source_manifest": source_manifest, "source_receipt": source_receipt, "expected_source_receipt_sha256": hashlib.sha256(source_receipt.read_bytes()).hexdigest(), "expected_capture_authorization_ref": AUTHORIZATION_REF, "manifest_sql": manifest_sql, "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, "gce_compute_identity", lambda project, instance, zone, expected_private_network, **_kwargs: { "project": project, "project_number": "785938879453", "instance": instance, "instance_id": "123456789", "zone": zone, "expected_private_network": expected_private_network, "private_ip": "10.61.0.2", "network_resource": "projects/123456789/networks/teleo-staging-net", "service_account": "teleo-runtime@teleo-501523.iam.gserviceaccount.com", "checks": { "project": True, "project_number": True, "expected_network_project": True, "instance": True, "instance_id": True, "zone": True, "network": True, "private_ip": True, "service_account": True, }, }, ) monkeypatch.setattr( restore, "cloudsql_control_plane_state", lambda project, instance, host, expected_private_network, **_kwargs: { "project": project, "instance": instance, "private_addresses": [host], "expected_private_host": host, "private_network": expected_private_network, "expected_private_network": expected_private_network, "public_ip_disabled": True, "checks": { "instance": True, "state": True, "postgres_16": True, "public_ip_disabled": True, "private_network": True, "expected_private_host": True, "no_non_private_address": True, }, }, ) 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["source"]["capture_receipt"]["schema"] == restore.SOURCE_SNAPSHOT_RECEIPT_SCHEMA assert result["source"]["capture_receipt"]["provenance_binding_sha256"] assert result["target"]["connectivity"]["private_connectivity"] is True assert result["target"]["cloudsql"]["public_ip_disabled"] 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 connectivity_path = run_dir / "gcp-private-connectivity.json" assert connectivity_path.stat().st_mode & 0o777 == 0o600 connectivity = json.loads(connectivity_path.read_text()) assert connectivity["target_database"] == args.target_db assert connectivity["source_compute"] == restore.DEFAULT_COMPUTE_INSTANCE assert connectivity["public_ip_disabled"] is True assert ( result["target"]["connectivity_proof"]["sha256"] == hashlib.sha256(connectivity_path.read_bytes()).hexdigest() ) 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_restore_rejects_manifest_not_bound_to_source_receipt_before_credentials( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: args = args_for(tmp_path) install_restore_fakes(monkeypatch, target_manifest=manifest_text(args.target_db, target=True)) receipt = json.loads(args.source_receipt.read_text()) receipt["manifest"]["sha256"] = "0" * 64 args.source_receipt.write_text(json.dumps(receipt)) args.expected_source_receipt_sha256 = hashlib.sha256(args.source_receipt.read_bytes()).hexdigest() credential_called = False def unexpected_credential(*_args, **_kwargs): nonlocal credential_called credential_called = True return "private-password" monkeypatch.setattr(restore, "resolve_password", unexpected_credential) result = restore.run_restore(args) assert result["status"] == "fail" assert result["error"]["phase"] == "validation" assert "source_manifest_sha256" in result["error"]["message"] assert credential_called is False def test_restore_rejects_unexpected_capture_authorization_before_credentials( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: args = args_for(tmp_path, expected_capture_authorization_ref="codex-delegation:another-authorized-run") install_restore_fakes(monkeypatch, target_manifest=manifest_text(args.target_db, target=True)) credential_called = False def unexpected_credential(*_args, **_kwargs): nonlocal credential_called credential_called = True return "private-password" monkeypatch.setattr(restore, "resolve_password", unexpected_credential) result = restore.run_restore(args) assert result["status"] == "fail" assert result["error"]["phase"] == "validation" assert "capture_authorization_ref" in result["error"]["message"] assert credential_called is False def test_restore_rejects_noncanonical_source_identity_before_credentials( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: args = args_for(tmp_path, expected_source_container="different-postgres") install_restore_fakes(monkeypatch, target_manifest=manifest_text(args.target_db, target=True)) credential_called = False def unexpected_credential(*_args, **_kwargs): nonlocal credential_called credential_called = True return "private-password" monkeypatch.setattr(restore, "resolve_password", unexpected_credential) result = restore.run_restore(args) assert result["status"] == "fail" assert result["error"]["phase"] == "validation" assert "source_identity" in result["error"]["message"] assert credential_called is False 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", "cloudsql_instance", "expected_private_network", "expected_compute_instance", "expected_compute_zone", "expected_source_ssh_target", "expected_source_container", "expected_source_database", "expected_source_service", "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_cleanup_rejects_target_topology_drift_before_credentials( 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", "cloudsql_instance", "expected_private_network", "expected_compute_instance", "expected_compute_zone", "expected_source_ssh_target", "expected_source_container", "expected_source_database", "expected_source_service", "password_secret", "service", "command_timeout", "psql_bin", "run_root", ) }, operation="cleanup", execute=True, ) cleanup_args.host = "10.61.0.4" credential_called = False def unexpected_credential(*_args, **_kwargs): nonlocal credential_called credential_called = True return "private-password" monkeypatch.setattr(restore, "resolve_password", unexpected_credential) with pytest.raises(restore.RestoreError, match="execution contract"): restore.run_cleanup(cleanup_args) assert credential_called is False assert state == {"exists": True, "dropped": False} def test_cleanup_retains_failure_receipt_when_post_drop_readback_fails( 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", "cloudsql_instance", "expected_private_network", "expected_compute_instance", "expected_compute_zone", "expected_source_ssh_target", "expected_source_container", "expected_source_database", "expected_source_service", "password_secret", "service", "command_timeout", "psql_bin", "run_root", ) }, operation="cleanup", execute=True, ) monkeypatch.setattr( restore, "rollback_database_state", lambda *_args, **_kwargs: (_ for _ in ()).throw(restore.RestoreError("rollback readback unavailable")), ) result = restore.run_cleanup(cleanup_args) assert result["status"] == "fail" assert result["error"]["phase"] == "rollback_readback" assert result["database"]["clone_database_remaining"] == 0 assert state == {"exists": False, "dropped": True} receipt_path = args.run_root / args.run_id / "cleanup-receipt.json" assert receipt_path.stat().st_mode & 0o777 == 0o600 assert json.loads(receipt_path.read_text())["error"]["phase"] == "rollback_readback" 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_cloudsql_control_plane_preflight_requires_private_only_expected_host( monkeypatch: pytest.MonkeyPatch, ) -> None: payload = { "name": "teleo-pgvector-standby", "state": "RUNNABLE", "databaseVersion": "POSTGRES_16", "region": "europe-west6", "settings": { "tier": "db-custom-1-3840", "ipConfiguration": { "ipv4Enabled": False, "privateNetwork": "projects/teleo-501523/global/networks/teleo-staging-net", }, }, "ipAddresses": [{"type": "PRIVATE", "ipAddress": "10.61.0.3"}], } monkeypatch.setattr( restore, "_run", lambda *_args, **_kwargs: subprocess.CompletedProcess([], 0, json.dumps(payload), ""), ) result = restore.cloudsql_control_plane_state( "teleo-501523", "teleo-pgvector-standby", "10.61.0.3", restore.DEFAULT_PRIVATE_NETWORK, timeout=2.0, ) assert result["public_ip_disabled"] is True assert result["private_addresses"] == ["10.61.0.3"] payload["settings"]["ipConfiguration"]["ipv4Enabled"] = True payload["ipAddresses"].append({"type": "PRIMARY", "ipAddress": "34.1.2.3"}) with pytest.raises(restore.RestoreError, match="no_non_private_address, public_ip_disabled"): restore.cloudsql_control_plane_state( "teleo-501523", "teleo-pgvector-standby", "10.61.0.3", restore.DEFAULT_PRIVATE_NETWORK, timeout=2.0, ) def test_gce_compute_identity_is_bound_to_metadata_instance_zone_and_network( monkeypatch: pytest.MonkeyPatch, ) -> None: metadata = { "project/project-id": "teleo-501523", "project/numeric-project-id": "785938879453", "instance/name": "teleo-prod-1", "instance/id": "123456789", "instance/zone": "projects/123456789/zones/europe-west6-a", "instance/network-interfaces/0/network": "projects/785938879453/networks/teleo-staging-net", "instance/network-interfaces/0/ip": "10.61.0.2", "instance/service-accounts/default/email": "teleo-runtime@teleo-501523.iam.gserviceaccount.com", } monkeypatch.setattr(restore, "gce_metadata_value", lambda path, **_kwargs: metadata[path]) result = restore.gce_compute_identity( "teleo-501523", "teleo-prod-1", "europe-west6-a", restore.DEFAULT_PRIVATE_NETWORK, timeout=2.0, ) assert result["instance_id"] == "123456789" assert result["project_number"] == "785938879453" assert result["zone"] == "europe-west6-a" assert all(result["checks"].values()) with pytest.raises(restore.RestoreError, match="instance"): restore.gce_compute_identity( "teleo-501523", "different-instance", "europe-west6-a", restore.DEFAULT_PRIVATE_NETWORK, timeout=2.0, ) metadata["instance/name"] = "teleo-prod-1" metadata["instance/network-interfaces/0/network"] = "projects/999999999999/networks/teleo-staging-net" with pytest.raises(restore.RestoreError, match="network"): restore.gce_compute_identity( "teleo-501523", "teleo-prod-1", "europe-west6-a", restore.DEFAULT_PRIVATE_NETWORK, timeout=2.0, ) metadata["instance/network-interfaces/0/network"] = "projects/785938879453/networks/teleo-staging-net" with pytest.raises(restore.RestoreError, match="expected_network_project"): restore.gce_compute_identity( "teleo-501523", "teleo-prod-1", "europe-west6-a", "projects/other-project/global/networks/teleo-staging-net", timeout=2.0, ) 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", ] )