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 | subprocess.run(["rsync", "-a", f"{host}:{remote_dir}/pgdump/{latest}", f"{dest}/"], check=True, env=_env(), timeout=600) | 75 | subprocess.run(["rsync", "-a", f"{host}:{remote_dir}/pgdump/{latest}", f"{dest}/"], check=True, env=_env(), timeout=600) | ||
| 76 | base_url = s.database_url.get_secret_value().rsplit("/", 1)[0] | 76 | base_url = s.database_url.get_secret_value().rsplit("/", 1)[0] | ||
| 77 | scratch = f"findata_restore_{datetime.now(UTC).strftime('%H%M%S')}" | 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"], | 78 | + env = _env() | ||
| 79 | − check=True, env=_env(), timeout=3600) | 79 | + env["FD_PG_ADMIN"] = f"{base_url}/postgres" | ||
| 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\""], | 80 | + env["FD_PG_SCRATCH"] = f"{base_url}/{scratch}" | ||
| 81 | − capture_output=True, text=True, env=_env(), timeout=120).stdout.strip() | 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 | manifests = sorted((dest / "manifests" / dataset_id).glob("v*.json"), key=lambda p: int(p.stem[1:])) | 92 | manifests = sorted((dest / "manifests" / dataset_id).glob("v*.json"), key=lambda p: int(p.stem[1:])) | ||
| 83 | verified = 0 | 93 | verified = 0 | ||
| 84 | missing = 0 | 94 | missing = 0 | ||
@@ -91,6 +101,6 @@ def restore_drill(target: str, dest: Path, dataset_id: str = "market_bars.stock. | |||||
| 91 | continue | 101 | continue | ||
| 92 | h = hashlib.sha256(p.read_bytes()).hexdigest() | 102 | h = hashlib.sha256(p.read_bytes()).hexdigest() | ||
| 93 | verified += int(h == f["sha256"]) | 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 | return {"pg_dump": latest, "restored_lake_files_in_db": out, "manifest_version": manifests[-1].stem if manifests else None, | 105 | return {"pg_dump": latest, "restored_lake_files_in_db": out, "manifest_version": manifests[-1].stem if manifests else None, | ||
| 96 | "files_hash_verified": verified, "files_missing": missing, "dest": str(dest)} | 106 | "files_hash_verified": verified, "files_missing": missing, "dest": str(dest)} | ||
| 97 | 107 | ||||