teleo-infrastructure/ops/capture_vps_canonical_postgres_snapshot.py

491 lines
19 KiB
Python
Executable file

#!/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 .private_receipt_io import print_private_receipt_summary, write_private_json
from .verify_postgres_parity_manifest import load_manifest
except ImportError: # pragma: no cover - direct script execution
from private_receipt_io import print_private_receipt_summary, write_private_json
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_AUTHORIZATION_REF_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.:/@+-]{7,255}\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")
SNAPSHOT_ID_RE = re.compile(r"[0-9A-F]+-[0-9A-F]+-[0-9]+\Z")
TXID_SNAPSHOT_RE = re.compile(r"[0-9]+:[0-9]+:(?:[0-9]+(?:,[0-9]+)*)?\Z")
WAL_LSN_RE = re.compile(r"[0-9A-F]+/[0-9A-F]+\Z")
SYSTEM_IDENTIFIER_RE = re.compile(r"[0-9]+\Z")
SOURCE_SNAPSHOT_RECEIPT_SCHEMA = "livingip.sourceSnapshotReceipt.v2"
EXPECTED_FILES = {
"teleo-canonical.dump",
"teleo-canonical.dump.sha256",
"teleo-canonical.toc",
"source-manifest.jsonl",
"source-context.txt",
"source-snapshot.json",
"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 sha256_bytes(payload: bytes) -> str:
return hashlib.sha256(payload).hexdigest()
def canonical_sha256(value: Any) -> str:
payload = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
return sha256_bytes(payload.encode())
def load_snapshot_metadata(path: Path) -> dict[str, str]:
try:
payload = json.loads(path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as exc:
raise RuntimeError("source snapshot metadata is unreadable") from exc
if not isinstance(payload, dict):
raise RuntimeError("source snapshot metadata must be an object")
fields = {
"exported_snapshot_id": SNAPSHOT_ID_RE,
"txid_snapshot": TXID_SNAPSHOT_RE,
"wal_lsn": WAL_LSN_RE,
"system_identifier": SYSTEM_IDENTIFIER_RE,
}
normalized: dict[str, str] = {}
for field, pattern in fields.items():
value = str(payload.get(field) or "")
if not pattern.fullmatch(value):
raise RuntimeError(f"source snapshot metadata has invalid {field}")
normalized[field] = value
captured_at_utc = str(payload.get("captured_at_utc") or "")
if not captured_at_utc or any(character in captured_at_utc for character in "\r\n\x00"):
raise RuntimeError("source snapshot metadata has invalid captured_at_utc")
normalized["captured_at_utc"] = captured_at_utc
return normalized
def load_service_state(path: Path) -> dict[str, str]:
try:
state = dict(line.split("=", 1) for line in path.read_text(encoding="utf-8").splitlines() if "=" in line)
main_pid = int(state.get("MainPID") or 0)
restarts = int(state.get("NRestarts") or 0)
except (OSError, TypeError, ValueError) as exc:
raise RuntimeError("source service state is unreadable") from exc
if state.get("ActiveState") != "active" or state.get("SubState") != "running" or main_pid <= 0 or restarts < 0:
raise RuntimeError("source service is not active/running with a positive MainPID")
return state
def build_provenance_binding(
*,
run_id: str,
authorization_ref: str,
source: dict[str, str],
snapshot: dict[str, str],
dump_sha256: str,
manifest_sql_sha256: str,
source_manifest_sha256: str,
source_context_sha256: str,
) -> dict[str, Any]:
payload = {
"receipt_schema": SOURCE_SNAPSHOT_RECEIPT_SCHEMA,
"run_id": run_id,
"capture_authorization_ref": authorization_ref,
"source": source,
"exported_snapshot_id": snapshot["exported_snapshot_id"],
"txid_snapshot": snapshot["txid_snapshot"],
"wal_lsn": snapshot["wal_lsn"],
"source_system_identifier": snapshot["system_identifier"],
"dump_sha256": dump_sha256,
"manifest_sql_sha256": manifest_sql_sha256,
"source_manifest_sha256": source_manifest_sha256,
"source_context_sha256": source_context_sha256,
}
return {"algorithm": "sha256", "payload": payload, "sha256": canonical_sha256(payload)}
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
assert_service_healthy() {{
local state_file="$1"
grep -qx 'ActiveState=active' "$state_file"
grep -qx 'SubState=running' "$state_file"
local main_pid restarts
main_pid="$(sed -n 's/^MainPID=//p' "$state_file")"
restarts="$(sed -n 's/^NRestarts=//p' "$state_file")"
[[ "$main_pid" =~ ^[1-9][0-9]*$ ]]
[[ "$restarts" =~ ^[0-9]+$ ]]
}}
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"
assert_service_healthy "$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
printf '%s\n' "select jsonb_build_object(
'exported_snapshot_id', '$snapshot_id',
'txid_snapshot', txid_current_snapshot()::text,
'wal_lsn', pg_current_wal_lsn()::text,
'system_identifier', (select system_identifier::text from pg_control_system()),
'captured_at_utc', clock_timestamp()
)::text;" >&3
IFS= read -r snapshot_metadata <&4
if [[ -z "$snapshot_metadata" ]]; then
echo "invalid source snapshot metadata" >&2
exit 1
fi
printf '%s\n' "$snapshot_metadata" > "$remote_root/source-snapshot.json"
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"
assert_service_healthy "$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 \
source-snapshot.json \
service-before.txt \
service-after.txt
"""
def upload_manifest(ssh: list[str], target: str, remote_root: str, manifest: bytes) -> 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,
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())
manifest_sql = args.manifest.resolve().read_bytes()
manifest_sql_sha256 = sha256_bytes(manifest_sql)
try:
upload_manifest(ssh, args.ssh_target, remote_root, manifest_sql)
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")
snapshot = load_snapshot_metadata(output_dir / "source-snapshot.json")
before = load_service_state(output_dir / "service-before.txt")
after = load_service_state(output_dir / "service-after.txt")
source = {
"ssh_target": args.ssh_target,
"container": args.container,
"database": args.database,
"service": args.service,
}
source_manifest_sha256 = sha256(output_dir / "source-manifest.jsonl")
source_context_sha256 = sha256(output_dir / "source-context.txt")
provenance_binding = build_provenance_binding(
run_id=args.run_id,
authorization_ref=args.authorization_ref,
source=source,
snapshot=snapshot,
dump_sha256=actual_sha,
manifest_sql_sha256=manifest_sql_sha256,
source_manifest_sha256=source_manifest_sha256,
source_context_sha256=source_context_sha256,
)
return {
"artifact": "vps_canonical_postgres_exported_snapshot",
"schema": SOURCE_SNAPSHOT_RECEIPT_SCHEMA,
"receipt_version": 2,
"generated_at_utc": datetime.now(UTC).isoformat(),
"status": "pass",
"run_id": args.run_id,
"capture_authorization_ref": args.authorization_ref,
"source": source,
"snapshot_exported": True,
"snapshot": snapshot,
"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"),
"sha256": source_manifest_sha256,
"manifest_sql_sha256": manifest_sql_sha256,
"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"],
},
"source_context": {
"path": str(output_dir / "source-context.txt"),
"sha256": source_context_sha256,
},
"source_service": {"before": before, "after": after, "unchanged": before == after},
"provenance_binding": provenance_binding,
"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(argv: list[str] | None = None) -> 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("--authorization-ref", required=True)
parser.add_argument("--output-dir", required=True, type=Path)
parser.add_argument("--timeout", default=180, type=int)
args = parser.parse_args(argv)
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 SAFE_AUTHORIZATION_REF_RE.fullmatch(args.authorization_ref):
parser.error("--authorization-ref must be 8-256 safe metadata 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",
"schema": SOURCE_SNAPSHOT_RECEIPT_SCHEMA,
"status": "dry_run",
"run_id": args.run_id,
"capture_authorization_ref": args.authorization_ref,
"source": f"{args.ssh_target}:{args.container}/{args.database}",
"manifest_sql_sha256": sha256(args.manifest.resolve()),
"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",
"schema": SOURCE_SNAPSHOT_RECEIPT_SCHEMA,
"generated_at_utc": datetime.now(UTC).isoformat(),
"status": "fail",
"run_id": args.run_id,
"capture_authorization_ref": args.authorization_ref,
"error": str(exc),
}
receipt = (
output_dir.parent / f".{output_dir.name}-{args.run_id}-failure.json"
if output_preexisted
else output_dir / "receipt.json"
)
write_private_json(receipt, payload)
print_private_receipt_summary(receipt, payload)
return 1
receipt = args.output_dir.resolve() / "receipt.json"
write_private_json(receipt, payload)
print_private_receipt_summary(receipt, payload)
return 0
if __name__ == "__main__":
raise SystemExit(main())