storage/database.py

"""
SQLite-backed persistence for jobs + results metadata.

Schema:
  jobs         — one row per job
  job_results  — one row per completed job (JSON blob)
  audit_events — mirrored from utils.audit for queryability

This class is injectable: pass a `path` (or use :memory: for tests).
"""

from __future__ import annotations

import json
import sqlite3
import threading
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

from loguru import logger


_SCHEMA = """
CREATE TABLE IF NOT EXISTS jobs (
    id           TEXT PRIMARY KEY,
    kind         TEXT NOT NULL,
    status       TEXT NOT NULL,
    created_at   TEXT NOT NULL,
    started_at   TEXT,
    completed_at TEXT,
    request      TEXT NOT NULL,
    image_hash   TEXT,
    error        TEXT
);

CREATE TABLE IF NOT EXISTS job_results (
    job_id     TEXT PRIMARY KEY,
    status     TEXT NOT NULL,
    report     TEXT,
    error      TEXT,
    elapsed_ms REAL,
    created_at TEXT NOT NULL,
    FOREIGN KEY (job_id) REFERENCES jobs(id)
);

CREATE INDEX IF NOT EXISTS idx_jobs_status ON jobs(status);
CREATE INDEX IF NOT EXISTS idx_jobs_created ON jobs(created_at);
"""


class Database:
    """Thread-safe SQLite wrapper.  One connection per instance."""

    def __init__(self, path: str = ":memory:") -> None:
        self._path = path
        self._lock = threading.Lock()
        self._conn = sqlite3.connect(path, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._conn.executescript(_SCHEMA)
        self._conn.commit()
        logger.info(f"Database initialized at {path}")

    # ------------------------------------------------------------------ #
    # Jobs
    # ------------------------------------------------------------------ #
    def save_job(self, job: dict) -> None:
        with self._lock:
            self._conn.execute(
                """INSERT OR REPLACE INTO jobs
                   (id, kind, status, created_at, started_at, completed_at,
                    request, image_hash, error)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    job["id"],
                    job["kind"],
                    job["status"],
                    job["created_at"],
                    job.get("started_at"),
                    job.get("completed_at"),
                    json.dumps(job.get("request", {}), default=str),
                    job.get("image_hash"),
                    job.get("error"),
                ),
            )
            self._conn.commit()

    def get_job(self, job_id: str) -> Optional[dict]:
        with self._lock:
            row = self._conn.execute(
                "SELECT * FROM jobs WHERE id = ?", (job_id,)
            ).fetchone()
        if not row:
            return None
        d = dict(row)
        d["request"] = json.loads(d["request"] or "{}")
        return d

    def list_jobs(self, limit: int = 50, status: Optional[str] = None) -> list[dict]:
        with self._lock:
            if status:
                cur = self._conn.execute(
                    "SELECT * FROM jobs WHERE status = ? ORDER BY created_at DESC LIMIT ?",
                    (status, limit),
                )
            else:
                cur = self._conn.execute(
                    "SELECT * FROM jobs ORDER BY created_at DESC LIMIT ?",
                    (limit,),
                )
            rows = [dict(r) for r in cur.fetchall()]
        for r in rows:
            r["request"] = json.loads(r["request"] or "{}")
        return rows

    # ------------------------------------------------------------------ #
    # Results
    # ------------------------------------------------------------------ #
    def save_result(self, job_id: str, status: str, report: Any,
                    error: Optional[str], elapsed_ms: float) -> None:
        with self._lock:
            self._conn.execute(
                """INSERT OR REPLACE INTO job_results
                   (job_id, status, report, error, elapsed_ms, created_at)
                   VALUES (?, ?, ?, ?, ?, ?)""",
                (
                    job_id, status,
                    json.dumps(report, default=str) if report is not None else None,
                    error, elapsed_ms,
                    datetime.now(timezone.utc).isoformat(),
                ),
            )
            self._conn.commit()

    def get_result(self, job_id: str) -> Optional[dict]:
        with self._lock:
            row = self._conn.execute(
                "SELECT * FROM job_results WHERE job_id = ?", (job_id,)
            ).fetchone()
        if not row:
            return None
        d = dict(row)
        if d.get("report"):
            try:
                d["report"] = json.loads(d["report"])
            except json.JSONDecodeError:
                pass
        return d

    # ------------------------------------------------------------------ #
    # Maintenance
    # ------------------------------------------------------------------ #
    def cleanup_old_jobs(self, retention_days: int) -> int:
        """Delete jobs older than retention_days.  Returns count deleted."""
        cutoff = datetime.now(timezone.utc).timestamp() - (retention_days * 86400)
        with self._lock:
            cur = self._conn.execute(
                "DELETE FROM jobs WHERE strftime('%s', created_at) < ?",
                (cutoff,),
            )
            self._conn.commit()
            return cur.rowcount

    def close(self) -> None:
        with self._lock:
            self._conn.close()