local-translator/local_translator/storage.py

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);
                CREATE TABLE IF NOT EXISTS provider_calls (
                    id INTEGER PRIMARY KEY, job_id TEXT NOT NULL, item_count INTEGER NOT NULL,
                    seconds REAL NOT NULL, prompt_tokens INTEGER, completion_tokens INTEGER,
                    error TEXT
                );
                CREATE INDEX IF NOT EXISTS calls_job ON provider_calls(job_id);
            """)

    @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
        # Avoid system-drive temporary files during large result scans/sorts.
        db.execute("PRAGMA temp_store=MEMORY")
        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,))]
            result["metrics"] = dict(db.execute("""SELECT COUNT(*) AS provider_calls,
                COALESCE(SUM(seconds),0) AS provider_seconds,
                SUM(prompt_tokens) AS prompt_tokens, SUM(completion_tokens) AS completion_tokens,
                COUNT(prompt_tokens) AS calls_with_prompt_usage,
                COUNT(completion_tokens) AS calls_with_completion_usage,
                COALESCE(SUM(error IS NOT NULL),0) AS provider_errors
                FROM provider_calls WHERE job_id=?""", (identity,)).fetchone())
        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:
            # A cancelled job stays cancelled even if the worker finishes a batch afterwards.
            db.execute("UPDATE jobs SET status=?,error=?,model=COALESCE(?,model) WHERE id=? AND status!='cancelled'", (status, error, model, identity))

    def job_state(self, identity):
        with self.connect() as db:
            row = db.execute("SELECT status FROM jobs WHERE id=?", (identity,)).fetchone()
        return row["status"] if row else None

    def cancel(self, identity):
        with self.connect() as db:
            db.execute("UPDATE items SET status='failed',error='Cancelled' WHERE job_id=? AND status='pending'", (identity,))
            db.execute("UPDATE jobs SET status='cancelled',error=NULL WHERE id=? AND status IN ('queued','running')", (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 record_call(self, job_id, item_count, metrics, error=None):
        with self.connect() as db:
            db.execute("INSERT INTO provider_calls(job_id,item_count,seconds,prompt_tokens,completion_tokens,error) VALUES (?,?,?,?,?,?)",
                       (job_id, item_count, metrics.get("seconds", 0), metrics.get("prompt_tokens"), metrics.get("completion_tokens"), error))

    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,))