"""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