Files

193 lines
7.2 KiB
Python

"""Small SQLite-backed metadata store for jobs and events.
This is intentionally dependency-free and acts as the first durable layer for
business metadata. The schema is shaped so it can be moved to PostgreSQL later
without changing callers.
"""
from __future__ import annotations
import json
import sqlite3
import threading
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable, Optional
from .schemas import JobEvent, JobStatus
def _now_iso() -> str:
return datetime.now(timezone.utc).isoformat()
class MetadataStore:
def __init__(self, db_path: Path) -> None:
self.db_path = db_path
self.db_path.parent.mkdir(parents=True, exist_ok=True)
self._lock = threading.Lock()
self._init_db()
def _connect(self) -> sqlite3.Connection:
conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA busy_timeout=5000")
return conn
def _init_db(self) -> None:
with self._lock, self._connect() as conn:
conn.executescript(
"""
CREATE TABLE IF NOT EXISTS jobs (
id TEXT PRIMARY KEY,
status TEXT NOT NULL,
stage TEXT NOT NULL,
progress_percent INTEGER NOT NULL DEFAULT 0,
message TEXT NOT NULL DEFAULT '',
artifacts TEXT NOT NULL DEFAULT '{}',
error TEXT NOT NULL DEFAULT '',
elapsed_seconds REAL,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS job_events (
id INTEGER PRIMARY KEY AUTOINCREMENT,
job_id TEXT NOT NULL,
type TEXT NOT NULL,
stage TEXT NOT NULL,
progress_percent INTEGER NOT NULL DEFAULT 0,
message TEXT NOT NULL DEFAULT '',
timestamp TEXT NOT NULL,
elapsed_seconds REAL
);
CREATE INDEX IF NOT EXISTS idx_job_events_job_timestamp
ON job_events(job_id, id);
"""
)
def upsert_job(self, status: JobStatus) -> None:
with self._lock, self._connect() as conn:
conn.execute(
"""
INSERT INTO jobs (
id, status, stage, progress_percent, message,
artifacts, error, elapsed_seconds, created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
status = excluded.status,
stage = excluded.stage,
progress_percent = excluded.progress_percent,
message = excluded.message,
artifacts = excluded.artifacts,
error = excluded.error,
elapsed_seconds = excluded.elapsed_seconds,
updated_at = excluded.updated_at
""",
(
status.job_id,
status.status,
status.stage,
status.progress_percent,
status.message,
json.dumps(status.artifacts, ensure_ascii=False),
status.error,
status.elapsed_seconds,
status.created_at.isoformat(),
status.updated_at.isoformat(),
),
)
def add_event(self, job_id: str, event: JobEvent) -> None:
with self._lock, self._connect() as conn:
conn.execute(
"""
INSERT INTO job_events (
job_id, type, stage, progress_percent, message, timestamp, elapsed_seconds
) VALUES (?, ?, ?, ?, ?, ?, ?)
""",
(
job_id,
event.type,
event.stage,
event.progress_percent,
event.message,
event.timestamp.isoformat(),
event.elapsed_seconds,
),
)
def load_jobs(self) -> list[JobStatus]:
with self._lock, self._connect() as conn:
rows = conn.execute(
"""
SELECT * FROM jobs
"""
).fetchall()
result: list[JobStatus] = []
for row in rows:
try:
result.append(
JobStatus(
job_id=row["id"],
status=row["status"],
stage=row["stage"],
progress_percent=row["progress_percent"],
message=row["message"],
artifacts=json.loads(row["artifacts"] or "{}"),
error=row["error"],
created_at=datetime.fromisoformat(row["created_at"]),
updated_at=datetime.fromisoformat(row["updated_at"]),
elapsed_seconds=row["elapsed_seconds"],
)
)
except Exception:
continue
return result
def load_events(self, job_id: str, limit: int = 100) -> list[JobEvent]:
with self._lock, self._connect() as conn:
rows = conn.execute(
"""
SELECT type, stage, progress_percent, message, timestamp, elapsed_seconds
FROM job_events
WHERE job_id = ?
ORDER BY id DESC
LIMIT ?
""",
(job_id, limit),
).fetchall()
return [
JobEvent(
type=row["type"],
stage=row["stage"],
progress_percent=row["progress_percent"],
message=row["message"],
timestamp=datetime.fromisoformat(row["timestamp"]),
elapsed_seconds=row["elapsed_seconds"],
)
for row in reversed(rows)
]
def delete_job(self, job_id: str) -> None:
with self._lock, self._connect() as conn:
conn.execute("DELETE FROM jobs WHERE id = ?", (job_id,))
conn.execute("DELETE FROM job_events WHERE job_id = ?", (job_id,))
def job_ids(self) -> set[str]:
with self._lock, self._connect() as conn:
rows = conn.execute("SELECT id FROM jobs").fetchall()
return {row["id"] for row in rows}
def summarize(self) -> dict:
with self._lock, self._connect() as conn:
jobs = conn.execute("SELECT COUNT(*) AS n FROM jobs").fetchone()
events = conn.execute("SELECT COUNT(*) AS n FROM job_events").fetchone()
return {
"db_size_bytes": self.db_path.stat().st_size if self.db_path.exists() else 0,
"jobs": jobs["n"] if jobs else 0,
"events": events["n"] if events else 0,
}