MediaShelf/mediashelf/db.py

"""SQLite access and migrations.

One writer (the scan job), many readers. WAL mode makes that work without
readers blocking. Migrations are plain numbered .sql files applied in order at
startup and tracked in schema_version — no ORM migration framework (§5).
"""

from __future__ import annotations

import logging
import os
import sqlite3
import threading
from pathlib import Path

log = logging.getLogger(__name__)

MIGRATIONS_DIR = Path(__file__).parent / "migrations"

_local = threading.local()


def connect(path: str, *, read_only: bool = False) -> sqlite3.Connection:
    """Open a tuned connection. Callers own the connection lifecycle."""
    if path != ":memory:":
        os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
    conn = sqlite3.connect(path, timeout=30.0, isolation_level=None)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    conn.execute("PRAGMA synchronous=NORMAL")
    conn.execute("PRAGMA busy_timeout=30000")
    conn.execute("PRAGMA temp_store=MEMORY")
    if read_only:
        conn.execute("PRAGMA query_only=ON")
    return conn


def migrate(conn: sqlite3.Connection) -> int:
    """Apply any unapplied migrations. Returns the resulting version."""
    conn.execute(
        "CREATE TABLE IF NOT EXISTS schema_version ("
        " version INTEGER PRIMARY KEY, applied_at INTEGER NOT NULL)"
    )
    row = conn.execute("SELECT COALESCE(MAX(version), 0) AS v FROM schema_version").fetchone()
    current = row["v"]

    files = sorted(MIGRATIONS_DIR.glob("*.sql"))
    applied = current
    for f in files:
        try:
            version = int(f.name.split("_", 1)[0])
        except ValueError:
            log.warning("skipping unnumbered migration %s", f.name)
            continue
        if version <= current:
            continue
        log.info("applying migration %s", f.name)
        sql = f.read_text()
        # Migrations contain triggers with internal semicolons, so the whole file
        # runs as a script rather than being split on ';'. executescript() commits
        # any pending transaction before it runs, so the BEGIN/COMMIT must live
        # inside the script itself for the migration to be atomic.
        script = (
            "BEGIN;\n"
            + sql
            + f"\nINSERT INTO schema_version (version, applied_at)"
            f" VALUES ({version}, strftime('%s','now'));\nCOMMIT;"
        )
        try:
            conn.executescript(script)
        except Exception:
            if conn.in_transaction:
                conn.execute("ROLLBACK")
            raise
        applied = version
    return applied


class Database:
    """Thin wrapper giving each thread its own connection.

    Flask serves requests on multiple threads and the scheduler runs on its own,
    and sqlite3 connections are not shareable across threads.
    """

    def __init__(self, path: str):
        self.path = path

    @property
    def conn(self) -> sqlite3.Connection:
        c = getattr(_local, "conn", None)
        if c is None or getattr(_local, "path", None) != self.path:
            c = connect(self.path)
            _local.conn = c
            _local.path = self.path
        return c

    def close(self) -> None:
        c = getattr(_local, "conn", None)
        if c is not None:
            c.close()
            _local.conn = None

    # ── convenience ──────────────────────────────────────────────────────

    def query(self, sql: str, params: tuple | dict = ()) -> list[sqlite3.Row]:
        return self.conn.execute(sql, params).fetchall()

    def one(self, sql: str, params: tuple | dict = ()) -> sqlite3.Row | None:
        return self.conn.execute(sql, params).fetchone()

    def scalar(self, sql: str, params: tuple | dict = ()):
        row = self.conn.execute(sql, params).fetchone()
        return None if row is None else row[0]

    def execute(self, sql: str, params: tuple | dict = ()) -> sqlite3.Cursor:
        return self.conn.execute(sql, params)

    def executemany(self, sql: str, seq) -> sqlite3.Cursor:
        return self.conn.executemany(sql, seq)

    def migrate(self) -> int:
        return migrate(self.conn)

    # ── settings ─────────────────────────────────────────────────────────

    def get_setting(self, key: str, default=None):
        row = self.one("SELECT value FROM setting WHERE key = ?", (key,))
        return default if row is None else row["value"]

    def set_setting(self, key: str, value: str) -> None:
        self.execute(
            "INSERT INTO setting (key, value) VALUES (?, ?) "
            "ON CONFLICT(key) DO UPDATE SET value = excluded.value",
            (key, str(value)),
        )


class transaction:
    """Context manager for an explicit write transaction.

    isolation_level is None (autocommit), so transactions are explicit here.
    """

    def __init__(self, db: Database):
        self.db = db

    def __enter__(self) -> sqlite3.Connection:
        self.db.conn.execute("BEGIN IMMEDIATE")
        return self.db.conn

    def __exit__(self, exc_type, exc, tb) -> bool:
        if exc_type is None:
            self.db.conn.execute("COMMIT")
        else:
            self.db.conn.execute("ROLLBACK")
        return False