SlabUploader/backend/db/__init__.py

"""S0 db factory package: stdlib sqlite3 ONLY. ... (same as before)"""
from __future__ import annotations

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

DB_PATH_ENV = "SLAB_DB_PATH"
IMAGES_DIR_ENV = "SLAB_IMAGES_DIR"
SCHEMA_FILE = "schema.sql"

_lock = threading.Lock()
_conn: sqlite3.Connection | None = None
_current_path: str | None = None


def _get_schema_path() -> Path:
    # backend/db/__init__.py -> backend/db/schema.sql
    return Path(__file__).parent / SCHEMA_FILE


def get_connection(db_path: str | None = None) -> sqlite3.Connection:
    global _conn, _current_path
    if db_path is None:
        db_path = os.environ.get(DB_PATH_ENV)
        if not db_path:
            raise RuntimeError(f"{DB_PATH_ENV} env var is required")
    db_path = str(Path(db_path).resolve())

    with _lock:
        if _conn is None or _current_path != db_path:
            if _conn is not None:
                _conn.close()
            _conn = sqlite3.connect(db_path, check_same_thread=False)
            _conn.row_factory = sqlite3.Row
            _conn.execute("PRAGMA journal_mode=WAL")
            _conn.execute("PRAGMA foreign_keys=ON")
            _conn.execute("PRAGMA busy_timeout=5000")
            _conn.execute("PRAGMA synchronous=NORMAL")
            _current_path = db_path
        return _conn

def _has_schema_version(conn: sqlite3.Connection) -> bool:
    try:
        return conn.execute("SELECT 1 FROM schema_version LIMIT 1").fetchone() is not None
    except sqlite3.OperationalError:
        return False

# (table, column) pairs that must exist for the DB to be a valid M3 schema.
# A pre-M3 (M2) volume lacks these, so any missing sentinel signals drift.
_SCHEMA_SENTINELS = [
    ("settings", "value"),
    ("slabs", "deleted_at"),
    ("slabs", "server_rev"),
    ("slabs", "metadata_locked"),
    ("slab_photos", "original_blob"),
    ("sync_log", "payload_hash"),
]


def _schema_matches(conn: sqlite3.Connection) -> bool:
    for table, column in _SCHEMA_SENTINELS:
        try:
            cols = {r[1] for r in conn.execute(f"PRAGMA table_info({table})").fetchall()}
        except sqlite3.OperationalError:
            return False
        if column not in cols:
            return False
    return True


# Tables the app owns and that must be rebuilt together on drift.
_M3_TABLES = [
    "slabs", "slab_photos", "pricing_rules", "settings", "woo_categories",
    "woo_attributes", "woo_attribute_terms", "woo_tags", "sync_log",
    "inference_log", "prompts", "prompt_revisions",
]


def _migrate(conn: sqlite3.Connection) -> None:
    """Idempotent column additions for volumes created before a column existed.

    `CREATE TABLE IF NOT EXISTS` in schema.sql is a no-op on an existing table,
    so new columns must be added here. Each guard is a no-op once the column is
    present. #90 adds slab_photos.role (lineup kind) for the Review photo strip.
    """
    photo_cols = {r[1] for r in conn.execute("PRAGMA table_info(slab_photos)").fetchall()}
    if photo_cols and "role" not in photo_cols:
        conn.execute(
            "ALTER TABLE slab_photos ADD COLUMN role TEXT NOT NULL DEFAULT 'batch' "
            "CHECK(role IN ('calibration','overlay','batch'))"
        )


def init_db(db_path: str | None = None) -> sqlite3.Connection:
    conn = get_connection(db_path)
    schema_path = _get_schema_path()
    if not schema_path.exists():
        raise FileNotFoundError(f"Schema not found: {schema_path}")
    with open(schema_path, "r", encoding="utf-8") as f:
        schema = f.read()
    with _lock:
        conn.executescript(schema)
        conn.commit()
    if not _has_schema_version(conn):
        if not _schema_matches(conn):
            logger = logging.getLogger("slab.db")
            logger.info("Detected pre-M3 schema drift, rebuilding app tables")
            for t in _M3_TABLES:
                conn.execute(f"DROP TABLE IF EXISTS {t}")
            conn.executescript(schema)
        conn.execute("INSERT OR REPLACE INTO schema_version (version) VALUES (1)")
        conn.commit()
    # Safe to run every boot: no-op once the columns exist.
    with _lock:
        _migrate(conn)
        conn.commit()

    return conn


def get_images_dir() -> Path:
    img = os.environ.get(IMAGES_DIR_ENV)
    if not img:
        raise RuntimeError(f"{IMAGES_DIR_ENV} env var is required")
    return Path(img).resolve()


def probe_writable(conn: sqlite3.Connection | None = None) -> bool:
    try:
        c = conn or get_connection()
        c.execute("SELECT 1").fetchone()
        return True
    except Exception:
        return False


def probe_images_writable() -> bool:
    try:
        img_dir = get_images_dir()
        if not img_dir.exists():
            img_dir.mkdir(parents=True, exist_ok=True)
        probe = img_dir / ".health-probe"
        probe.write_text("ok", encoding="utf-8")
        probe.unlink(missing_ok=True)
        return True
    except Exception:
        return False


def _reset_for_tests() -> None:
    """Close global conn so temp dirs can be cleaned on Windows (test only)."""
    global _conn, _current_path
    with _lock:
        if _conn is not None:
            try:
                _conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
            except Exception:
                pass
            try:
                _conn.close()
            except Exception:
                pass
            _conn = None
            _current_path = None