Make the SQLite job claim atomic across processes

haiku_rag_slim/haiku/rag/ingester/queue/repository.py

@@ -127,34 +127,34 @@ class JobRepo:
 
     async def claim_next(self, worker_id: str) -> Job | None:
         """Atomically claim the oldest queued job whose scheduled_at <= now.
-        A `SELECT ... FOR UPDATE SKIP LOCKED` picks the row (multi-process safe
-        on Postgres; the clause is omitted on SQLite, where pool_size=1 keeps
-        the select-then-update atomic), then a guarded UPDATE claims it."""
+        A single `UPDATE ... WHERE id = (SELECT ... LIMIT 1) RETURNING` keeps
+        the claim atomic across connections: on Postgres the subquery adds
+        `FOR UPDATE SKIP LOCKED`; on SQLite the whole statement runs under one
+        write lock, so a racing connection re-evaluates the subquery against
+        the committed state and finds the row already claimed."""
         now = _utcnow_iso()
-        select_candidate = (
+        candidate = (
             sa.select(jobs.c.id)
             .where(jobs.c.status == "queued", jobs.c.scheduled_at <= now)
-            .order_by(jobs.c.scheduled_at)
+            .order_by(jobs.c.scheduled_at, jobs.c.id)
             .limit(1)
             .with_for_update(skip_locked=True)
+            .scalar_subquery()
+        )
+        claim = (
+            sa.update(jobs)
+            .where(jobs.c.id == candidate)
+            .values(
+                status="claimed",
+                claimed_at=now,
+                claimed_by=worker_id,
+                attempts=jobs.c.attempts + 1,
+            )
+            .returning(*jobs.c)
         )
         async with self._engine.begin() as conn:
-            job_id = (await conn.execute(select_candidate)).scalar_one_or_none()
-            if job_id is None:
-                return None
-            claim = (
-                sa.update(jobs)
-                .where(jobs.c.id == job_id)
-                .values(
-                    status="claimed",
-                    claimed_at=now,
-                    claimed_by=worker_id,
-                    attempts=jobs.c.attempts + 1,
-                )
-                .returning(*jobs.c)
-            )
-            row = (await conn.execute(claim)).mappings().one()
-        return _row_to_job(row)
+            row = (await conn.execute(claim)).mappings().one_or_none()
+        return _row_to_job(row) if row else None
 
     async def get_job(self, job_id: str) -> Job | None:
         async with self._engine.connect() as conn:

tests/ingester/test_queue.py

@@ -7,6 +7,7 @@ import sqlalchemy as sa
 from haiku.rag.config import QueueConfig
 from haiku.rag.ingester.queue.migrations import apply_migrations, open_queue
 from haiku.rag.ingester.queue.models import JobOp, JobStatus, SyncRow
+from haiku.rag.ingester.queue.repository import JobRepo
 
 # --- migrations / schema ---
 
@@ -229,6 +230,43 @@ async def test_claim_next_atomic_under_concurrency(jobs):
     assert {c.id for c in claimed} == {j.id for j in enqueued}
 
 
+@pytest.mark.asyncio
+async def test_claim_next_atomic_across_independent_engines(tmp_path):
+    """Two engines on one SQLite file stand in for two ingester processes.
+    The claim is a single UPDATE...WHERE id=(subquery), so it stays atomic
+    across connections — pool_size=1 only serializes within one engine."""
+    cfg = QueueConfig(path=tmp_path / "shared-queue.db")
+    engine_a = await open_queue(cfg)
+    engine_b = await open_queue(cfg)
+    try:
+        repo_a = JobRepo(engine_a)
+        repo_b = JobRepo(engine_b)
+
+        enqueued = []
+        for i in range(10):
+            job = await repo_a.enqueue("s", f"u{i}", JobOp.UPSERT)
+            assert job is not None
+            enqueued.append(job)
+
+        # Claims alternate across the two engines, contending on the same file.
+        results = await asyncio.gather(
+            *((repo_a if i % 2 == 0 else repo_b).claim_next(f"w{i}") for i in range(20))
+        )
+        claimed = [r for r in results if r is not None]
+
+        assert len(claimed) == 10
+        assert len({c.id for c in claimed}) == 10
+        assert {c.id for c in claimed} == {j.id for j in enqueued}
+        for job in enqueued:
+            refreshed = await repo_a.get_job(job.id)
+            assert refreshed is not None
+            assert refreshed.status is JobStatus.CLAIMED
+            assert refreshed.attempts == 1
+    finally:
+        await engine_a.dispose()
+        await engine_b.dispose()
+
+
 # --- terminal transitions ---