193 lines
7.5 KiB
Python
193 lines
7.5 KiB
Python
import asyncio
|
|
import logging
|
|
from datetime import UTC, datetime
|
|
|
|
from haiku.rag.config import SourceConfig
|
|
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
|
|
from haiku.rag.ingester.queue.models import JobOp
|
|
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
|
|
from haiku.rag.ingester.sources.base import (
|
|
Source,
|
|
SourceEvent,
|
|
SourceEventKind,
|
|
)
|
|
from haiku.rag.telemetry import get_context, logfire
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _enqueue_extra(cfg: SourceConfig) -> dict | None:
|
|
"""Per-source state worth carrying into the job (so the worker can rebuild
|
|
the same fetch context when it processes), plus the current logfire trace
|
|
context so the worker's `ingester.job` span nests under the
|
|
`ingester.poller.sweep` that enqueued it."""
|
|
extra: dict = {}
|
|
storage_options = getattr(cfg, "storage_options", None)
|
|
if storage_options:
|
|
extra["storage_options"] = dict(storage_options)
|
|
headers = getattr(cfg, "headers", None)
|
|
if headers:
|
|
extra["headers"] = dict(headers)
|
|
carrier = get_context()
|
|
if carrier:
|
|
extra["_otel"] = dict(carrier)
|
|
return extra or None
|
|
|
|
|
|
def _max_attempts(cfg: SourceConfig, default: int) -> int:
|
|
return cfg.retry.max_attempts if cfg.retry is not None else default
|
|
|
|
|
|
class BasePoller:
|
|
"""Shared lifecycle: build a discover() coroutine + process its events
|
|
into queue jobs and sync_state updates. Subclasses provide the loop
|
|
(FS uses watchfiles + initial discover; periodic uses sleep+discover)."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
source: Source,
|
|
config: SourceConfig,
|
|
job_repo: JobRepo,
|
|
sync_repo: SyncStateRepo,
|
|
breaker: CircuitBreaker | None = None,
|
|
default_max_attempts: int = 5,
|
|
):
|
|
self.source = source
|
|
self.config = config
|
|
self._jobs = job_repo
|
|
self._sync = sync_repo
|
|
self._breaker = breaker or CircuitBreaker(config.circuit_breaker)
|
|
self._stop = asyncio.Event()
|
|
self._task: asyncio.Task | None = None
|
|
self._last_polled_at: datetime | None = None
|
|
self._last_skip_reason: str | None = None
|
|
self._default_max_attempts = default_max_attempts
|
|
|
|
@property
|
|
def source_id(self) -> str:
|
|
return self.source.source_id
|
|
|
|
@property
|
|
def last_polled_at(self) -> datetime | None:
|
|
return self._last_polled_at
|
|
|
|
@property
|
|
def is_circuit_open(self) -> bool:
|
|
return self._breaker.is_open
|
|
|
|
@property
|
|
def last_skip_reason(self) -> str | None:
|
|
"""Reason the most recent sweep attempt skipped (e.g. "pending_work",
|
|
"circuit_open"), or None when the most recent attempt actually polled.
|
|
Cleared on the next successful sweep."""
|
|
return self._last_skip_reason
|
|
|
|
async def run(self) -> None: # pragma: no cover - subclasses override
|
|
raise NotImplementedError
|
|
|
|
async def stop(self) -> None:
|
|
self._stop.set()
|
|
if self._task is not None:
|
|
await asyncio.gather(self._task, return_exceptions=True)
|
|
self._task = None
|
|
|
|
async def _sweep_once(self) -> bool:
|
|
"""One discover() sweep. Returns True on success, False if the
|
|
breaker is open, the source has pending work already queued, or the
|
|
sweep failed (and was recorded)."""
|
|
if self._breaker.is_open:
|
|
self._last_skip_reason = "circuit_open"
|
|
logger.debug(
|
|
"Skipping discover() — circuit breaker open for %s", self.source_id
|
|
)
|
|
return False
|
|
with logfire.span("ingester.poller.sweep", source_id=self.source_id) as span:
|
|
if await self._jobs.has_pending(self.source_id):
|
|
# The unique index would dedupe a re-sweep into a saturated
|
|
# queue anyway; skipping saves the listing round-trip
|
|
# (PROPFIND / S3 LIST / FS walk) and keeps Logfire readable.
|
|
self._last_skip_reason = "pending_work"
|
|
span.set_attribute("skipped", True)
|
|
span.set_attribute("skip_reason", "pending_work")
|
|
logger.debug(
|
|
"Skipping discover() — %s has pending work in the queue",
|
|
self.source_id,
|
|
)
|
|
return False
|
|
try:
|
|
snapshot = await self._sync.get_snapshot(self.source_id)
|
|
counts = {
|
|
SourceEventKind.UPSERT: 0,
|
|
SourceEventKind.DELETE: 0,
|
|
SourceEventKind.UNCHANGED: 0,
|
|
}
|
|
async for event in self.source.discover(since=snapshot):
|
|
counts[event.kind] += 1
|
|
await self._handle_event(event)
|
|
self._breaker.record_success()
|
|
self._last_polled_at = datetime.now(UTC)
|
|
self._last_skip_reason = None
|
|
span.set_attribute("upsert", counts[SourceEventKind.UPSERT])
|
|
span.set_attribute("delete", counts[SourceEventKind.DELETE])
|
|
span.set_attribute("unchanged", counts[SourceEventKind.UNCHANGED])
|
|
if counts[SourceEventKind.UPSERT] or counts[SourceEventKind.DELETE]:
|
|
logger.info(
|
|
"Swept %s: %d upsert, %d delete, %d unchanged",
|
|
self.source_id,
|
|
counts[SourceEventKind.UPSERT],
|
|
counts[SourceEventKind.DELETE],
|
|
counts[SourceEventKind.UNCHANGED],
|
|
)
|
|
return True
|
|
except Exception as exc:
|
|
self._breaker.record_failure()
|
|
span.set_attribute(
|
|
"consecutive_failures", self._breaker.consecutive_failures
|
|
)
|
|
span.record_exception(exc)
|
|
logger.exception(
|
|
"discover() failed for %s (consecutive=%d): %s",
|
|
self.source_id,
|
|
self._breaker.consecutive_failures,
|
|
exc,
|
|
)
|
|
return False
|
|
|
|
async def _handle_event(self, event: SourceEvent) -> None:
|
|
if event.kind is SourceEventKind.UPSERT:
|
|
await self._jobs.enqueue(
|
|
event.source_id,
|
|
event.uri,
|
|
op=JobOp.UPSERT,
|
|
revision=event.revision,
|
|
max_attempts=_max_attempts(self.config, self._default_max_attempts),
|
|
extra=_enqueue_extra(self.config),
|
|
)
|
|
# Don't write revision to sync_state here — the worker writes it
|
|
# after a successful ingestion. last_seen_at gets bumped to keep
|
|
# orphan detection accurate.
|
|
await self._sync.upsert(
|
|
event.source_id,
|
|
event.uri,
|
|
revision=None,
|
|
content_hash=None,
|
|
)
|
|
elif event.kind is SourceEventKind.UNCHANGED:
|
|
# Touch last_seen_at without changing the stored revision.
|
|
await self._sync.upsert(
|
|
event.source_id,
|
|
event.uri,
|
|
revision=event.revision,
|
|
content_hash=None,
|
|
)
|
|
elif event.kind is SourceEventKind.DELETE:
|
|
if not self.config.delete_orphans:
|
|
return
|
|
await self._jobs.enqueue(
|
|
event.source_id,
|
|
event.uri,
|
|
op=JobOp.DELETE,
|
|
max_attempts=_max_attempts(self.config, self._default_max_attempts),
|
|
extra=_enqueue_extra(self.config),
|
|
)
|