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