feat(db): gated WAL-safe restore of seismo_relay.db + waveforms

sfm/db_snapshot.py

@@ -5,8 +5,11 @@ the waveform store, backing the /db/snapshot, /db/waveforms/recent.zip and
 unit-testable without HTTP.
 """
 from __future__ import annotations
+import os
+import shutil
 import sqlite3
 import zipfile
+import datetime
 from pathlib import Path
 
 
@@ -48,3 +51,56 @@ def build_recent_waveforms_zip(db, store, n: int, dest_path: Path) -> dict:
                 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}")
+        shutil.copy2(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}

tests/test_db_snapshot.py

@@ -79,3 +79,55 @@ def test_recent_waveforms_zip_bundles_all_per_event_files(tmp_path):
         "BE111/BBBB2222.OE0H.a5.pkl",
         "BE111/BBBB2222.OE0H.sfm.json",
     }
+
+
+def test_sfm_db_restore_enabled_reads_env(monkeypatch):
+    from sfm.db_snapshot import sfm_db_restore_enabled
+    monkeypatch.delenv("SFM_DB_RESTORE_ENABLED", raising=False)
+    assert sfm_db_restore_enabled() is False
+    monkeypatch.setenv("SFM_DB_RESTORE_ENABLED", "true")
+    assert sfm_db_restore_enabled() is True
+    monkeypatch.setenv("SFM_DB_RESTORE_ENABLED", "0")
+    assert sfm_db_restore_enabled() is False
+
+
+def test_restore_swaps_db_and_unpacks_waveforms(tmp_path):
+    import zipfile
+    from sfm.db_snapshot import restore_db_and_waveforms
+
+    db_path = tmp_path / "seismo_relay.db"
+    _make_db(db_path, rows=1)
+    (tmp_path / "seismo_relay.db-wal").write_bytes(b"stale")   # must be cleared
+
+    new_db = tmp_path / "incoming.db"
+    _make_db(new_db, rows=9)
+
+    wf_root = tmp_path / "waveforms"
+    wf_root.mkdir()
+    wf_zip = tmp_path / "wf.zip"
+    with zipfile.ZipFile(wf_zip, "w") as zf:
+        zf.writestr("BE111/AAAA.OE0H", b"x")
+        zf.writestr("BE111/AAAA.OE0H.h5", b"x")
+
+    result = restore_db_and_waveforms(db_path, wf_root, new_db, wf_zip)
+
+    conn = sqlite3.connect(str(db_path))
+    assert conn.execute("SELECT COUNT(*) FROM t").fetchone()[0] == 9
+    conn.close()
+    assert not (tmp_path / "seismo_relay.db-wal").exists()
+    assert (wf_root / "BE111" / "AAAA.OE0H.h5").exists()
+    assert result["ok"] and result["waveforms_unpacked"] == 2
+    assert result["safety_backup"] is not None
+
+
+def test_restore_rejects_non_sqlite(tmp_path):
+    import pytest
+    from sfm.db_snapshot import restore_db_and_waveforms
+    db_path = tmp_path / "seismo_relay.db"
+    _make_db(db_path, rows=1)
+    bogus = tmp_path / "bogus.db"
+    bogus.write_bytes(b"not a sqlite file at all")
+    wf_root = tmp_path / "waveforms"
+    wf_root.mkdir()
+    with pytest.raises(ValueError):
+        restore_db_and_waveforms(db_path, wf_root, bogus, None)