374 lines
12 KiB
Python
374 lines
12 KiB
Python
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",
|
|
]
|
|
)
|