#!/usr/bin/env python3 """Capture one pg_dump and parity manifest from the same exported VPS snapshot.""" from __future__ import annotations import argparse import hashlib import io import json import os import re import shlex import shutil import subprocess import tarfile from datetime import UTC, datetime from pathlib import Path from typing import Any try: from .verify_postgres_parity_manifest import load_manifest except ImportError: # pragma: no cover - direct script execution from verify_postgres_parity_manifest import load_manifest SAFE_NAME_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}\Z") SAFE_RUN_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_-]{7,63}\Z") SAFE_SSH_TARGET_RE = re.compile( r"(?:[A-Za-z0-9][A-Za-z0-9._-]{0,63}@)?[A-Za-z0-9][A-Za-z0-9.-]{0,252}\Z" ) EXPECTED_FILES = { "teleo-canonical.dump", "teleo-canonical.dump.sha256", "teleo-canonical.toc", "source-manifest.jsonl", "source-context.txt", "service-before.txt", "service-after.txt", } def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def ssh_base(identity_file: Path) -> list[str]: return [ "ssh", "-T", "-i", str(identity_file), "-o", "BatchMode=yes", "-o", "ConnectTimeout=12", ] def remote_capture_script(*, run_id: str, container: str, database: str, service: str) -> str: remote_root = f"/tmp/teleo-canonical-snapshot-{run_id}" container_manifest = f"/tmp/teleo-postgres-parity-{run_id}.sql" container_dump = f"/tmp/teleo-canonical-{run_id}.dump" quoted = { "remote_root": shlex.quote(remote_root), "container": shlex.quote(container), "database": shlex.quote(database), "service": shlex.quote(service), "container_manifest": shlex.quote(container_manifest), "container_dump": shlex.quote(container_dump), } return f"""set -euo pipefail remote_root={quoted['remote_root']} container={quoted['container']} database={quoted['database']} service={quoted['service']} container_manifest={quoted['container_manifest']} container_dump={quoted['container_dump']} cleanup() {{ docker exec "$container" rm -f "$container_manifest" "$container_dump" >/dev/null 2>&1 || true rm -rf "$remote_root" }} trap cleanup EXIT test -d "$remote_root" test -f "$remote_root/postgres-parity-manifest.sql" docker inspect "$container" >/dev/null docker exec "$container" pg_isready -U postgres -d "$database" >/dev/null systemctl show "$service" -p ActiveState -p SubState -p MainPID -p NRestarts -p ActiveEnterTimestamp --no-pager > "$remote_root/service-before.txt" docker inspect "$container" --format 'container_id={{{{.Id}}}}\nrunning={{{{.State.Running}}}}\nnetwork_mode={{{{.HostConfig.NetworkMode}}}}' > "$remote_root/source-context.txt" printf 'database=%s\ncaptured_at_utc=%s\n' "$database" "$(date -u +%Y-%m-%dT%H:%M:%SZ)" >> "$remote_root/source-context.txt" docker cp "$remote_root/postgres-parity-manifest.sql" "$container:$container_manifest" >/dev/null coproc EXPORTER {{ docker exec -i "$container" psql -X -U postgres -d "$database" -Atq -v ON_ERROR_STOP=1 }} exporter_pid="$EXPORTER_PID" exec 3>&"${{EXPORTER[1]}}" exec 4<&"${{EXPORTER[0]}}" printf '%s\n' 'begin transaction isolation level repeatable read read only;' 'select pg_export_snapshot();' >&3 IFS= read -r snapshot_id <&4 if [[ ! "$snapshot_id" =~ ^[0-9A-F]+-[0-9A-F]+-[0-9]+$ ]]; then echo "invalid exported snapshot id" >&2 exit 1 fi docker exec "$container" pg_dump \ -U postgres -d "$database" \ --format=custom --compress=9 --no-owner --no-acl \ --snapshot="$snapshot_id" \ --file="$container_dump" docker exec "$container" psql \ -X -U postgres -d "$database" -Atq -v ON_ERROR_STOP=1 \ -v "snapshot_id=$snapshot_id" \ -f "$container_manifest" \ > "$remote_root/source-manifest.jsonl" printf '%s\n' 'commit;' '\\q' >&3 exec 3>&- while IFS= read -r _ <&4; do :; done exec 4<&- wait "$exporter_pid" docker exec "$container" pg_restore --list "$container_dump" > "$remote_root/teleo-canonical.toc" docker cp "$container:$container_dump" "$remote_root/teleo-canonical.dump" >/dev/null sha256sum "$remote_root/teleo-canonical.dump" > "$remote_root/teleo-canonical.dump.sha256" printf 'snapshot_exported=true\n' >> "$remote_root/source-context.txt" systemctl show "$service" -p ActiveState -p SubState -p MainPID -p NRestarts -p ActiveEnterTimestamp --no-pager > "$remote_root/service-after.txt" cmp -s "$remote_root/service-before.txt" "$remote_root/service-after.txt" tar -C "$remote_root" -cf - \ teleo-canonical.dump \ teleo-canonical.dump.sha256 \ teleo-canonical.toc \ source-manifest.jsonl \ source-context.txt \ service-before.txt \ service-after.txt """ def upload_manifest(ssh: list[str], target: str, remote_root: str, manifest: Path) -> None: command = ( f"umask 077; rm -rf {shlex.quote(remote_root)}; " f"mkdir -m 700 {shlex.quote(remote_root)}; " f"tee {shlex.quote(remote_root + '/postgres-parity-manifest.sql')} >/dev/null" ) completed = subprocess.run( [*ssh, target, "--", command], input=manifest.read_bytes(), capture_output=True, check=False, ) if completed.returncode != 0: raise RuntimeError("manifest upload failed: " + completed.stderr.decode(errors="replace")[-1000:]) def cleanup_remote(ssh: list[str], target: str, remote_root: str) -> None: subprocess.run( [*ssh, target, "--", f"rm -rf {shlex.quote(remote_root)}"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=False, ) def extract_capture(payload: bytes, output_dir: Path) -> None: seen: set[str] = set() with tarfile.open(fileobj=io.BytesIO(payload), mode="r:") as archive: for member in archive.getmembers(): if not member.isfile() or member.name not in EXPECTED_FILES: raise RuntimeError(f"unexpected remote capture member: {member.name}") if member.name in seen: raise RuntimeError(f"duplicate remote capture member: {member.name}") source = archive.extractfile(member) if source is None: raise RuntimeError(f"missing remote capture payload: {member.name}") destination = output_dir / member.name descriptor = os.open(destination, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) with os.fdopen(descriptor, "wb") as handle: handle.write(source.read()) seen.add(member.name) if seen != EXPECTED_FILES: raise RuntimeError(f"incomplete remote capture: missing={sorted(EXPECTED_FILES - seen)}") def toc_counts(path: Path) -> dict[str, int]: lines = [line for line in path.read_text(errors="replace").splitlines() if line and not line.startswith(";")] return { "entries": len(lines), "table_data": sum(" TABLE DATA " in line for line in lines), "constraints": sum(" CONSTRAINT " in line or " FK CONSTRAINT " in line for line in lines), "indexes": sum(" INDEX " in line for line in lines), } def capture(args: argparse.Namespace) -> dict[str, Any]: output_dir = args.output_dir.resolve() if output_dir.exists(): raise ValueError("--output-dir must not already exist") output_dir.parent.mkdir(parents=True, exist_ok=True) output_dir.mkdir(mode=0o700) remote_root = f"/tmp/teleo-canonical-snapshot-{args.run_id}" ssh = ssh_base(args.ssh_key.resolve()) try: upload_manifest(ssh, args.ssh_target, remote_root, args.manifest.resolve()) script = remote_capture_script( run_id=args.run_id, container=args.container, database=args.database, service=args.service, ) completed = subprocess.run( [*ssh, args.ssh_target, "--", "bash", "-s"], input=script.encode(), capture_output=True, check=False, timeout=args.timeout, ) if completed.returncode != 0: raise RuntimeError("remote snapshot capture failed: " + completed.stderr.decode(errors="replace")[-2000:]) extract_capture(completed.stdout, output_dir) finally: cleanup_remote(ssh, args.ssh_target, remote_root) dump_path = output_dir / "teleo-canonical.dump" expected_sha = (output_dir / "teleo-canonical.dump.sha256").read_text().split()[0] actual_sha = sha256(dump_path) if expected_sha != actual_sha: raise RuntimeError("captured dump SHA-256 mismatch") manifest = load_manifest(output_dir / "source-manifest.jsonl") before = (output_dir / "service-before.txt").read_text() after = (output_dir / "service-after.txt").read_text() return { "artifact": "vps_canonical_postgres_exported_snapshot", "generated_at_utc": datetime.now(UTC).isoformat(), "status": "pass", "source": { "ssh_target": args.ssh_target, "container": args.container, "database": args.database, "service": args.service, }, "snapshot_exported": True, "dump": { "path": str(dump_path), "bytes": dump_path.stat().st_size, "sha256": actual_sha, "toc": toc_counts(output_dir / "teleo-canonical.toc"), }, "manifest": { "path": str(output_dir / "source-manifest.jsonl"), "table_count": len(manifest["tables"]), "total_rows": sum(int(row["row_count"]) for row in manifest["tables"].values()), "application_roles": manifest["singleton"]["application_roles"]["items"], }, "service_unchanged": before == after, "production_db_mutated": False, "telegram_send_attempted": False, } def cleanup_failed_output(output_dir: Path, *, preexisting: bool) -> None: if not preexisting and output_dir.exists(): shutil.rmtree(output_dir) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--execute", action="store_true") parser.add_argument("--ssh-target", required=True) parser.add_argument("--ssh-key", required=True, type=Path) parser.add_argument("--container", default="teleo-pg") parser.add_argument("--database", default="teleo") parser.add_argument("--service", default="leoclean-gateway.service") parser.add_argument("--manifest", default=Path("ops/postgres_parity_manifest.sql"), type=Path) parser.add_argument("--run-id", required=True) parser.add_argument("--output-dir", required=True, type=Path) parser.add_argument("--timeout", default=180, type=int) args = parser.parse_args() for value, flag in ((args.container, "--container"), (args.database, "--database"), (args.service, "--service")): if not SAFE_NAME_RE.fullmatch(value): parser.error(f"{flag} contains unsafe characters") if not SAFE_SSH_TARGET_RE.fullmatch(args.ssh_target): parser.error("--ssh-target must be a host or user@host without options") if not SAFE_RUN_RE.fullmatch(args.run_id): parser.error("--run-id must be 8-64 safe characters") if not args.ssh_key.is_file(): parser.error("--ssh-key must be an existing file") if not args.manifest.is_file(): parser.error("--manifest must be an existing file") if not args.execute: print( json.dumps( { "artifact": "vps_canonical_postgres_exported_snapshot_plan", "status": "dry_run", "source": f"{args.ssh_target}:{args.container}/{args.database}", "output_dir": str(args.output_dir), "mutates_production_db": False, }, indent=2, sort_keys=True, ) ) raise SystemExit(0) return args def main() -> int: args = parse_args() output_dir = args.output_dir.resolve() output_preexisted = output_dir.exists() try: payload = capture(args) except (OSError, ValueError, RuntimeError, subprocess.TimeoutExpired) as exc: cleanup_failed_output(output_dir, preexisting=output_preexisted) payload = { "artifact": "vps_canonical_postgres_exported_snapshot", "generated_at_utc": datetime.now(UTC).isoformat(), "status": "fail", "error": str(exc), } print(json.dumps(payload, indent=2, sort_keys=True)) return 1 receipt = args.output_dir.resolve() / "receipt.json" receipt.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") os.chmod(receipt, 0o600) print(json.dumps(payload, indent=2, sort_keys=True)) return 0 if __name__ == "__main__": raise SystemExit(main())