141 lines
5.4 KiB
Python
141 lines
5.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Run one empty-environment restore gate against a completed backup.
|
|
|
|
Implements ``runbooks/RESTORE.md``: restores a verified backup into a
|
|
disposable database and disposable Qdrant collections, reconciles point
|
|
counts, re-verifies known chunks/citations, and proves cross-space denial
|
|
still holds -- all without ever writing to the live source. See
|
|
``src/ragcore/application/restore/service.py`` for the actual pipeline;
|
|
this script only wires real settings and infrastructure to it, and reports
|
|
each stage's outcome.
|
|
|
|
Golden-query probes and the security-denial probe are supplied via a small
|
|
JSON file (``--probes``) since they depend on what was actually indexed in
|
|
the backed-up environment -- this script has no way to infer them on its
|
|
own. See ``runbooks/RESTORE.md`` for the expected shape.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
from uuid import UUID
|
|
|
|
from qdrant_client import AsyncQdrantClient
|
|
|
|
from ragcore.application.restore import GoldenQueryProbe, RestoreGateService, SecurityDenialProbe
|
|
from ragcore.config import get_settings
|
|
from ragcore.infrastructure.backup import LocalFileCopier
|
|
from ragcore.infrastructure.restore import (
|
|
QdrantHttpRestorer,
|
|
SqlAlchemyControlPlaneReader,
|
|
SqlAlchemyPostgresRestorer,
|
|
)
|
|
|
|
|
|
def _load_probes(path: Path) -> tuple[tuple[GoldenQueryProbe, ...], SecurityDenialProbe]:
|
|
body = json.loads(path.read_text())
|
|
golden_queries = tuple(
|
|
GoldenQueryProbe(
|
|
collection_alias=item["collection_alias"],
|
|
point_id=UUID(item["point_id"]),
|
|
query_vector=tuple(item["query_vector"]),
|
|
expected_excerpt_checksum=item["expected_excerpt_checksum"],
|
|
)
|
|
for item in body.get("golden_queries", ())
|
|
)
|
|
security = body["security_denial"]
|
|
security_probe = SecurityDenialProbe(
|
|
principal_id=UUID(security["principal_id"]),
|
|
workspace_id=UUID(security["workspace_id"]),
|
|
granted_space_id=UUID(security["granted_space_id"]),
|
|
denied_space_id=UUID(security["denied_space_id"]),
|
|
)
|
|
return golden_queries, security_probe
|
|
|
|
|
|
async def run(*, backup_dir: Path, probes_path: Path, disposable_suffix: str) -> None:
|
|
settings = get_settings()
|
|
golden_queries, security_probe = _load_probes(probes_path)
|
|
|
|
# Snapshot recovery can legitimately take minutes for large collections
|
|
# or slow disks; a short default timeout would abort an otherwise-healthy
|
|
# restore.
|
|
qdrant_client = AsyncQdrantClient(
|
|
url=settings.qdrant_url, api_key=settings.qdrant_api_key.get_secret_value(), timeout=300
|
|
)
|
|
postgres_restorer = SqlAlchemyPostgresRestorer(admin_database_url=settings.database_url)
|
|
disposable_database_name = f"restore_gate_{disposable_suffix}"
|
|
disposable_collection_prefix = f"restored-{disposable_suffix}"
|
|
try:
|
|
with tempfile.TemporaryDirectory(prefix="ragcore-restore-gate-") as disposable_root:
|
|
service = RestoreGateService(
|
|
postgres=postgres_restorer,
|
|
qdrant=QdrantHttpRestorer(client=qdrant_client),
|
|
control_plane=SqlAlchemyControlPlaneReader(),
|
|
files=LocalFileCopier(),
|
|
)
|
|
report = await service.run(
|
|
backup_dir=backup_dir,
|
|
source_database_url=settings.database_url,
|
|
disposable_database_name=disposable_database_name,
|
|
disposable_collection_prefix=disposable_collection_prefix,
|
|
disposable_blob_root=Path(disposable_root) / "blobs",
|
|
disposable_parsed_root=Path(disposable_root) / "parsed",
|
|
golden_queries=golden_queries,
|
|
security_probe=security_probe,
|
|
)
|
|
finally:
|
|
await postgres_restorer.drop_database(database_name=disposable_database_name)
|
|
await qdrant_client.close()
|
|
|
|
print(f"restore gate for backup {report.backup_id}: {'PASS' if report.passed else 'FAIL'}")
|
|
print(f" schema revision: {report.schema_revision}")
|
|
print(f" postgres row counts: {report.postgres_row_counts}")
|
|
for result in report.reconciliation:
|
|
status = "match" if result.matches else "MISMATCH"
|
|
print(
|
|
f" reconciliation[{result.collection_alias}]: {status} "
|
|
f"(missing={len(result.missing_point_ids)}, orphan={len(result.orphan_point_ids)})"
|
|
)
|
|
for golden in report.golden_queries:
|
|
print(f" golden query[{golden.collection_alias}]: {'pass' if golden.passed else 'FAIL'}")
|
|
if report.security_denial is not None:
|
|
print(f" security denial: {'pass' if report.security_denial.passed else 'FAIL'}")
|
|
|
|
if not report.passed:
|
|
sys.exit(1)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--backup-dir", type=Path, required=True, help="Path to one backup set")
|
|
parser.add_argument(
|
|
"--probes",
|
|
type=Path,
|
|
required=True,
|
|
help="JSON file with golden_queries and security_denial probes",
|
|
)
|
|
parser.add_argument(
|
|
"--disposable-suffix",
|
|
default="cli",
|
|
help="Suffix for the disposable database/collection names this run creates",
|
|
)
|
|
args = parser.parse_args()
|
|
asyncio.run(
|
|
run(
|
|
backup_dir=args.backup_dir,
|
|
probes_path=args.probes,
|
|
disposable_suffix=args.disposable_suffix,
|
|
)
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|