File size: 5,602 Bytes
aac350d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | """
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()
|