import asyncio from contextlib import nullcontext from typing import TYPE_CHECKING import httpx from pydantic import BaseModel from haiku.rag.client.exceptions import UnsupportedSourceError from haiku.rag.ingester.exceptions import PermanentError, TransientError from haiku.rag.ingester.queue.models import Job, JobOp from haiku.rag.ingester.sources.base import FileTooLargeError from haiku.rag.ingester.sources.registry import resolve_configured_source from haiku.rag.telemetry import attach_context, logfire if TYPE_CHECKING: from collections.abc import Mapping from haiku.rag.client import HaikuRAG from haiku.rag.ingester.metadata import MetadataProvider from haiku.rag.ingester.sources.base import Source 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 # UnsupportedSourceError is the typed signal from client/* that the # source will never ingest successfully on a retry (bad URI scheme, # missing file, unsupported extension, etc.). if isinstance(exc, UnsupportedSourceError | FileTooLargeError): return PermanentError(str(exc)) if isinstance(exc, ValueError): # Some downstream libraries (e.g. docling) raise plain ValueError # for "couldn't parse this file"; default to transient so the queue # retries up to max_attempts in case the issue is intermittent. return TransientError(str(exc)) 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.TransportError): # Umbrella for ConnectError, NetworkError, TimeoutException, ProtocolError, # ProxyError — every transport-layer failure that's worth retrying. return TransientError(f"network: {exc}") if isinstance(exc, FileNotFoundError): return PermanentError(f"file not found: {exc}") if isinstance(exc, PermissionError): return PermanentError(f"permission denied: {exc}") if isinstance(exc, IsADirectoryError | NotADirectoryError): return PermanentError(f"path error: {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}") async def run_job( client: "HaikuRAG", job: Job, *, sources: list["Source"] | None = None, metadata_providers: "Mapping[str, MetadataProvider] | None" = None, ) -> JobResult: """Execute the work described by `job`. `sources` is the list of configured Source adapters; the client looks up `job.source_id` against it via `resolve_configured_source` so workers reuse the authenticated/pre-configured fetch context the pollers used at discovery. `metadata_providers` maps source_id to a provider whose output is attached as document metadata on UPSERT. Raises PermanentError or TransientError; the worker uses that to decide dead vs retry.""" extra = job.extra or {} parent_ctx = extra.get("_otel") attach = attach_context(parent_ctx) if parent_ctx else nullcontext() with ( attach, logfire.span( "ingester.job", job_id=job.id, source_id=job.source_id, uri=job.uri, op=job.op.value, attempt=job.attempts, ), ): try: manifest_context = extra.get("_manifest") if job.op is JobOp.DELETE: # An atomic-rename save can let a spurious DELETE win the # enqueue race while the file is mid-rewrite. If the resource # is already back, skip the delete (it would blackhole a live # document) and let the next sweep re-ingest it. Manifest # replay intentionally follows the frozen dry-run changeset. if manifest_context is None: try: source = resolve_configured_source( job.uri, job.source_id, sources ) restored = await source.head(job.uri) is not None except Exception: restored = False if restored: return JobResult(deleted=False) 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) if manifest_context is not None and job.revision is not None: source = resolve_configured_source(job.uri, job.source_id, sources) current_revision = await source.head(job.uri) if current_revision != job.revision: raise PermanentError( "manifest revision is stale for " f"{job.uri}: expected {job.revision!r}, " f"current {current_revision!r}" ) result = await client.create_document_from_source( job.uri, sources=sources, source_id=job.source_id, metadata_provider=(metadata_providers or {}).get(job.source_id), ) # 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("source_revision"), content_hash=metadata.get("md5"), ) except Exception as exc: # CancelledError, KeyboardInterrupt, SystemExit are BaseException # subclasses; they signal the runtime is shutting us down, not a # job-level failure, so we let them propagate untouched. raise _classify(exc) from exc