modelbeast/server/db.py
2026-07-12 21:05:22 +10:00

66 lines
1.5 KiB
Python

import json
import sqlite3
import time
import uuid
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
DATA = ROOT / "data"
DB_PATH = DATA / "modelbeast.db"
SCHEMA = """
CREATE TABLE IF NOT EXISTS assets (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
kind TEXT NOT NULL,
path TEXT NOT NULL,
size INTEGER NOT NULL DEFAULT 0,
meta TEXT NOT NULL DEFAULT '{}',
parent_job TEXT,
created_at REAL NOT NULL
);
CREATE TABLE IF NOT EXISTS jobs (
id TEXT PRIMARY KEY,
operator TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'queued',
asset_id TEXT,
params TEXT NOT NULL DEFAULT '{}',
outdir TEXT,
log TEXT NOT NULL DEFAULT '',
error TEXT,
created_at REAL NOT NULL,
started_at REAL,
finished_at REAL
);
"""
def connect() -> sqlite3.Connection:
DATA.mkdir(parents=True, exist_ok=True)
# FastAPI sync endpoints run in a threadpool; python sqlite3 is built in
# serialized threading mode, so sharing one connection across threads is safe.
con = sqlite3.connect(DB_PATH, check_same_thread=False)
con.row_factory = sqlite3.Row
con.execute("PRAGMA journal_mode=WAL")
con.executescript(SCHEMA)
return con
def new_id() -> str:
return uuid.uuid4().hex[:12]
def now() -> float:
return time.time()
def row_to_dict(row: sqlite3.Row) -> dict:
d = dict(row)
for key in ("meta", "params"):
if key in d and isinstance(d[key], str):
try:
d[key] = json.loads(d[key])
except ValueError:
pass
return d