restore drill: credentials via environment, redacted errors
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1 changed file +15 −5
modified
services/ingest/src/findata_ingest/tasks/backup.py
+15 −5
@@ -75,10 +75,20 @@ def restore_drill(target: str, dest: Path, dataset_id: str = "market_bars.stock. | ||
| 75 | 75 | subprocess.run(["rsync", "-a", f"{host}:{remote_dir}/pgdump/{latest}", f"{dest}/"], check=True, env=_env(), timeout=600) |
| 76 | 76 | base_url = s.database_url.get_secret_value().rsplit("/", 1)[0] |
| 77 | 77 | scratch = f"findata_restore_{datetime.now(UTC).strftime('%H%M%S')}" |
| 78 | − subprocess.run(["bash", "-c", f"psql '{base_url}/postgres' -qc 'CREATE DATABASE {scratch}' && zstd -dc '{dest}/{latest}' | psql -q '{base_url}/{scratch}' >/dev/null"], | |
| 79 | − check=True, env=_env(), timeout=3600) | |
| 80 | − out = subprocess.run(["bash", "-c", f"psql -tA '{base_url}/{scratch}' -c \"SELECT count(*) FROM lake_files WHERE dataset_id='{dataset_id}' AND version_removed IS NULL\""], | |
| 81 | − capture_output=True, text=True, env=_env(), timeout=120).stdout.strip() | |
| 78 | + env = _env() | |
| 79 | + env["FD_PG_ADMIN"] = f"{base_url}/postgres" | |
| 80 | + env["FD_PG_SCRATCH"] = f"{base_url}/{scratch}" | |
| 81 | + env["FD_DUMP"] = f"{dest}/{latest}" | |
| 82 | + | |
| 83 | + def _sh(cmd: str, timeout: int) -> subprocess.CompletedProcess: | |
| 84 | + r = subprocess.run(["bash", "-c", cmd], capture_output=True, text=True, env=env, timeout=timeout) | |
| 85 | + if r.returncode != 0: | |
| 86 | + from findata_core.redact import redact | |
| 87 | + raise RuntimeError(redact(f"restore step failed: {r.stderr.strip()[-400:]}", s.secret_values())) | |
| 88 | + return r | |
| 89 | + | |
| 90 | + _sh('psql "$FD_PG_ADMIN" -qc "CREATE DATABASE \"$(basename "$FD_PG_SCRATCH")\"" && zstd -dc "$FD_DUMP" | psql -q "$FD_PG_SCRATCH" >/dev/null', 3600) | |
| 91 | + out = _sh(f'psql -tA "$FD_PG_SCRATCH" -c "SELECT count(*) FROM lake_files WHERE dataset_id=\'{dataset_id}\' AND version_removed IS NULL"', 120).stdout.strip() | |
| 82 | 92 | manifests = sorted((dest / "manifests" / dataset_id).glob("v*.json"), key=lambda p: int(p.stem[1:])) |
| 83 | 93 | verified = 0 |
| 84 | 94 | missing = 0 |
@@ -91,6 +101,6 @@ def restore_drill(target: str, dest: Path, dataset_id: str = "market_bars.stock. | ||
| 91 | 101 | continue |
| 92 | 102 | h = hashlib.sha256(p.read_bytes()).hexdigest() |
| 93 | 103 | verified += int(h == f["sha256"]) |
| 94 | − subprocess.run(["bash", "-c", f"psql '{base_url}/postgres' -qc 'DROP DATABASE {scratch}'"], env=_env(), timeout=120) | |
| 104 | + subprocess.run(["bash", "-c", 'psql "$FD_PG_ADMIN" -qc "DROP DATABASE \"$(basename "$FD_PG_SCRATCH")\""'], env=env, timeout=120, capture_output=True) | |
| 95 | 105 | return {"pg_dump": latest, "restored_lake_files_in_db": out, "manifest_version": manifests[-1].stem if manifests else None, |
| 96 | 106 | "files_hash_verified": verified, "files_missing": missing, "dest": str(dest)} |
| 97 | 107 | |