persistence_manager.py

#!/usr/bin/env python3
"""Single owner for verified OmniRoute SQLite recovery and snapshots."""

from __future__ import annotations

import argparse
import contextlib
import hashlib
import json
import logging
import os
import shutil
import sqlite3
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterator

RUNTIME_DB = Path(os.getenv("OMNIROUTE_RUNTIME_DB", "/root/.omniroute/storage.sqlite"))
DATA_DIR = Path(os.getenv("OMNIROUTE_DATA_DIR", "/data/omniroute"))
PERSISTENT_DB = DATA_DIR / "storage.sqlite"
BACKUP_DIR = DATA_DIR / "backups"
INCOMING_DIR = DATA_DIR / ".incoming"
LKG_DB = BACKUP_DIR / "last-known-good.sqlite"
LOCK_FILE = Path(os.getenv("OMNIROUTE_LOCK_FILE", "/tmp/.persistence.lock"))
ROLLING_LIMIT = max(1, int(os.getenv("OMNIROUTE_ROLLING_BACKUPS", "5")))
log = logging.getLogger("persistence")


def configure_logging() -> None:
    logging.basicConfig(
        level=os.getenv("PERSISTENCE_LOG_LEVEL", "INFO"),
        format="%(asctime)s [%(levelname)s] [PERSISTENCE] %(message)s",
        stream=sys.stdout,
    )


def ensure_dirs() -> None:
    for path in (RUNTIME_DB.parent, DATA_DIR, BACKUP_DIR, INCOMING_DIR):
        path.mkdir(parents=True, exist_ok=True)


@contextlib.contextmanager
def persistence_lock(blocking: bool = False) -> Iterator[bool]:
    ensure_dirs()
    import fcntl

    with LOCK_FILE.open("a+", encoding="utf-8") as handle:
        flags = fcntl.LOCK_EX | (0 if blocking else fcntl.LOCK_NB)
        try:
            fcntl.flock(handle.fileno(), flags)
        except BlockingIOError:
            yield False
            return
        try:
            yield True
        finally:
            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)


def sqlite_status(path: Path, deep: bool = False) -> dict:
    result = {
        "path": str(path),
        "exists": path.exists(),
        "size": path.stat().st_size if path.exists() else 0,
        "status": "MISSING",
    }
    if not path.exists():
        return result
    connection = None
    try:
        connection = sqlite3.connect(str(Path(path).resolve()), timeout=5)
        quick = connection.execute("PRAGMA quick_check").fetchone()
        result["quick_check"] = quick[0] if quick else "empty result"
        if result["quick_check"] != "ok":
            result["status"] = "CORRUPT"
            return result
        integrity = connection.execute("PRAGMA integrity_check").fetchone()
        result["integrity_check"] = integrity[0] if integrity else "empty result"
        if result["integrity_check"] != "ok":
            result["status"] = "CORRUPT"
            return result
        result["journal_mode"] = connection.execute("PRAGMA journal_mode").fetchone()[0]
        result["page_count"] = connection.execute("PRAGMA page_count").fetchone()[0]
        result["status"] = "HEALTHY"
    except sqlite3.OperationalError as exc:
        result["status"] = "IO_ERROR" if "disk i/o" in str(exc).lower() else "UNAVAILABLE"
        result["error"] = str(exc)
    except sqlite3.DatabaseError as exc:
        result["status"] = "CORRUPT"
        result["error"] = str(exc)
    except OSError as exc:
        result["status"] = "IO_ERROR"
        result["error"] = str(exc)
    finally:
        if connection is not None:
            connection.close()
    return result


def _sync_file(path: Path) -> None:
    try:
        with path.open("rb") as handle:
            os.fsync(handle.fileno())
    except OSError:
        log.warning("File fsync unavailable for %s", path)


def _sync_dir(path: Path) -> None:
    try:
        fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
        try:
            os.fsync(fd)
        finally:
            os.close(fd)
    except OSError:
        log.warning("Directory fsync unavailable for %s", path)


def sqlite_backup(source: Path, destination: Path) -> dict:
    started = time.monotonic()
    temporary = destination.with_name(f"{destination.name}.tmp-{os.getpid()}-{time.time_ns()}")
    try:
        source_connection = sqlite3.connect(str(Path(source).resolve()), timeout=10)
        destination_connection = sqlite3.connect(str(Path(temporary).resolve()), timeout=10)
        try:
            source_connection.backup(destination_connection, pages=-1)
            destination_connection.commit()
        finally:
            destination_connection.close()
            source_connection.close()
        status = sqlite_status(temporary)
        if status["status"] != "HEALTHY":
            raise RuntimeError(f"snapshot validation failed: {status}")
        _sync_file(temporary)
        os.replace(temporary, destination)
        _sync_dir(destination.parent)
        return {
            "status": "HEALTHY",
            "path": str(destination),
            "size": destination.stat().st_size,
            "sha256": hashlib.sha256(destination.read_bytes()).hexdigest(),
            "duration_seconds": round(time.monotonic() - started, 3),
        }
    finally:
        temporary.unlink(missing_ok=True)


