141 lines
5 KiB
Python
141 lines
5 KiB
Python
import hashlib
|
|
from collections.abc import AsyncIterator
|
|
from datetime import UTC, datetime
|
|
from urllib.parse import urlparse
|
|
|
|
import httpx
|
|
|
|
from haiku.rag.ingester.sources.base import (
|
|
FetchResult,
|
|
RevisionSnapshot,
|
|
SourceEvent,
|
|
SourceEventKind,
|
|
)
|
|
|
|
|
|
def _extract_revision(headers: httpx.Headers) -> tuple[str | None, dict[str, str]]:
|
|
extra: dict[str, str] = {}
|
|
etag = (headers.get("etag") or "").strip('"').strip()
|
|
last_modified = (headers.get("last-modified") or "").strip()
|
|
if etag:
|
|
extra["etag"] = etag
|
|
if last_modified:
|
|
extra["last_modified"] = last_modified
|
|
# Prefer ETag — it's a stronger validator. Fall back to Last-Modified.
|
|
revision = etag or last_modified or None
|
|
return revision, extra
|
|
|
|
|
|
class HTTPSource:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
source_id: str,
|
|
urls: list[str] | None = None,
|
|
headers: dict[str, str] | None = None,
|
|
transport: httpx.AsyncBaseTransport | None = None,
|
|
) -> None:
|
|
self.source_id = source_id
|
|
self.urls = list(urls or [])
|
|
self.headers = dict(headers or {})
|
|
# transport is for testing — production callers leave it None and httpx
|
|
# uses its real transport.
|
|
self._transport = transport
|
|
|
|
def supports(self, uri: str) -> bool:
|
|
return urlparse(uri).scheme in ("http", "https")
|
|
|
|
def _client(self) -> httpx.AsyncClient:
|
|
return httpx.AsyncClient(headers=self.headers, transport=self._transport)
|
|
|
|
async def head(self, uri: str) -> str | None:
|
|
# HTTP doesn't get a cheap revision lookup in v1: the existing
|
|
# ingestion flow always GETs and the dedup uses MD5. A HEAD-first
|
|
# optimization could land later without changing this contract —
|
|
# callers just need to tolerate the extra HEAD.
|
|
return None
|
|
|
|
async def fetch(self, uri: str) -> FetchResult:
|
|
async with self._client() as http:
|
|
response = await http.get(uri)
|
|
response.raise_for_status()
|
|
body = response.content
|
|
content_type = (
|
|
response.headers.get("content-type", "application/octet-stream")
|
|
.split(";")[0]
|
|
.strip()
|
|
.lower()
|
|
)
|
|
revision, extra = _extract_revision(response.headers)
|
|
return FetchResult(
|
|
uri=uri,
|
|
body=body,
|
|
content_type=content_type,
|
|
content_hash=hashlib.md5(body, usedforsecurity=False).hexdigest(),
|
|
revision=revision,
|
|
extra_metadata=extra,
|
|
)
|
|
|
|
async def discover(
|
|
self, since: RevisionSnapshot | None = None
|
|
) -> AsyncIterator[SourceEvent]:
|
|
# HTTP has no listing concept — discover() only reports on what is
|
|
# currently configured in self.urls. Config drift (URIs that were
|
|
# configured before but aren't now) is not visible here; the poller
|
|
# layer detects that by diffing self.urls across sweeps.
|
|
#
|
|
# 410 Gone is the one real source-side deletion signal: the origin
|
|
# explicitly says "permanently gone". 404 and other failures are
|
|
# ambiguous (transient outage, misconfigured URL, auth blip), so we
|
|
# fall back to UPSERT with no revision and let the worker decide
|
|
# via GET.
|
|
snapshot: dict[str, str] = dict(since) if since else {}
|
|
now = datetime.now(UTC)
|
|
|
|
async with self._client() as http:
|
|
for url in self.urls:
|
|
try:
|
|
head = await http.head(url)
|
|
except Exception:
|
|
yield SourceEvent(
|
|
source_id=self.source_id,
|
|
uri=url,
|
|
kind=SourceEventKind.UPSERT,
|
|
revision=None,
|
|
discovered_at=now,
|
|
)
|
|
continue
|
|
|
|
if head.status_code == 410:
|
|
yield SourceEvent(
|
|
source_id=self.source_id,
|
|
uri=url,
|
|
kind=SourceEventKind.DELETE,
|
|
revision=None,
|
|
discovered_at=now,
|
|
)
|
|
continue
|
|
|
|
if head.is_error:
|
|
yield SourceEvent(
|
|
source_id=self.source_id,
|
|
uri=url,
|
|
kind=SourceEventKind.UPSERT,
|
|
revision=None,
|
|
discovered_at=now,
|
|
)
|
|
continue
|
|
|
|
revision, _ = _extract_revision(head.headers)
|
|
if revision is not None and snapshot.get(url) == revision:
|
|
kind = SourceEventKind.UNCHANGED
|
|
else:
|
|
kind = SourceEventKind.UPSERT
|
|
|
|
yield SourceEvent(
|
|
source_id=self.source_id,
|
|
uri=url,
|
|
kind=kind,
|
|
revision=revision,
|
|
discovered_at=now,
|
|
)
|