107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
"""
|
|
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 '<serial>/<filename>'. 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}
|