api/store.py

"""SQLite persistence for projects, records, jobs, logs, and artifacts."""

from __future__ import annotations

import json
import sqlite3
from datetime import datetime, timezone
from pathlib import Path
from threading import Lock
from typing import Any
from uuid import uuid4

DB_PATH = Path(__file__).resolve().parent.parent / "data" / "platform.db"
_LOCK = Lock()


def utcnow() -> str:
    return datetime.now(timezone.utc).isoformat()


def init_db(path: Path | None = None) -> None:
    path = path or DB_PATH
    path.parent.mkdir(parents=True, exist_ok=True)
    with sqlite3.connect(path) as db:
        db.executescript(
            """
            CREATE TABLE IF NOT EXISTS projects (
                project_id TEXT PRIMARY KEY, name TEXT NOT NULL, cancer_type TEXT NOT NULL,
                metadata TEXT NOT NULL, created_at TEXT NOT NULL
            );
            CREATE TABLE IF NOT EXISTS records (
                record_id TEXT PRIMARY KEY, project_id TEXT NOT NULL, kind TEXT NOT NULL,
                payload TEXT NOT NULL, created_at TEXT NOT NULL,
                FOREIGN KEY(project_id) REFERENCES projects(project_id)
            );
            CREATE INDEX IF NOT EXISTS idx_records_project_kind
                ON records(project_id, kind);
            CREATE TABLE IF NOT EXISTS jobs (
                job_id TEXT PRIMARY KEY, project_id TEXT NOT NULL, operation TEXT NOT NULL,
                status TEXT NOT NULL, input TEXT NOT NULL, output TEXT,
                error TEXT, logs TEXT NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL
            );
            CREATE TABLE IF NOT EXISTS artifacts (
                artifact_id TEXT PRIMARY KEY, project_id TEXT NOT NULL, job_id TEXT,
                kind TEXT NOT NULL, name TEXT NOT NULL, media_type TEXT NOT NULL,
                payload TEXT NOT NULL, created_at TEXT NOT NULL
            );
            """
        )


def _execute(sql: str, args: tuple = (), *, fetchone=False, fetchall=False):
    with _LOCK, sqlite3.connect(DB_PATH) as db:
        db.row_factory = sqlite3.Row
        cur = db.execute(sql, args)
        db.commit()
        if fetchone:
            row = cur.fetchone()
            return dict(row) if row else None
        if fetchall:
            return [dict(row) for row in cur.fetchall()]
        return cur.lastrowid


def create_project(name: str, cancer_type: str, metadata: dict) -> dict:
    project_id = f"PRJ-{uuid4().hex[:12]}"
    now = utcnow()
    _execute(
        "INSERT INTO projects VALUES (?,?,?,?,?)",
        (project_id, name, cancer_type, json.dumps(metadata), now),
    )
    return get_project(project_id)


def get_project(project_id: str) -> dict | None:
    row = _execute("SELECT * FROM projects WHERE project_id=?", (project_id,), fetchone=True)
    if row:
        row["metadata"] = json.loads(row["metadata"])
    return row


def add_records(project_id: str, kind: str, payloads: list[dict]) -> list[str]:
    ids = []
    id_fields = {
        "variant": "variant_id", "structure": "structure_id", "pocket": "pocket_id",
        "compound": "compound_id", "docking": "docking_id", "md": "md_id",
        "ranking": "ranking_id",
    }
    with _LOCK, sqlite3.connect(DB_PATH) as db:
        for payload in payloads:
            rid = payload.get(id_fields.get(kind, "record_id"))
            rid = rid or f"{kind.upper()}-{uuid4().hex[:12]}"
            db.execute(
                "INSERT OR REPLACE INTO records VALUES (?,?,?,?,?)",
                (rid, project_id, kind, json.dumps(payload, allow_nan=False), utcnow()),
            )
            ids.append(rid)
        db.commit()
    return ids


def get_record(project_id: str, kind: str, record_id: str) -> dict | None:
    row = _execute(
        "SELECT payload FROM records WHERE project_id=? AND kind=? AND record_id=?",
        (project_id, kind, record_id),
        fetchone=True,
    )
    return json.loads(row["payload"]) if row else None


def list_records(project_id: str, kind: str) -> list[dict]:
    rows = _execute(
        "SELECT payload FROM records WHERE project_id=? AND kind=? ORDER BY created_at",
        (project_id, kind),
        fetchall=True,
    )
    return [json.loads(row["payload"]) for row in rows]


def create_job(project_id: str, operation: str, input_data: dict) -> dict:
    job_id = f"JOB-{uuid4().hex[:12]}"
    now = utcnow()
    _execute(
        "INSERT INTO jobs VALUES (?,?,?,?,?,?,?,?,?,?)",
        (job_id, project_id, operation, "queued", json.dumps(input_data), None, None, "[]", now, now),
    )
    return get_job(job_id)


def update_job(job_id: str, status: str, *, output: Any = None, error: str | None = None, log: str | None = None):
    job = get_job(job_id)
    logs = job["logs"] if job else []
    if log:
        logs.append({"at": utcnow(), "message": log})
    _execute(
        "UPDATE jobs SET status=?, output=?, error=?, logs=?, updated_at=? WHERE job_id=?",
        (
            status,
            json.dumps(output, allow_nan=False) if output is not None else None,
            error,
            json.dumps(logs),
            utcnow(),
            job_id,
        ),
    )


def get_job(job_id: str) -> dict | None:
    row = _execute("SELECT * FROM jobs WHERE job_id=?", (job_id,), fetchone=True)
    if not row:
        return None
    row["input"] = json.loads(row["input"])
    row["output"] = json.loads(row["output"]) if row["output"] else None
    row["logs"] = json.loads(row["logs"])
    return row


def list_jobs(project_id: str) -> list[dict]:
    rows = _execute(
        "SELECT * FROM jobs WHERE project_id=? ORDER BY created_at DESC",
        (project_id,),
        fetchall=True,
    )
    for row in rows:
        row["input"] = json.loads(row["input"])
        row["output"] = json.loads(row["output"]) if row["output"] else None
        row["logs"] = json.loads(row["logs"])
    return rows


def add_artifact(project_id: str, kind: str, name: str, payload: Any, media_type: str, job_id: str | None = None) -> dict:
    artifact_id = f"ART-{uuid4().hex[:12]}"
    _execute(
        "INSERT INTO artifacts VALUES (?,?,?,?,?,?,?,?)",
        (artifact_id, project_id, job_id, kind, name, media_type, json.dumps(payload, allow_nan=False), utcnow()),
    )
    return get_artifact(artifact_id)


def get_artifact(artifact_id: str) -> dict | None:
    row = _execute("SELECT * FROM artifacts WHERE artifact_id=?", (artifact_id,), fetchone=True)
    if row:
        row["payload"] = json.loads(row["payload"])
    return row


def list_artifacts(project_id: str) -> list[dict]:
    rows = _execute(
        "SELECT artifact_id,project_id,job_id,kind,name,media_type,created_at FROM artifacts WHERE project_id=?",
        (project_id,),
        fetchall=True,
    )
    return rows