""" db_snapshot.py — WAL-safe snapshot / restore helpers for seismo_relay.db and the waveform store, backing the /db/snapshot, /db/waveforms/recent.zip and /db/restore endpoints in server.py. Kept out of server.py so the logic is unit-testable without HTTP. """ from __future__ import annotations import os import shutil import sqlite3 import zipfile import datetime from pathlib import Path def snapshot_db(src_path: Path, dest_path: Path) -> Path: """WAL-safe copy of a live SQLite DB via the SQLite online backup API. Consistent even while the source is being written under WAL, and never blocks the source's writers.""" src = sqlite3.connect(str(src_path)) dest = sqlite3.connect(str(dest_path)) try: with dest: src.backup(dest) finally: src.close() dest.close() return Path(dest_path) def build_recent_waveforms_zip(db, store, n: int, dest_path: Path) -> dict: """Zip every on-disk file (event file + .h5/.a5.pkl/.sfm.json siblings) for the n most-recent events, arcname '/'. Events with no resolvable files on disk are skipped. Returns {'events': int, 'files': int}.""" rows = db.query_events(limit=n, offset=0) events = 0 files = 0 with zipfile.ZipFile(dest_path, "w", compression=zipfile.ZIP_DEFLATED) as zf: for row in rows: serial = row.get("serial") fn = row.get("blastware_filename") if not serial or not fn: continue serial_dir = store.root / serial matched = sorted(serial_dir.glob(f"{fn}*")) if serial_dir.is_dir() else [] matched = [p for p in matched if p.is_file()] if not matched: continue events += 1 for p in matched: zf.write(p, arcname=f"{serial}/{p.name}") files += 1 return {"events": events, "files": files} def sfm_db_restore_enabled() -> bool: return os.getenv("SFM_DB_RESTORE_ENABLED", "").strip().lower() in ("1", "true", "yes", "on") def _assert_sqlite(path: Path) -> None: with open(path, "rb") as f: head = f.read(16) if head != b"SQLite format 3\x00": raise ValueError(f"{path} is not a SQLite database") def restore_db_and_waveforms(db_path: Path, waveform_root: Path, new_db_path: Path, waveforms_zip_path) -> dict: """Safety-backup the current seismo_relay.db, atomically swap in new_db_path, clear stale -wal/-shm sidecars, then unpack the waveforms zip under waveform_root. SFM opens a fresh connection per call, so an in-place atomic replace is safe.""" db_path = Path(db_path) new_db_path = Path(new_db_path) _assert_sqlite(new_db_path) safety_name = None if db_path.exists(): ts = datetime.datetime.utcnow().strftime("%Y%m%d_%H%M%S") safety = db_path.with_name(f"{db_path.name}.pre-restore-{ts}") snapshot_db(db_path, safety) safety_name = safety.name # Clear the OUTGOING db's WAL sidecars so they cannot attach to the incoming file. for sfx in ("-wal", "-shm"): side = Path(str(db_path) + sfx) if side.exists(): side.unlink() incoming = db_path.with_name(db_path.name + ".incoming") shutil.copy2(new_db_path, incoming) os.replace(incoming, db_path) # atomic within the same directory unpacked = 0 if waveforms_zip_path and Path(waveforms_zip_path).exists(): waveform_root = Path(waveform_root) waveform_root.mkdir(parents=True, exist_ok=True) with zipfile.ZipFile(waveforms_zip_path) as zf: for m in zf.namelist(): if m.endswith("/") or m.startswith("/") or ".." in Path(m).parts: continue zf.extract(m, waveform_root) unpacked += 1 return {"ok": True, "db_restored": True, "waveforms_unpacked": unpacked, "safety_backup": safety_name}