def snapshot(blocking: bool = False) -> dict | None:
    if not RUNTIME_DB.exists():
        log.warning("Runtime database is missing; snapshot skipped")
        return None
    with persistence_lock(blocking=blocking) as acquired:
        if not acquired:
            log.info("Snapshot skipped because another persistence operation holds the lock")
            return None
        stamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
        incoming = INCOMING_DIR / f"storage-{stamp}.sqlite"
        result = sqlite_backup(RUNTIME_DB, incoming)
        rolling = BACKUP_DIR / incoming.name
        os.replace(incoming, rolling)
        _sync_dir(BACKUP_DIR)
        os.replace(rolling, PERSISTENT_DB)
        _sync_dir(DATA_DIR)
        sqlite_backup(PERSISTENT_DB, BACKUP_DIR / incoming.name)
        sqlite_backup(BACKUP_DIR / incoming.name, LKG_DB)
        snapshots = sorted(BACKUP_DIR.glob("storage-*.sqlite"), key=lambda p: p.stat().st_mtime, reverse=True)
        for old in snapshots[ROLLING_LIMIT:]:
            old.unlink(missing_ok=True)
        log.info("Snapshot promoted: %s (%s bytes)", PERSISTENT_DB, result["size"])
        return result


def recover() -> dict:
    ensure_dirs()
    with persistence_lock(blocking=True):
        current = sqlite_status(RUNTIME_DB)
        if current["status"] == "HEALTHY":
            return {"action": "kept", "source": str(RUNTIME_DB), "status": "HEALTHY"}
        candidates = [PERSISTENT_DB, LKG_DB]
        candidates.extend(sorted(BACKUP_DIR.glob("storage-*.sqlite"), key=lambda p: p.stat().st_mtime, reverse=True))
        for candidate in dict.fromkeys(candidates):
            if sqlite_status(candidate, deep=False)["status"] != "HEALTHY":
                continue
            if RUNTIME_DB.exists():
                evidence = BACKUP_DIR / f"runtime-corrupt-{datetime.now(timezone.utc).strftime('%Y%m%d-%H%M%S')}.sqlite"
                shutil.copy2(RUNTIME_DB, evidence)
            recovered = RUNTIME_DB.with_suffix(".recovered")
            sqlite_backup(candidate, recovered)
            for ext in ("-wal", "-shm"):
                wal_f = RUNTIME_DB.parent / f"{RUNTIME_DB.name}{ext}"
                wal_f.unlink(missing_ok=True)
            os.replace(recovered, RUNTIME_DB)
            try:
                r_conn = sqlite3.connect(str(Path(RUNTIME_DB).resolve()), timeout=5)
                r_conn.execute("REINDEX;")
                r_conn.commit()
                r_conn.close()
            except Exception:
                pass
            _sync_dir(RUNTIME_DB.parent)
            log.info("Recovered runtime DB from verified snapshot %s", candidate)
            return {"action": "restored", "source": str(candidate), "status": "HEALTHY"}
        # Try cloud vault restore if all local snapshots failed
        log.warning("No verified local snapshot exists; attempting emergency cloud vault restore...")
        try:
            sys.path.insert(0, "/")
            import vault_sync
            vault_sync.restore_from_vault(force=True)
            if PERSISTENT_DB.exists() and sqlite_status(PERSISTENT_DB)["status"] == "HEALTHY":
                sqlite_backup(PERSISTENT_DB, RUNTIME_DB)
                log.info("Successfully recovered runtime DB from cloud vault snapshot")
                return {"action": "vault_restored", "source": str(PERSISTENT_DB), "status": "HEALTHY"}
        except Exception as e:
            log.error("Emergency cloud vault restore failed: %s", e)

        if current["status"] == "MISSING":
            sqlite3.connect(RUNTIME_DB).close()
            log.warning("No verified snapshot exists; initialized an empty runtime database")
            return {"action": "initialized", "source": None, "status": "NOT_INITIALIZED"}
        raise RuntimeError("No verified SQLite recovery source is available")


def main() -> int:
    configure_logging()
    parser = argparse.ArgumentParser()
    parser.add_argument("command", choices=("recover", "snapshot", "verify", "diagnose"))
    parser.add_argument("--path", default=str(RUNTIME_DB))
    args = parser.parse_args()
    if args.command == "recover":
        output = recover()
    elif args.command == "snapshot":
        output = snapshot() or {"status": "SKIPPED"}
    else:
        output = sqlite_status(Path(args.path), deep=True)
        if args.command == "diagnose":
            output["header_valid"] = Path(args.path).exists() and Path(args.path).read_bytes()[:16] == b"SQLite format 3\x00"
    print(json.dumps(output, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())