165 lines
7.6 KiB
Python
165 lines
7.6 KiB
Python
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,))
|