diff --git a/job_store.py b/job_store.py new file mode 100644 index 0000000..c1ed154 --- /dev/null +++ b/job_store.py @@ -0,0 +1,264 @@ +"""SQLite-backed durable job store for async/worker job tracking. + +The async web layer currently tracks jobs in an in-memory dict, which +loses all state on restart and cannot be shared across processes. This +module provides :class:`JobStore`, a small stdlib-only persistence layer +(``sqlite3`` with WAL journaling) that any async or worker process can +use so jobs survive restarts and are visible across processes. + +Design notes: + +- Every public method opens a short-lived connection guarded by a lock, + so a single ``JobStore`` instance is safe to share across threads. +- WAL mode allows concurrent readers alongside a writer, which suits a + web process polling job status while a worker updates it. +- Callers pass ``now`` (a :class:`datetime.datetime`) explicitly; the + store never calls ``datetime.now()`` itself, keeping tests + deterministic. + +Example: + >>> from datetime import datetime, timezone + >>> store = JobStore("/tmp/jobs.db") # doctest: +SKIP + >>> store.create("job-1", temp_dir="/tmp/job-1", + ... now=datetime.now(timezone.utc)) # doctest: +SKIP +""" + +from __future__ import annotations + +import sqlite3 +import threading +from datetime import datetime + +#: Allowed job lifecycle states. +VALID_STATUSES = frozenset({"queued", "processing", "done", "failed"}) + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS jobs ( + id TEXT PRIMARY KEY, + status TEXT NOT NULL, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + output_path TEXT, + output_name TEXT, + error TEXT, + temp_dir TEXT +) +""" + +_COLUMNS = ( + "id", + "status", + "created_at", + "updated_at", + "output_path", + "output_name", + "error", + "temp_dir", +) + + +class DuplicateJobError(ValueError): + """Raised by :meth:`JobStore.create` when the job id already exists.""" + + +class JobStore: + """Durable, thread-safe job store backed by a SQLite database file. + + Multiple processes may open independent ``JobStore`` instances on the + same ``db_path``; SQLite's file locking plus WAL mode keeps their + reads and writes consistent. Within one process, a single instance + may be shared freely across threads. + + Args: + db_path: Filesystem path of the SQLite database. Created (along + with the schema) if it does not exist. ``":memory:"`` is not + supported because each operation opens a fresh connection, + which would discard an in-memory database every time. + """ + + def __init__(self, db_path: str) -> None: + """Initialize the store and create the schema if needed. + + Args: + db_path: Path to the SQLite database file. + + Raises: + ValueError: If ``db_path`` is ``":memory:"``. + """ + if db_path == ":memory:": + raise ValueError( + "JobStore requires a file path; ':memory:' databases do " + "not survive the short-lived connections this store uses" + ) + self._db_path = str(db_path) + self._lock = threading.Lock() + with self._connect() as conn: + conn.execute(_SCHEMA) + + def _connect(self) -> sqlite3.Connection: + """Open a new WAL-mode connection to the underlying database. + + Returns: + A ``sqlite3.Connection`` with WAL journaling and a row + factory that yields ``sqlite3.Row`` objects. + """ + conn = sqlite3.connect(self._db_path, timeout=30.0) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA journal_mode=WAL") + return conn + + @staticmethod + def _validate_status(status: str) -> None: + """Reject statuses outside the allowed lifecycle set. + + Args: + status: Candidate status string. + + Raises: + ValueError: If ``status`` is not one of ``VALID_STATUSES``. + """ + if status not in VALID_STATUSES: + allowed = ", ".join(sorted(VALID_STATUSES)) + raise ValueError( + f"invalid status {status!r}; must be one of: {allowed}" + ) + + @staticmethod + def _row_to_dict(row: sqlite3.Row) -> dict: + """Convert a database row into a plain job dict. + + Args: + row: A ``sqlite3.Row`` from the ``jobs`` table. + + Returns: + A dict with the keys ``id``, ``status``, ``created_at``, + ``updated_at``, ``output_path``, ``output_name``, ``error``, + and ``temp_dir``. + """ + return {key: row[key] for key in _COLUMNS} + + def create(self, job_id: str, *, temp_dir: str, now: datetime) -> None: + """Insert a new job in the ``queued`` state. + + Args: + job_id: Unique identifier for the job. + temp_dir: Working directory associated with the job (stored + so a cleanup pass can remove it later). + now: Timestamp recorded as both ``created_at`` and + ``updated_at`` (ISO 8601 via ``datetime.isoformat()``). + + Raises: + DuplicateJobError: If a job with ``job_id`` already exists. + (Subclass of ``ValueError``, so callers may catch either.) + """ + timestamp = now.isoformat() + with self._lock, self._connect() as conn: + try: + conn.execute( + "INSERT INTO jobs (id, status, created_at, updated_at," + " temp_dir) VALUES (?, 'queued', ?, ?, ?)", + (job_id, timestamp, timestamp, temp_dir), + ) + except sqlite3.IntegrityError as exc: + raise DuplicateJobError( + f"job {job_id!r} already exists" + ) from exc + + def set_status( + self, + job_id: str, + status: str, + *, + now: datetime, + output_path: str | None = None, + output_name: str | None = None, + error: str | None = None, + ) -> None: + """Update a job's status, timestamp, and optional result fields. + + Only the fields passed as non-``None`` keyword arguments are + overwritten; previously stored values for ``output_path``, + ``output_name``, and ``error`` are preserved otherwise. + + Args: + job_id: Identifier of the job to update. + status: New status; one of ``queued``, ``processing``, + ``done``, or ``failed``. + now: Timestamp recorded as ``updated_at``. + output_path: Path of the finished output file, if any. + output_name: Client-facing download name, if any. + error: Human-readable failure message, if any. + + Raises: + ValueError: If ``status`` is not allowed. + KeyError: If no job with ``job_id`` exists. + """ + self._validate_status(status) + with self._lock, self._connect() as conn: + cursor = conn.execute( + "UPDATE jobs SET status = ?, updated_at = ?," + " output_path = COALESCE(?, output_path)," + " output_name = COALESCE(?, output_name)," + " error = COALESCE(?, error)" + " WHERE id = ?", + (status, now.isoformat(), output_path, output_name, + error, job_id), + ) + if cursor.rowcount == 0: + raise KeyError(f"job {job_id!r} does not exist") + + def get(self, job_id: str) -> dict | None: + """Fetch a single job by id. + + Args: + job_id: Identifier of the job to look up. + + Returns: + The job as a dict (see :meth:`_row_to_dict` for keys), or + ``None`` if no such job exists. + """ + with self._lock, self._connect() as conn: + row = conn.execute( + "SELECT * FROM jobs WHERE id = ?", (job_id,) + ).fetchone() + return self._row_to_dict(row) if row is not None else None + + def list_jobs(self, status: str | None = None) -> list[dict]: + """List jobs, optionally filtered by status. + + Args: + status: If given, only jobs in this state are returned; must + be one of the allowed statuses. + + Returns: + Jobs as dicts, ordered by ``created_at`` then id for a + stable listing. + + Raises: + ValueError: If ``status`` is given but not allowed. + """ + with self._lock, self._connect() as conn: + if status is None: + rows = conn.execute( + "SELECT * FROM jobs ORDER BY created_at, id" + ).fetchall() + else: + self._validate_status(status) + rows = conn.execute( + "SELECT * FROM jobs WHERE status = ?" + " ORDER BY created_at, id", + (status,), + ).fetchall() + return [self._row_to_dict(row) for row in rows] + + def delete(self, job_id: str) -> None: + """Remove a job record if it exists. + + Deleting an unknown id is a no-op, so cleanup passes can call + this without checking existence first. + + Args: + job_id: Identifier of the job to remove. + """ + with self._lock, self._connect() as conn: + conn.execute("DELETE FROM jobs WHERE id = ?", (job_id,)) diff --git a/tests/test_job_store.py b/tests/test_job_store.py new file mode 100644 index 0000000..7a343e1 --- /dev/null +++ b/tests/test_job_store.py @@ -0,0 +1,203 @@ +"""Tests for the SQLite-backed durable job store (job_store.py).""" + +import os +import tempfile +import threading +import unittest +from datetime import datetime, timedelta, timezone + +from job_store import DuplicateJobError, JobStore + +T0 = datetime(2026, 7, 6, 12, 0, 0, tzinfo=timezone.utc) +T1 = T0 + timedelta(minutes=5) +T2 = T0 + timedelta(minutes=10) + + +class JobStoreTestCase(unittest.TestCase): + """Base fixture: a fresh JobStore on a temp SQLite file.""" + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + self.db_path = os.path.join(self._tmp.name, "jobs.db") + self.store = JobStore(self.db_path) + + +class TestCreateAndGet(JobStoreTestCase): + def test_create_get_roundtrip(self): + self.store.create("job-1", temp_dir="/tmp/job-1", now=T0) + job = self.store.get("job-1") + self.assertEqual( + job, + { + "id": "job-1", + "status": "queued", + "created_at": T0.isoformat(), + "updated_at": T0.isoformat(), + "output_path": None, + "output_name": None, + "error": None, + "temp_dir": "/tmp/job-1", + }, + ) + + def test_get_unknown_returns_none(self): + self.assertIsNone(self.store.get("nope")) + + def test_create_duplicate_raises(self): + self.store.create("job-1", temp_dir="/tmp/a", now=T0) + with self.assertRaises(DuplicateJobError): + self.store.create("job-1", temp_dir="/tmp/b", now=T1) + # DuplicateJobError is a ValueError, so generic handlers work too. + self.assertTrue(issubclass(DuplicateJobError, ValueError)) + # Original record is untouched. + self.assertEqual(self.store.get("job-1")["temp_dir"], "/tmp/a") + + def test_memory_path_rejected(self): + with self.assertRaises(ValueError): + JobStore(":memory:") + + +class TestSetStatus(JobStoreTestCase): + def test_transitions_update_status_and_timestamp(self): + self.store.create("job-1", temp_dir="/tmp/j", now=T0) + + self.store.set_status("job-1", "processing", now=T1) + job = self.store.get("job-1") + self.assertEqual(job["status"], "processing") + self.assertEqual(job["created_at"], T0.isoformat()) + self.assertEqual(job["updated_at"], T1.isoformat()) + + self.store.set_status( + "job-1", "done", now=T2, + output_path="/tmp/out.mp4", output_name="video.mp4", + ) + job = self.store.get("job-1") + self.assertEqual(job["status"], "done") + self.assertEqual(job["updated_at"], T2.isoformat()) + self.assertEqual(job["output_path"], "/tmp/out.mp4") + self.assertEqual(job["output_name"], "video.mp4") + self.assertIsNone(job["error"]) + + def test_failed_records_error(self): + self.store.create("job-1", temp_dir="/tmp/j", now=T0) + self.store.set_status("job-1", "failed", now=T1, error="boom") + job = self.store.get("job-1") + self.assertEqual(job["status"], "failed") + self.assertEqual(job["error"], "boom") + + def test_omitted_fields_are_preserved(self): + self.store.create("job-1", temp_dir="/tmp/j", now=T0) + self.store.set_status( + "job-1", "done", now=T1, + output_path="/tmp/out.mp4", output_name="video.mp4", + ) + # A later update without output fields must not erase them. + self.store.set_status("job-1", "done", now=T2) + job = self.store.get("job-1") + self.assertEqual(job["output_path"], "/tmp/out.mp4") + self.assertEqual(job["output_name"], "video.mp4") + + def test_invalid_status_raises_value_error(self): + self.store.create("job-1", temp_dir="/tmp/j", now=T0) + with self.assertRaises(ValueError): + self.store.set_status("job-1", "exploded", now=T1) + # Job unchanged after the rejected update. + self.assertEqual(self.store.get("job-1")["status"], "queued") + + def test_unknown_job_raises_key_error(self): + with self.assertRaises(KeyError): + self.store.set_status("ghost", "done", now=T0) + + +class TestListAndDelete(JobStoreTestCase): + def test_list_all_and_filter_by_status(self): + self.store.create("a", temp_dir="/tmp/a", now=T0) + self.store.create("b", temp_dir="/tmp/b", now=T1) + self.store.create("c", temp_dir="/tmp/c", now=T2) + self.store.set_status("b", "processing", now=T2) + + all_ids = [job["id"] for job in self.store.list_jobs()] + self.assertEqual(all_ids, ["a", "b", "c"]) + + queued = [job["id"] for job in self.store.list_jobs(status="queued")] + self.assertEqual(queued, ["a", "c"]) + + processing = self.store.list_jobs(status="processing") + self.assertEqual([job["id"] for job in processing], ["b"]) + + self.assertEqual(self.store.list_jobs(status="failed"), []) + + def test_list_invalid_status_raises(self): + with self.assertRaises(ValueError): + self.store.list_jobs(status="bogus") + + def test_delete_removes_job(self): + self.store.create("a", temp_dir="/tmp/a", now=T0) + self.store.delete("a") + self.assertIsNone(self.store.get("a")) + self.assertEqual(self.store.list_jobs(), []) + + def test_delete_unknown_is_noop(self): + self.store.delete("ghost") # must not raise + + +class TestDurability(JobStoreTestCase): + def test_job_survives_store_reopen(self): + self.store.create("job-1", temp_dir="/tmp/j", now=T0) + self.store.set_status( + "job-1", "done", now=T1, + output_path="/tmp/out.mp4", output_name="video.mp4", + ) + del self.store # simulate process exit + + reopened = JobStore(self.db_path) + job = reopened.get("job-1") + self.assertIsNotNone(job) + self.assertEqual(job["status"], "done") + self.assertEqual(job["output_path"], "/tmp/out.mp4") + self.assertEqual(job["output_name"], "video.mp4") + self.assertEqual(job["created_at"], T0.isoformat()) + self.assertEqual(job["updated_at"], T1.isoformat()) + + +class TestConcurrency(JobStoreTestCase): + def test_concurrent_set_status_is_consistent(self): + num_jobs = 8 + num_threads = 8 + for i in range(num_jobs): + self.store.create(f"job-{i}", temp_dir=f"/tmp/{i}", now=T0) + + errors = [] + barrier = threading.Barrier(num_threads) + + def worker(thread_idx): + try: + barrier.wait(timeout=10) + for i in range(num_jobs): + status = "done" if (i + thread_idx) % 2 else "processing" + self.store.set_status(f"job-{i}", status, now=T1) + except Exception as exc: # pragma: no cover - failure path + errors.append(exc) + + threads = [ + threading.Thread(target=worker, args=(t,)) + for t in range(num_threads) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=30) + + self.assertEqual(errors, []) + for i in range(num_jobs): + job = self.store.get(f"job-{i}") + # Every job ends in exactly one of the two written states, + # with the updated timestamp applied — no torn writes. + self.assertIn(job["status"], {"processing", "done"}) + self.assertEqual(job["updated_at"], T1.isoformat()) + self.assertEqual(job["created_at"], T0.isoformat()) + + +if __name__ == "__main__": + unittest.main()