From 6add2780e83d1e5fb522a4baf9b02aa9be9e197c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 16 Jun 2026 15:47:41 +0300 Subject: [PATCH] Make ingester job retry idempotent against a live sibling --- CHANGELOG.md | 1 + .../haiku/rag/ingester/queue/repository.py | 45 ++++++++++++++++--- tests/ingester/test_api.py | 33 +++++++++++--- tests/ingester/test_queue.py | 25 +++++++++++ 4 files changed, 91 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a40739be..99c71774 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ - SQLite ingester queue runs with a multi-connection pool (`pool_size=5, max_overflow=5`) instead of a single connection. API reads (`/stats`, `/jobs`) no longer time out with `QueuePool limit of size 1 reached` while workers hold the connection. - Permanently-failed ingester documents (revisioned sources) record their revision in `sync_state`, so discovery no longer re-enqueues them every sweep; re-attempted only when the file's revision changes or via explicit retry/rebuild. +- Retrying a dead ingester job (`/jobs/{id}/retry`, `/dlq/{id}/retry`) when a live job already exists for the same `(source_id, uri)` returns the live job instead of failing with a 500 (`uq_jobs_live` violation). ## [0.58.0] - 2026-06-15 diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index a4dc8c76..9ad32d84 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py @@ -7,6 +7,7 @@ from datetime import UTC, datetime, timedelta import sqlalchemy as sa from sqlalchemy.dialects import postgresql from sqlalchemy.dialects import sqlite as sqlite_dialect +from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncEngine from haiku.rag.ingester.queue.db import jobs, sync_state @@ -242,7 +243,12 @@ class JobRepo: error cleared, scheduled for immediate re-claim. Refuses `claimed` rows (would race with the worker still processing) and `succeeded` rows (re-ingest via UPSERT instead). Raises KeyError when the row - is missing or in a non-retryable state.""" + is missing or in a non-retryable state. + + Idempotent against a live sibling: if a live (queued/claimed) job + already exists for the same (source_id, uri), reviving the dead row + would violate `uq_jobs_live`; instead the existing live job is returned + and the dead row is left dead.""" now = _utcnow_iso() stmt = ( sa.update(jobs) @@ -258,11 +264,38 @@ class JobRepo: ) .returning(*jobs.c) ) - async with self._engine.begin() as conn: - row = (await conn.execute(stmt)).mappings().one_or_none() - if not row: - raise KeyError(f"Job {job_id!r} not found or not retryable") - return _row_to_job(row) + collided = False + try: + async with self._engine.begin() as conn: + row = (await conn.execute(stmt)).mappings().one_or_none() + except IntegrityError: + collided = True + row = None + if row is not None: + return _row_to_job(row) + if collided: + target = await self.get_job(job_id) + if target is not None: + live = await self._live_sibling(target.source_id, target.uri) + if live is not None: + return live + raise KeyError(f"Job {job_id!r} not found or not retryable") + + async def _live_sibling(self, source_id: str, uri: str) -> Job | None: + """The live (queued/claimed) job for a (source_id, uri), if any. + `uq_jobs_live` guarantees at most one.""" + query = ( + sa.select(jobs) + .where( + jobs.c.source_id == source_id, + jobs.c.uri == uri, + jobs.c.status.in_(["queued", "claimed"]), + ) + .limit(1) + ) + async with self._engine.connect() as conn: + row = (await conn.execute(query)).mappings().one_or_none() + return _row_to_job(row) if row else None async def cancel(self, job_id: str) -> bool: """True iff a queued/claimed row was removed; terminal jobs aren't diff --git a/tests/ingester/test_api.py b/tests/ingester/test_api.py index fdf5b001..fa5b365a 100644 --- a/tests/ingester/test_api.py +++ b/tests/ingester/test_api.py @@ -139,7 +139,7 @@ async def test_protected_endpoint_accepts_correct_token(state): @pytest.mark.asyncio -async def test_no_auth_token_allows_everything(state, jobs): +async def test_no_auth_token_allows_everything(state): async with _client(state, auth_token=None) as client: assert (await client.get("/jobs")).status_code == 200 assert (await client.get("/health")).status_code == 200 @@ -332,6 +332,24 @@ async def test_dlq_retry_404_on_missing_job(state): assert resp.status_code == 404 +@pytest.mark.asyncio +async def test_dlq_retry_with_live_sibling_returns_200(state, jobs): + """Retrying a dead job when a live job already exists for the same + (source_id, uri) returns 200 with the live job, not a 500.""" + first = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert first is not None + claimed = await jobs.claim_next("w") + assert claimed is not None + await jobs.mark_dead(first.id, "err", "w") + live = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert live is not None + + async with _client(state) as client: + resp = await client.post(f"/dlq/{first.id}/retry") + assert resp.status_code == 200 + assert resp.json()["id"] == live.id + + # --- /sources --- @@ -340,16 +358,17 @@ class _StubSource: self.source_id = source_id self._sweeps = list(sweeps) - def supports(self, uri): # pragma: no cover + def supports(self, _uri): # pragma: no cover return True - async def head(self, uri): # pragma: no cover + async def head(self, _uri): # pragma: no cover return None - async def fetch(self, uri) -> FetchResult: # pragma: no cover + async def fetch(self, _uri) -> FetchResult: # pragma: no cover raise NotImplementedError async def discover(self, since=None, *, known_uris=None): + del since, known_uris # interface-required, unused by this stub events = self._sweeps.pop(0) if self._sweeps else [] for event in events: yield event @@ -501,7 +520,7 @@ async def test_providers_probes_each_docling_serve_url(state, monkeypatch): from haiku.rag.ingester.api.routes import providers as providers_mod from haiku.rag.ingester.api.schemas import ProviderEndpoint - async def _fake_probe(client, base_url): + async def _fake_probe(_client, base_url): if "down" in base_url: return ProviderEndpoint( base_url=base_url, @@ -539,7 +558,7 @@ async def test_providers_skips_docling_serve_when_not_in_use(state, monkeypatch) probed: list[str] = [] - async def _spy_probe(client, base_url): # pragma: no cover - asserted not called + async def _spy_probe(_client, base_url): # pragma: no cover - asserted not called probed.append(base_url) raise AssertionError("probe should not run when docling-serve is not in use") @@ -562,7 +581,7 @@ async def test_providers_probes_when_only_chunker_uses_docling_serve( from haiku.rag.ingester.api.routes import providers as providers_mod from haiku.rag.ingester.api.schemas import ProviderEndpoint - async def _fake_probe(client, base_url): + async def _fake_probe(_client, base_url): return ProviderEndpoint(base_url=base_url, reachable=True, status_code=200) monkeypatch.setattr(providers_mod, "_probe", _fake_probe) diff --git a/tests/ingester/test_queue.py b/tests/ingester/test_queue.py index c6419758..bf020c51 100644 --- a/tests/ingester/test_queue.py +++ b/tests/ingester/test_queue.py @@ -530,6 +530,31 @@ async def test_retry_refuses_claimed_job(jobs): assert refreshed.attempts == 1 +@pytest.mark.asyncio +async def test_retry_returns_live_sibling_instead_of_colliding(jobs): + """Retrying a dead job when a live job already exists for the same + (source_id, uri) is idempotent: it returns the live sibling rather than + violating uq_jobs_live (which previously surfaced as a 500).""" + first = await jobs.enqueue("s", "u", JobOp.UPSERT) + assert first is not None + claimed = await jobs.claim_next("w") + assert claimed is not None + await jobs.mark_dead(claimed.id, "boom", "w") + + # A fresh live job for the same (source, uri) — e.g. re-discovered by a sweep. + live = await jobs.enqueue("s", "u", JobOp.UPSERT) + assert live is not None + assert live.id != first.id + + result = await jobs.retry(first.id) + assert result.id == live.id + assert result.status is JobStatus.QUEUED + # The originally-dead row stays dead — not duplicated into a second live job. + refreshed = await jobs.get_job(first.id) + assert refreshed is not None + assert refreshed.status is JobStatus.DEAD + + @pytest.mark.asyncio async def test_retry_refuses_succeeded_job(jobs): """Succeeded rows should be re-ingested through the UPSERT path, not