from contextlib import contextmanager import json import os from pathlib import Path import sqlite3 import uuid def encode(value): return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) class Storage: def __init__(self, path: str): self.path = path Path(path).parent.mkdir(parents=True, exist_ok=True) with self.connect() as db: db.executescript(""" PRAGMA journal_mode=WAL; CREATE TABLE IF NOT EXISTS projects ( id TEXT PRIMARY KEY, revision INTEGER NOT NULL, config TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS jobs ( id TEXT PRIMARY KEY, project_id TEXT NOT NULL, revision INTEGER NOT NULL, snapshot TEXT NOT NULL, request TEXT NOT NULL, status TEXT NOT NULL, model TEXT, error TEXT, created_at TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')) ); CREATE TABLE IF NOT EXISTS items ( job_id TEXT NOT NULL, position INTEGER NOT NULL, id TEXT NOT NULL, source TEXT NOT NULL, translation TEXT, status TEXT NOT NULL DEFAULT 'pending', error TEXT, attempts INTEGER NOT NULL DEFAULT 0, cached INTEGER NOT NULL DEFAULT 0, PRIMARY KEY (job_id, position) ); CREATE TABLE IF NOT EXISTS cache ( fingerprint TEXT NOT NULL, source TEXT NOT NULL, translation TEXT NOT NULL, PRIMARY KEY (fingerprint, source) ); CREATE INDEX IF NOT EXISTS jobs_project ON jobs(project_id, created_at); """) @contextmanager def exclusive_worker(self): """Prevent two processes from recovering/processing the same queue.""" lock_path = str(Path(self.path).resolve()) + ".lock" stream = open(lock_path, "a+b") acquired = False try: if os.name == "nt": import msvcrt stream.seek(0, 2) if stream.tell() == 0: stream.write(b"0") stream.flush() stream.seek(0) msvcrt.locking(stream.fileno(), msvcrt.LK_NBLCK, 1) else: import fcntl fcntl.flock(stream.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) acquired = True yield except OSError as exc: if not acquired: raise RuntimeError("Another translator process owns this database") from exc raise finally: if acquired: if os.name == "nt": stream.seek(0) msvcrt.locking(stream.fileno(), msvcrt.LK_UNLCK, 1) else: fcntl.flock(stream.fileno(), fcntl.LOCK_UN) stream.close() @contextmanager def connect(self): db = sqlite3.connect(self.path, timeout=10) db.row_factory = sqlite3.Row try: with db: yield db finally: db.close() def create_project(self, config): identity = uuid.uuid4().hex with self.connect() as db: db.execute("INSERT INTO projects VALUES (?, 1, ?)", (identity, encode(config))) return self.get_project(identity) def get_project(self, identity): with self.connect() as db: row = db.execute("SELECT * FROM projects WHERE id=?", (identity,)).fetchone() return {"id": row["id"], "revision": row["revision"], **json.loads(row["config"])} if row else None def projects(self): with self.connect() as db: return [{"id": row["id"], "revision": row["revision"], **json.loads(row["config"])} for row in db.execute("SELECT * FROM projects ORDER BY rowid")] def update_project(self, identity, config): with self.connect() as db: db.execute("UPDATE projects SET config=?, revision=revision+1 WHERE id=?", (encode(config), identity)) return self.get_project(identity) def create_job(self, project, provider, request): identity = uuid.uuid4().hex snapshot = {"project": project, "provider": provider} with self.connect() as db: db.execute("INSERT INTO jobs (id,project_id,revision,snapshot,request,status) VALUES (?,?,?,?,?,'queued')", (identity, project["id"], project["revision"], encode(snapshot), encode(request))) db.executemany("INSERT INTO items(job_id,position,id,source) VALUES (?,?,?,?)", [(identity, i, item["id"], item["text"]) for i, item in enumerate(request["items"])]) return self.get_job(identity) def get_job(self, identity, internal=False): with self.connect() as db: row = db.execute("SELECT * FROM jobs WHERE id=?", (identity,)).fetchone() if not row: return None result = dict(row) result["items"] = [dict(item) for item in db.execute("SELECT id,source,translation,status,error,attempts,cached,position FROM items WHERE job_id=? ORDER BY position", (identity,))] if internal: result["snapshot"] = json.loads(result["snapshot"]) result["request"] = json.loads(result["request"]) else: del result["snapshot"], result["request"] result["progress"] = {status: sum(item["status"] == status for item in result["items"]) for status in ("pending", "completed", "failed")} return result def list_jobs(self, project_id, limit=50, offset=0): with self.connect() as db: return [dict(row) for row in db.execute("SELECT id,project_id,revision,status,model,error,created_at FROM jobs WHERE project_id=? ORDER BY rowid DESC LIMIT ? OFFSET ?", (project_id, limit, offset))] def recover(self): with self.connect() as db: db.execute("UPDATE jobs SET status='queued' WHERE status='running'") def next_job(self): with self.connect() as db: row = db.execute("SELECT id FROM jobs WHERE status='queued' ORDER BY rowid LIMIT 1").fetchone() return row["id"] if row else None def job_status(self, identity, status, error=None, model=None): with self.connect() as db: db.execute("UPDATE jobs SET status=?,error=?,model=COALESCE(?,model) WHERE id=?", (status, error, model, identity)) def item_result(self, job_id, position, status, translation=None, error=None, cached=False): with self.connect() as db: db.execute("UPDATE items SET status=?,translation=?,error=?,cached=? WHERE job_id=? AND position=?", (status, translation, error, int(cached), job_id, position)) def attempted(self, job_id, positions): with self.connect() as db: db.executemany("UPDATE items SET attempts=attempts+1 WHERE job_id=? AND position=?", [(job_id, position) for position in positions]) def cache_get(self, fingerprint, source): with self.connect() as db: row = db.execute("SELECT translation FROM cache WHERE fingerprint=? AND source=?", (fingerprint, source)).fetchone() return row["translation"] if row else None def cache_put(self, fingerprint, source, translation): with self.connect() as db: db.execute("INSERT OR REPLACE INTO cache VALUES (?,?,?)", (fingerprint, source, translation)) def retry(self, identity): with self.connect() as db: db.execute("UPDATE items SET status='pending',error=NULL,attempts=0 WHERE job_id=? AND status!='completed'", (identity,)) db.execute("UPDATE jobs SET status='queued',error=NULL WHERE id=?", (identity,))