versioning/src/versiond/dbsafe.py

"""Enterprise-safe snapshots of live SQLite databases.

A raw byte copy of a live (especially WAL-mode) SQLite file can be torn or
silently stale: the main file, -wal and -shm must be captured in one instant,
which a plain read cannot do. So any file with the SQLite magic is snapshotted
through the SQLite Online Backup API instead of read directly:

1. Open the source read-only (never triggers WAL recovery or checkpointing,
   never takes a write lock, never mutates the live DB).
2. Copy page-by-page with ``Connection.backup()`` (restarts safely if the
   source is being written to; retries while the source is locked).
3. Run ``PRAGMA integrity_check`` on the snapshot; only ``ok`` is stored.
4. Anything else (locked past the deadline, unreadable, corrupt) raises
   ``DatabaseUnsafe`` and nothing is versioned -- a loud skip beats a silent
   corrupt version. Monitor/API surfaces it as ``database-locked`` /
   ``database-corrupt`` instead of storing garbage.

Journal files (``*-wal``, ``*-shm``, ``*-journal``) are never versioned on
their own (see filters: ``database-journal``); they are folded into the main
file's snapshot. For non-SQLite engines (postgres/mysql data files) only a
crash-consistent raw copy is possible -- dump-then-backup remains required.
"""

from __future__ import annotations

import os
import queue
import sqlite3
import tempfile
import threading
from pathlib import Path

SQLITE_MAGIC = b"SQLite format 3\x00"

SNAPSHOT_TIMEOUT_SECONDS = 5.0


class DatabaseUnsafe(Exception):
    """Raised when no consistent snapshot can be produced. Never store partial data."""

    def __init__(self, reason: str, detail: str = ""):
        super().__init__(detail or reason)
        self.reason = reason
        self.detail = detail


def is_sqlite_image(data: bytes) -> bool:
    return data[: len(SQLITE_MAGIC)] == SQLITE_MAGIC


def _integrity_ok(conn: sqlite3.Connection) -> bool:
    row = conn.execute("PRAGMA integrity_check").fetchone()
    return row is not None and row[0] == "ok"


def verify_sqlite_bytes(data: bytes) -> None:
    """Raise DatabaseUnsafe unless data is a fully consistent SQLite image."""
    if not is_sqlite_image(data):
        raise DatabaseUnsafe("not-a-database", "missing SQLite magic")
    fd, tmp = tempfile.mkstemp(prefix="versiond-verify-", suffix=".sqlite")
    try:
        with os.fdopen(fd, "wb") as fh:
            fh.write(data)
        conn = sqlite3.connect(f"file:{tmp}?mode=ro", uri=True)
        try:
            if not _integrity_ok(conn):
                raise DatabaseUnsafe("database-corrupt", "integrity_check failed on pushed content")
        finally:
            conn.close()
    except DatabaseUnsafe:
        raise
    except Exception as exc:
        raise DatabaseUnsafe("database-unreadable", str(exc)) from exc
    finally:
        try:
            os.unlink(tmp)
        except OSError:
            pass


def snapshot_sqlite(path: str, max_bytes: int, timeout: float = SNAPSHOT_TIMEOUT_SECONDS) -> bytes:
    """Consistent, verified snapshot of a live SQLite DB file. Read-only on source.

    The backup runs on a daemon worker thread bounded by ``timeout``: a source
    locked past the deadline fails loudly (``database-locked``) instead of
    blocking the monitor loop forever. SQLite connections are thread-local to
    the worker.
    """
    out: queue.Queue = queue.Queue(maxsize=1)

    def _work() -> None:
        try:
            out.put((True, _snapshot_once(path, max_bytes)))
        except Exception as exc:  # noqa: BLE001 - ferried back to the caller
            out.put((False, exc))

    worker = threading.Thread(target=_work, daemon=True)
    worker.start()
    try:
        ok, payload = out.get(timeout=timeout)
    except queue.Empty as exc:
        raise DatabaseUnsafe("database-locked", f"{path}: still locked after {timeout}s") from exc
    if ok:
        return payload
    if isinstance(payload, DatabaseUnsafe):
        raise payload
    raise DatabaseUnsafe("database-unreadable", f"{path}: {payload}") from payload


def _snapshot_once(path: str, max_bytes: int) -> bytes:
    fd, tmp = tempfile.mkstemp(prefix="versiond-db-", suffix=".sqlite")
    os.close(fd)
    try:
        try:
            src = sqlite3.connect(f"file:{path}?mode=ro", uri=True, timeout=1.0)
        except Exception as exc:
            raise DatabaseUnsafe("database-unreadable", f"{path}: {exc}") from exc
        try:
            dst = sqlite3.connect(tmp)
            try:
                src.backup(dst)
            except sqlite3.OperationalError as exc:
                raise DatabaseUnsafe("database-locked", f"{path}: {exc}") from exc
            finally:
                dst.close()
        finally:
            src.close()
        conn = sqlite3.connect(f"file:{tmp}?mode=ro", uri=True)
        try:
            if not _integrity_ok(conn):
                raise DatabaseUnsafe("database-corrupt", f"integrity_check failed for {path}")
        finally:
            conn.close()
        data = Path(tmp).read_bytes()
        if len(data) > max_bytes:
            raise DatabaseUnsafe("file-too-large", f"{len(data)} bytes")
        return data
    finally:
        try:
            os.unlink(tmp)
        except OSError:
            pass