338 lines
13 KiB
Python
Executable file
338 lines
13 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 .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())
|