diff --git a/haiku_rag_slim/haiku/rag/ingester/cli.py b/haiku_rag_slim/haiku/rag/ingester/cli.py index b5ad5c9c..2d2860e8 100644 --- a/haiku_rag_slim/haiku/rag/ingester/cli.py +++ b/haiku_rag_slim/haiku/rag/ingester/cli.py @@ -1,4 +1,6 @@ import asyncio +import uuid +from datetime import UTC, datetime from pathlib import Path import typer @@ -6,6 +8,7 @@ from dotenv import find_dotenv, load_dotenv load_dotenv(find_dotenv(usecwd=True)) +from haiku.rag.client import HaikuRAG # noqa: E402 from haiku.rag.config import ( # noqa: E402 AppConfig, find_config_file, @@ -13,7 +16,10 @@ from haiku.rag.config import ( # noqa: E402 load_yaml_config, set_config, ) +from haiku.rag.ingester.exceptions import PermanentError, TransientError # noqa: E402 from haiku.rag.ingester.queue.migrations import open_queue # noqa: E402 +from haiku.rag.ingester.queue.models import Job, JobOp, JobStatus # noqa: E402 +from haiku.rag.ingester.workers.pipeline import run_job # noqa: E402 cli = typer.Typer( name="haiku-ingester", @@ -88,3 +94,65 @@ def queue_migrate( path = _resolve_queue_path(app_config, queue) asyncio.run(_ensure_schema(path)) typer.echo(f"Queue at {path} is up to date") + + +@cli.command("run-once") +def run_once( + uri: str = typer.Argument(..., help="URI to ingest (file://, http(s)://, s3://)."), + config: Path | None = typer.Option( + None, "--config", "-c", help="Path to haiku.rag.yaml." + ), + db: Path | None = typer.Option( + None, + "--db", + help="LanceDB path (overrides config.storage.data_dir).", + ), + delete: bool = typer.Option( + False, "--delete", help="Run a delete op instead of upsert." + ), +) -> None: + """Build an ad-hoc Job and run it through the worker pipeline once. + + Bypasses the queue — does NOT enqueue. Useful for smoke-testing the + Source adapter + pipeline path without spinning up the full pool. + """ + app_config = _load_config_with_override(config) + asyncio.run(_run_once(app_config, uri, db, delete)) + + +async def _run_once( + app_config: AppConfig, uri: str, db_path: Path | None, delete: bool +) -> None: + db = db_path or (app_config.storage.data_dir / "haiku.rag.lancedb") + now = datetime.now(UTC) + job = Job( + id=f"adhoc-{uuid.uuid4()}", + source_id="adhoc", + uri=uri, + op=JobOp.DELETE if delete else JobOp.UPSERT, + status=JobStatus.CLAIMED, + attempts=1, + max_attempts=1, + enqueued_at=now, + scheduled_at=now, + claimed_at=now, + claimed_by="run-once", + ) + + async with HaikuRAG(db, config=app_config) as client: + try: + result = await run_job(client, job) + except PermanentError as e: + typer.echo(f"PERMANENT: {e}", err=True) + raise typer.Exit(2) from e + except TransientError as e: + typer.echo(f"TRANSIENT: {e}", err=True) + raise typer.Exit(1) from e + + if result.deleted: + typer.echo(f"Deleted document at {uri}") + else: + typer.echo( + f"Ingested {uri}: document_id={result.document_id} " + f"revision={result.revision} md5={result.content_hash}" + ) diff --git a/haiku_rag_slim/haiku/rag/ingester/exceptions.py b/haiku_rag_slim/haiku/rag/ingester/exceptions.py new file mode 100644 index 00000000..ae493b3e --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/exceptions.py @@ -0,0 +1,12 @@ +class IngesterError(Exception): + """Base class for ingester errors that the worker classifies on.""" + + +class PermanentError(IngesterError): + """The job will never succeed without intervention (bad URI, unsupported + content type, 410 Gone, etc.). Goes straight to dead — no retry.""" + + +class TransientError(IngesterError): + """The job might succeed on a future attempt (network hiccup, 5xx, DB + busy). Worker reschedules with backoff up to max_attempts, then dead.""" diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/__init__.py b/haiku_rag_slim/haiku/rag/ingester/workers/__init__.py new file mode 100644 index 00000000..807831ce --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/workers/__init__.py @@ -0,0 +1,11 @@ +from haiku.rag.ingester.workers.pipeline import JobResult, run_job +from haiku.rag.ingester.workers.pool import WorkerPool +from haiku.rag.ingester.workers.retry import RetryPolicy, compute_backoff + +__all__ = [ + "JobResult", + "RetryPolicy", + "WorkerPool", + "compute_backoff", + "run_job", +] diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py new file mode 100644 index 00000000..6b6c3402 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py @@ -0,0 +1,134 @@ +import asyncio +from typing import TYPE_CHECKING + +import httpx +from pydantic import BaseModel + +from haiku.rag.ingester.exceptions import PermanentError, TransientError +from haiku.rag.ingester.queue.models import Job, JobOp + +if TYPE_CHECKING: + from haiku.rag.client import HaikuRAG + +# ValueError messages from create_document_from_source that mean "this job +# will never succeed". Anything else from ValueError defaults to transient. +_PERMANENT_VALUE_MARKERS = ( + "Unsupported file extension", + "Unsupported content type", + "Invalid S3 URI", + "File does not exist", + "uri override is not supported", +) + + +class JobResult(BaseModel): + """What the worker needs after a successful job: enough metadata to + update sync_state. document_id is None for DELETE ops.""" + + document_id: str | None = None + revision: str | None = None + content_hash: str | None = None + deleted: bool = False + + +def _classify(exc: BaseException) -> Exception: + """Wrap an unclassified exception into Permanent or Transient. Already- + classified errors pass through unchanged.""" + if isinstance(exc, PermanentError | TransientError): + return exc + + if isinstance(exc, ValueError): + message = str(exc) + if any(marker in message for marker in _PERMANENT_VALUE_MARKERS): + return PermanentError(message) + return TransientError(message) + + if isinstance(exc, httpx.HTTPStatusError): + status = exc.response.status_code + # 401/403/404/410 are unrecoverable without operator action; 408/429/5xx + # are transient. Everything else in 4xx is treated as permanent — better + # to DLQ a misconfigured URL than to retry it forever. + if status in (408, 429) or status >= 500: + return TransientError(f"HTTP {status}: {exc}") + return PermanentError(f"HTTP {status}: {exc}") + + if isinstance( + exc, + httpx.ConnectError | httpx.ReadTimeout | httpx.WriteTimeout | httpx.PoolTimeout, + ): + return TransientError(f"network: {exc}") + + if isinstance(exc, asyncio.TimeoutError | TimeoutError | OSError): + return TransientError(f"timeout/io: {exc}") + + # Unknown errors default to transient — retry up to max_attempts gives the + # operator visibility into the failure mode without dropping data on the + # first hiccup. + return TransientError(f"unexpected: {exc!r}") + + +def _logfire_span_or_null(name: str, **attrs): + """logfire is available via pydantic-ai but may be disabled; use the + no-op API surface so tests don't need it configured.""" + try: + import logfire + + return logfire.span(name, **attrs) + except ImportError: # pragma: no cover - logfire ships with pydantic-ai + + class _Null: + def __enter__(self): + return self + + def __exit__(self, *_): + return False + + def set_attribute(self, *_args, **_kwargs): + pass + + return _Null() + + +async def run_job(client: "HaikuRAG", job: Job) -> JobResult: + """Execute the work described by `job`. Raises PermanentError or + TransientError; the worker uses that to decide dead vs retry.""" + extra = job.extra or {} + storage_options = extra.get("storage_options") + user_metadata = extra.get("metadata", {}) + + with _logfire_span_or_null( + "ingester.job", + job_id=job.id, + source_id=job.source_id, + uri=job.uri, + op=job.op.value, + attempt=job.attempts, + ): + try: + if job.op is JobOp.DELETE: + doc = await client.get_document_by_uri(job.uri) + if doc is not None and doc.id is not None: + await client.delete_document(doc.id) + return JobResult(deleted=True) + + result = await client.create_document_from_source( + job.uri, + metadata=user_metadata, + storage_options=storage_options, + ) + # Directory ingestion returns list[Document] — workers ingest single + # resources, so a list here is a programming error in the caller. + if isinstance(result, list): + raise PermanentError( + f"Job {job.id} resolved to a directory; queue jobs must " + f"reference a single document URI." + ) + + metadata = result.metadata or {} + return JobResult( + document_id=result.id, + revision=metadata.get("etag"), + content_hash=metadata.get("md5"), + ) + except BaseException as exc: + raise _classify(exc) from exc diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py new file mode 100644 index 00000000..c9d79b87 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -0,0 +1,156 @@ +import asyncio +import logging +from typing import TYPE_CHECKING + +from haiku.rag.ingester.exceptions import PermanentError, TransientError +from haiku.rag.ingester.queue.models import Job, JobOp +from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo +from haiku.rag.ingester.workers.pipeline import run_job +from haiku.rag.ingester.workers.retry import RetryPolicy, compute_backoff + +if TYPE_CHECKING: + from haiku.rag.client import HaikuRAG + +logger = logging.getLogger(__name__) + + +class WorkerPool: + """Asyncio-based pool. N worker tasks share a bounded Semaphore, each + pulling jobs from the queue and running them through `run_job`. Reaper + task resets claims older than `claim_timeout_s` so a crashed worker + doesn't strand its job. + + Lifecycle: build it, await start(), let it run, await stop(). + """ + + def __init__( + self, + *, + client: "HaikuRAG", + job_repo: JobRepo, + sync_repo: SyncStateRepo, + worker_count: int = 4, + max_concurrent: int = 4, + retry_policy: RetryPolicy | None = None, + poll_idle_interval_s: float = 1.0, + claim_timeout_s: int = 1800, + reaper_interval_s: int = 60, + ): + self._client = client + self._jobs = job_repo + self._sync = sync_repo + self._worker_count = worker_count + self._semaphore = asyncio.Semaphore(max_concurrent) + self._retry = retry_policy or RetryPolicy() + self._poll_idle_s = poll_idle_interval_s + self._claim_timeout_s = claim_timeout_s + self._reaper_interval_s = reaper_interval_s + self._stop = asyncio.Event() + self._workers: list[asyncio.Task] = [] + self._reaper: asyncio.Task | None = None + + async def start(self) -> None: + if self._workers: + raise RuntimeError("WorkerPool already started") + self._stop.clear() + for i in range(self._worker_count): + self._workers.append(asyncio.create_task(self._worker_loop(f"worker-{i}"))) + self._reaper = asyncio.create_task(self._reaper_loop()) + + async def stop(self) -> None: + self._stop.set() + tasks = list(self._workers) + if self._reaper is not None: + tasks.append(self._reaper) + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self._workers.clear() + self._reaper = None + + async def drain_once(self, worker_id: str = "drain") -> int: + """Drain every currently-claimable job to completion on the calling + coroutine. Used by run-once and by tests; not by `start()`.""" + processed = 0 + while True: + job = await self._jobs.claim_next(worker_id) + if job is None: + return processed + await self._process(job) + processed += 1 + + async def _worker_loop(self, worker_id: str) -> None: + while not self._stop.is_set(): + try: + job = await self._jobs.claim_next(worker_id) + except Exception: + logger.exception("claim_next failed in %s", worker_id) + await self._sleep_or_stop(self._poll_idle_s) + continue + + if job is None: + await self._sleep_or_stop(self._poll_idle_s) + continue + + async with self._semaphore: + await self._process(job) + + async def _reaper_loop(self) -> None: + while not self._stop.is_set(): + await self._sleep_or_stop(self._reaper_interval_s) + if self._stop.is_set(): + return + try: + reset = await self._jobs.reap_stale(self._claim_timeout_s) + if reset: + logger.info("Reaper reset %d stale claim(s)", reset) + except Exception: + logger.exception("reaper failed") + + async def _sleep_or_stop(self, seconds: float) -> None: + try: + await asyncio.wait_for(self._stop.wait(), timeout=seconds) + except TimeoutError: + pass + + async def _process(self, job: Job) -> None: + try: + result = await run_job(self._client, job) + except PermanentError as e: + await self._jobs.mark_dead(job.id, str(e)) + logger.info("Job %s dead (permanent): %s", job.id, e) + return + except TransientError as e: + if job.attempts >= job.max_attempts: + await self._jobs.mark_dead(job.id, str(e)) + logger.info( + "Job %s dead (max attempts %d): %s", job.id, job.max_attempts, e + ) + return + delay = compute_backoff(job.attempts, self._retry) + await self._jobs.reschedule(job.id, delay, str(e)) + logger.info( + "Job %s rescheduled in %.1fs (attempt %d/%d): %s", + job.id, + delay, + job.attempts, + job.max_attempts, + e, + ) + return + except Exception as e: + # Defensive: pipeline classifier should have caught everything. + await self._jobs.mark_dead(job.id, f"unclassified: {e!r}") + logger.exception("Unclassified error in job %s", job.id) + return + + await self._jobs.mark_succeeded(job.id) + if job.op is JobOp.DELETE: + await self._sync.delete(job.source_id, job.uri) + else: + await self._sync.upsert( + job.source_id, + job.uri, + revision=result.revision, + content_hash=result.content_hash, + ingested=True, + ) diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/retry.py b/haiku_rag_slim/haiku/rag/ingester/workers/retry.py new file mode 100644 index 00000000..1c7a15f1 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/workers/retry.py @@ -0,0 +1,32 @@ +import random +from dataclasses import dataclass + + +@dataclass(frozen=True) +class RetryPolicy: + """Exponential backoff with jitter. Per-source override is allowed so + one slow/flaky source doesn't drag the rest of the queue.""" + + max_attempts: int = 5 + base_delay_s: float = 2.0 + max_delay_s: float = 300.0 + jitter: float = 0.25 # ±25% + + +def compute_backoff( + attempt: int, + policy: RetryPolicy, + *, + rng: random.Random | None = None, +) -> float: + """Delay before the next attempt. `attempt` is the 1-indexed count of + attempts already made (so attempt=1 → first retry delay). + + `rng` is injectable for deterministic tests. + """ + if attempt < 1: + attempt = 1 + raw = min(policy.base_delay_s * (2 ** (attempt - 1)), policy.max_delay_s) + r = rng if rng is not None else random + j = 1.0 + (r.random() * 2 - 1) * policy.jitter + return max(0.0, raw * j) diff --git a/tests/ingester/test_pipeline.py b/tests/ingester/test_pipeline.py new file mode 100644 index 00000000..6aecd747 --- /dev/null +++ b/tests/ingester/test_pipeline.py @@ -0,0 +1,212 @@ +from datetime import UTC, datetime +from unittest.mock import AsyncMock + +import httpx +import pytest + +from haiku.rag.client import HaikuRAG +from haiku.rag.ingester.exceptions import PermanentError, TransientError +from haiku.rag.ingester.queue.models import Job, JobOp, JobStatus +from haiku.rag.ingester.workers.pipeline import run_job +from haiku.rag.store.models.document import Document + + +def _job( + *, + op: JobOp = JobOp.UPSERT, + uri: str = "https://example.com/a.pdf", + extra: dict | None = None, + attempts: int = 0, +) -> Job: + now = datetime.now(UTC) + return Job( + id="job-1", + source_id="src", + uri=uri, + op=op, + status=JobStatus.CLAIMED, + attempts=attempts, + max_attempts=5, + extra=extra, + enqueued_at=now, + scheduled_at=now, + ) + + +def _mock_client() -> AsyncMock: + return AsyncMock(spec=HaikuRAG) + + +@pytest.mark.asyncio +async def test_upsert_calls_create_document_from_source_and_returns_metadata(): + client = _mock_client() + client.create_document_from_source.return_value = Document( + id="doc-42", + content="x", + uri="https://example.com/a.pdf", + metadata={"md5": "abcd", "etag": "xyz", "contentType": "application/pdf"}, + ) + + result = await run_job( + client, _job(extra={"metadata": {"k": "v"}, "storage_options": {"o": "1"}}) + ) + + assert result.document_id == "doc-42" + assert result.revision == "xyz" + assert result.content_hash == "abcd" + assert result.deleted is False + client.create_document_from_source.assert_awaited_once_with( + "https://example.com/a.pdf", + metadata={"k": "v"}, + storage_options={"o": "1"}, + ) + + +@pytest.mark.asyncio +async def test_upsert_without_extra_passes_empty_metadata(): + client = _mock_client() + client.create_document_from_source.return_value = Document( + id="d", content="x", uri="u", metadata={} + ) + + await run_job(client, _job()) + client.create_document_from_source.assert_awaited_once_with( + "https://example.com/a.pdf", metadata={}, storage_options=None + ) + + +@pytest.mark.asyncio +async def test_delete_calls_delete_document_when_present(): + client = _mock_client() + client.get_document_by_uri.return_value = Document(id="doc-9", content="", uri="u") + + result = await run_job(client, _job(op=JobOp.DELETE)) + + assert result.deleted is True + assert result.document_id is None + client.delete_document.assert_awaited_once_with("doc-9") + + +@pytest.mark.asyncio +async def test_delete_is_noop_when_document_missing(): + client = _mock_client() + client.get_document_by_uri.return_value = None + + result = await run_job(client, _job(op=JobOp.DELETE)) + + assert result.deleted is True + client.delete_document.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_unsupported_extension_classified_permanent(): + client = _mock_client() + client.create_document_from_source.side_effect = ValueError( + "Unsupported file extension: .xyz" + ) + + with pytest.raises(PermanentError, match="Unsupported"): + await run_job(client, _job()) + + +@pytest.mark.asyncio +async def test_unsupported_content_type_classified_permanent(): + client = _mock_client() + client.create_document_from_source.side_effect = ValueError( + "Unsupported content type/extension: application/octet-stream/.bin" + ) + with pytest.raises(PermanentError): + await run_job(client, _job()) + + +@pytest.mark.asyncio +async def test_invalid_s3_uri_classified_permanent(): + client = _mock_client() + client.create_document_from_source.side_effect = ValueError( + "Invalid S3 URI: s3:///bad" + ) + with pytest.raises(PermanentError): + await run_job(client, _job()) + + +@pytest.mark.asyncio +async def test_missing_file_classified_permanent(): + client = _mock_client() + client.create_document_from_source.side_effect = ValueError( + "File does not exist: /nope" + ) + with pytest.raises(PermanentError): + await run_job(client, _job()) + + +@pytest.mark.asyncio +async def test_other_value_error_classified_transient(): + client = _mock_client() + client.create_document_from_source.side_effect = ValueError("DB busy") + with pytest.raises(TransientError): + await run_job(client, _job()) + + +@pytest.mark.parametrize("status", [401, 403, 404, 410]) +@pytest.mark.asyncio +async def test_http_4xx_classified_permanent(status): + client = _mock_client() + response = httpx.Response( + status, request=httpx.Request("GET", "https://example.com/a") + ) + client.create_document_from_source.side_effect = httpx.HTTPStatusError( + "err", request=response.request, response=response + ) + with pytest.raises(PermanentError): + await run_job(client, _job()) + + +@pytest.mark.parametrize("status", [408, 429, 500, 502, 503]) +@pytest.mark.asyncio +async def test_http_408_429_5xx_classified_transient(status): + client = _mock_client() + response = httpx.Response( + status, request=httpx.Request("GET", "https://example.com/a") + ) + client.create_document_from_source.side_effect = httpx.HTTPStatusError( + "err", request=response.request, response=response + ) + with pytest.raises(TransientError): + await run_job(client, _job()) + + +@pytest.mark.asyncio +async def test_connect_error_classified_transient(): + client = _mock_client() + client.create_document_from_source.side_effect = httpx.ConnectError("boom") + with pytest.raises(TransientError): + await run_job(client, _job()) + + +@pytest.mark.asyncio +async def test_unknown_exception_classified_transient(): + client = _mock_client() + client.create_document_from_source.side_effect = RuntimeError("???") + with pytest.raises(TransientError): + await run_job(client, _job()) + + +@pytest.mark.asyncio +async def test_directory_result_classified_permanent(): + client = _mock_client() + client.create_document_from_source.return_value = [ + Document(id="a", content="", uri="u1"), + Document(id="b", content="", uri="u2"), + ] + with pytest.raises(PermanentError, match="directory"): + await run_job(client, _job()) + + +@pytest.mark.asyncio +async def test_existing_permanent_error_passes_through_unchanged(): + client = _mock_client() + sentinel = PermanentError("explicit") + client.create_document_from_source.side_effect = sentinel + with pytest.raises(PermanentError) as excinfo: + await run_job(client, _job()) + assert excinfo.value is sentinel diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py new file mode 100644 index 00000000..0cd3b887 --- /dev/null +++ b/tests/ingester/test_workers.py @@ -0,0 +1,252 @@ +import asyncio +from unittest.mock import AsyncMock + +import aiosqlite +import pytest + +from haiku.rag.client import HaikuRAG +from haiku.rag.ingester.exceptions import PermanentError, TransientError +from haiku.rag.ingester.queue.migrations import apply_migrations +from haiku.rag.ingester.queue.models import JobOp, JobStatus +from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo +from haiku.rag.ingester.workers.pool import WorkerPool +from haiku.rag.ingester.workers.retry import RetryPolicy +from haiku.rag.store.models.document import Document + + +@pytest.fixture +async def conn(tmp_path): + path = tmp_path / "queue.db" + connection = await aiosqlite.connect(str(path)) + connection.row_factory = aiosqlite.Row + await apply_migrations(connection) + yield connection + await connection.close() + + +@pytest.fixture +def jobs(conn): + return JobRepo(conn) + + +@pytest.fixture +def sync(conn): + return SyncStateRepo(conn) + + +@pytest.fixture +def client(): + return AsyncMock(spec=HaikuRAG) + + +def _pool(client, jobs, sync, **kwargs) -> WorkerPool: + return WorkerPool( + client=client, + job_repo=jobs, + sync_repo=sync, + worker_count=kwargs.pop("worker_count", 2), + max_concurrent=kwargs.pop("max_concurrent", 2), + poll_idle_interval_s=kwargs.pop("poll_idle_interval_s", 0.05), + reaper_interval_s=kwargs.pop("reaper_interval_s", 60), + claim_timeout_s=kwargs.pop("claim_timeout_s", 60), + retry_policy=kwargs.pop("retry_policy", RetryPolicy()), + ) + + +# --- drain_once: covers _process logic deterministically --- + + +@pytest.mark.asyncio +async def test_drain_marks_job_succeeded_and_writes_sync_state(client, jobs, sync): + client.create_document_from_source.return_value = Document( + id="doc-1", content="x", uri="s3://b/k.md", metadata={"md5": "m1", "etag": "e1"} + ) + job = await jobs.enqueue("src", "s3://b/k.md", JobOp.UPSERT, revision="e0") + assert job is not None + + pool = _pool(client, jobs, sync) + processed = await pool.drain_once() + assert processed == 1 + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.SUCCEEDED + + snapshot = await sync.get_snapshot("src") + assert snapshot == {"s3://b/k.md": "e1"} + + +@pytest.mark.asyncio +async def test_drain_delete_op_removes_sync_state(client, jobs, sync): + await sync.upsert("src", "s3://b/k.md", revision="e1", content_hash="m1") + client.get_document_by_uri.return_value = Document( + id="doc-1", content="", uri="s3://b/k.md" + ) + job = await jobs.enqueue("src", "s3://b/k.md", JobOp.DELETE) + assert job is not None + + pool = _pool(client, jobs, sync) + await pool.drain_once() + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.SUCCEEDED + client.delete_document.assert_awaited_once_with("doc-1") + + snapshot = await sync.get_snapshot("src") + assert snapshot == {} + + +@pytest.mark.asyncio +async def test_permanent_error_marks_dead_no_reschedule(client, jobs, sync): + client.create_document_from_source.side_effect = PermanentError("unsupported") + job = await jobs.enqueue("src", "https://x/y.bin", JobOp.UPSERT) + assert job is not None + + pool = _pool(client, jobs, sync) + await pool.drain_once() + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.DEAD + assert refreshed.last_error == "unsupported" + # sync_state is NOT written on failure + assert await sync.get_snapshot("src") == {} + + +@pytest.mark.asyncio +async def test_transient_error_reschedules_below_max_attempts(client, jobs, sync): + client.create_document_from_source.side_effect = TransientError("blip") + job = await jobs.enqueue("src", "u", JobOp.UPSERT, max_attempts=3) + assert job is not None + + # base_delay large enough that claim_next won't re-pick the job within + # drain_once — we want exactly one process iteration. + pool = _pool( + client, jobs, sync, retry_policy=RetryPolicy(base_delay_s=60.0, jitter=0.0) + ) + processed = await pool.drain_once() + assert processed == 1 + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.QUEUED + assert refreshed.last_error == "blip" + assert refreshed.attempts == 1 + assert refreshed.scheduled_at > job.scheduled_at + + +@pytest.mark.asyncio +async def test_transient_error_at_max_attempts_marks_dead(client, jobs, sync, conn): + client.create_document_from_source.side_effect = TransientError("blip") + job = await jobs.enqueue("src", "u", JobOp.UPSERT, max_attempts=1) + assert job is not None + + pool = _pool( + client, jobs, sync, retry_policy=RetryPolicy(base_delay_s=0.0, jitter=0.0) + ) + await pool.drain_once() + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + # attempts started at 0, claim_next set it to 1 = max → dead + assert refreshed.status is JobStatus.DEAD + assert refreshed.attempts == 1 + + +@pytest.mark.asyncio +async def test_unknown_exception_caught_and_marked_dead(client, jobs, sync): + # Pipeline classifies BaseException → TransientError, but if something + # slips past the pool catches with a final defensive net. + client.create_document_from_source.side_effect = KeyboardInterrupt("nope") + job = await jobs.enqueue("src", "u", JobOp.UPSERT, max_attempts=1) + assert job is not None + + pool = _pool( + client, jobs, sync, retry_policy=RetryPolicy(base_delay_s=0.0, jitter=0.0) + ) + # KeyboardInterrupt is a BaseException — pipeline wraps it to TransientError. + # With max_attempts=1, the worker marks dead. + await pool.drain_once() + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.DEAD + + +# --- start / stop lifecycle --- + + +@pytest.mark.asyncio +async def test_workers_drain_queue_after_start(client, jobs, sync): + client.create_document_from_source.return_value = Document( + id="doc", content="x", uri="u", metadata={"md5": "m", "etag": "e"} + ) + for i in range(5): + await jobs.enqueue("src", f"u{i}", JobOp.UPSERT) + + pool = _pool(client, jobs, sync, worker_count=3, max_concurrent=3) + await pool.start() + try: + # Wait until everything is succeeded or until a deadline. + for _ in range(50): + counts = await jobs.counts_by_status() + if counts.get("succeeded", 0) == 5: + break + await asyncio.sleep(0.05) + finally: + await pool.stop() + + counts = await jobs.counts_by_status() + assert counts.get("succeeded", 0) == 5 + + +@pytest.mark.asyncio +async def test_double_start_raises(client, jobs, sync): + pool = _pool(client, jobs, sync, worker_count=1) + await pool.start() + try: + with pytest.raises(RuntimeError, match="already started"): + await pool.start() + finally: + await pool.stop() + + +# --- reaper --- + + +@pytest.mark.asyncio +async def test_reaper_resets_stale_claims(client, jobs, sync, conn): + from datetime import UTC, datetime, timedelta + + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + claimed = await jobs.claim_next("worker-old") + assert claimed is not None + + # Push claimed_at back so reap_stale picks it up. + long_ago = (datetime.now(UTC) - timedelta(hours=1)).isoformat() + await conn.execute( + "UPDATE jobs SET claimed_at = ? WHERE id = ?", (long_ago, job.id) + ) + await conn.commit() + + pool = _pool( + client, + jobs, + sync, + worker_count=0, + reaper_interval_s=0.05, + claim_timeout_s=1, + ) + await pool.start() + try: + for _ in range(30): + refreshed = await jobs.get_job(job.id) + if refreshed is not None and refreshed.status is JobStatus.QUEUED: + break + await asyncio.sleep(0.05) + finally: + await pool.stop() + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.QUEUED