Merge pull request #452 from ggozad/feat/run-batch-dry-run
Add dry-run manifests for ingestor's run-batch
This commit is contained in:
commit
0be4f7e1bd
13 changed files with 1304 additions and 34 deletions
|
|
@ -1,6 +1,11 @@
|
||||||
# Changelog
|
# Changelog
|
||||||
## [Unreleased]
|
## [Unreleased]
|
||||||
|
|
||||||
|
### Added
|
||||||
|
|
||||||
|
- `haiku-ingester run-batch --dry-run` writes a YAML manifest of planned upserts/deletes without mutating queue jobs or `sync_state`; `run-batch --manifest <path>` replays that frozen changeset without another discovery sweep.
|
||||||
|
- `haiku-ingester run-batch` and `run-batch --manifest` show an interactive progress bar with ETA while draining queued jobs.
|
||||||
|
|
||||||
### Fixed
|
### Fixed
|
||||||
|
|
||||||
- CPU-bound ingest steps now run off the event loop: Docling document serialization, docling-serve zip parsing, split-PDF concatenation, fetched-body temp writes, and filesystem read/hash work no longer stall concurrent ingester workers on large image-bearing documents.
|
- CPU-bound ingest steps now run off the event loop: Docling document serialization, docling-serve zip parsing, split-PDF concatenation, fetched-body temp writes, and filesystem read/hash work no longer stall concurrent ingester workers on large image-bearing documents.
|
||||||
|
|
|
||||||
|
|
@ -524,6 +524,31 @@ haiku-ingester run-batch
|
||||||
haiku-ingester run-batch --db rag.lancedb
|
haiku-ingester run-batch --db rag.lancedb
|
||||||
```
|
```
|
||||||
|
|
||||||
|
To review a batch before it mutates the document store, use `--dry-run`.
|
||||||
|
Dry-run performs the same discovery checks but writes no queue jobs and does
|
||||||
|
not update `sync_state`. It writes a YAML manifest named
|
||||||
|
`manifest-<datestamp>.yaml` by default:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
haiku-ingester run-batch --dry-run
|
||||||
|
haiku-ingester run-batch --dry-run --output manifest-20260622.yaml
|
||||||
|
```
|
||||||
|
|
||||||
|
The manifest records the `upsert` and `delete` changes discovered for each
|
||||||
|
source. Replay it later to ingest exactly that changeset, without another
|
||||||
|
discovery sweep:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
haiku-ingester run-batch --manifest manifest-20260622.yaml
|
||||||
|
```
|
||||||
|
|
||||||
|
Manifest replay rejects sources with queued or claimed work, preserving the
|
||||||
|
one-active-changeset-per-source pattern. Revisioned upserts are checked
|
||||||
|
against the current upstream revision before fetch; if the resource changed
|
||||||
|
after dry-run, that job dead-letters and the newer version waits for the next
|
||||||
|
dry-run. Sources that provide no revision can freeze URI discovery but cannot
|
||||||
|
prove byte identity at replay time.
|
||||||
|
|
||||||
Orphan deletion compares each source against `sync_state` in the queue DB,
|
Orphan deletion compares each source against `sync_state` in the queue DB,
|
||||||
so persist `ingester.db` between runs for deletions to be detected. It exits
|
so persist `ingester.db` between runs for deletions to be detected. It exits
|
||||||
non-zero if any job dead-letters or a source's discovery sweep does not
|
non-zero if any job dead-letters or a source's discovery sweep does not
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,7 @@
|
||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
import signal
|
import signal
|
||||||
|
from collections.abc import Callable
|
||||||
from contextlib import asynccontextmanager
|
from contextlib import asynccontextmanager
|
||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
@ -9,9 +10,11 @@ from typing import TYPE_CHECKING
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
from haiku.rag.config import AppConfig
|
from haiku.rag.config import AppConfig
|
||||||
|
from haiku.rag.ingester.batch import BatchChange, BatchDryRunReport, BatchManifest
|
||||||
from haiku.rag.ingester.metadata import build_providers, load_metadata_providers
|
from haiku.rag.ingester.metadata import build_providers, load_metadata_providers
|
||||||
from haiku.rag.ingester.pollers.manager import PollerManager
|
from haiku.rag.ingester.pollers.manager import PollerManager
|
||||||
from haiku.rag.ingester.queue.migrations import open_queue
|
from haiku.rag.ingester.queue.migrations import open_queue
|
||||||
|
from haiku.rag.ingester.queue.models import Job, JobStatus
|
||||||
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
|
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
|
||||||
from haiku.rag.ingester.workers.pool import WorkerPool
|
from haiku.rag.ingester.workers.pool import WorkerPool
|
||||||
from haiku.rag.ingester.workers.retry import RetryPolicy
|
from haiku.rag.ingester.workers.retry import RetryPolicy
|
||||||
|
|
@ -21,6 +24,8 @@ if TYPE_CHECKING:
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
_MANIFEST_EXTRA_KEY = "_manifest"
|
||||||
|
|
||||||
|
|
||||||
def _api_access_log_enabled() -> bool:
|
def _api_access_log_enabled() -> bool:
|
||||||
"""Per-request access logging only when the haiku.rag logger is at DEBUG.
|
"""Per-request access logging only when the haiku.rag logger is at DEBUG.
|
||||||
|
|
@ -38,6 +43,28 @@ class BatchReport(BaseModel):
|
||||||
failed_sweeps: list[str] = []
|
failed_sweeps: list[str] = []
|
||||||
|
|
||||||
|
|
||||||
|
class BatchProgress(BaseModel):
|
||||||
|
"""Snapshot emitted while a one-shot batch drains queued work."""
|
||||||
|
|
||||||
|
total: int = 0
|
||||||
|
completed: int = 0
|
||||||
|
succeeded: int = 0
|
||||||
|
dead: int = 0
|
||||||
|
queued: int = 0
|
||||||
|
claimed: int = 0
|
||||||
|
|
||||||
|
|
||||||
|
BatchProgressCallback = Callable[[BatchProgress], None]
|
||||||
|
|
||||||
|
|
||||||
|
def _manifest_change_key(change: BatchChange) -> tuple[str, str, str, str | None]:
|
||||||
|
return (change.source_id, change.uri, change.op.value, change.revision)
|
||||||
|
|
||||||
|
|
||||||
|
def _manifest_job_key(job: Job) -> tuple[str, str, str, str | None]:
|
||||||
|
return (job.source_id, job.uri, job.op.value, job.revision)
|
||||||
|
|
||||||
|
|
||||||
class IngesterApp:
|
class IngesterApp:
|
||||||
"""Top-level lifecycle for the production ingester.
|
"""Top-level lifecycle for the production ingester.
|
||||||
|
|
||||||
|
|
@ -129,6 +156,35 @@ class IngesterApp:
|
||||||
await self._engine.dispose()
|
await self._engine.dispose()
|
||||||
self._engine = None
|
self._engine = None
|
||||||
|
|
||||||
|
@asynccontextmanager
|
||||||
|
async def _discovery_resources(self):
|
||||||
|
"""Open only the queue and source pollers needed for discovery.
|
||||||
|
Dry-runs must not create/open the LanceDB document store or worker
|
||||||
|
pool because they are upstream checks only."""
|
||||||
|
from haiku.rag.converters import get_converter
|
||||||
|
|
||||||
|
ingester_cfg = self._config.ingester
|
||||||
|
self._engine = await open_queue(ingester_cfg.queue)
|
||||||
|
try:
|
||||||
|
self._jobs = JobRepo(self._engine)
|
||||||
|
self._sync = SyncStateRepo(self._engine)
|
||||||
|
supported_extensions = get_converter(self._config).supported_extensions
|
||||||
|
self._pollers = PollerManager(
|
||||||
|
configs=ingester_cfg.sources,
|
||||||
|
job_repo=self._jobs,
|
||||||
|
sync_repo=self._sync,
|
||||||
|
supported_extensions=supported_extensions,
|
||||||
|
default_max_attempts=ingester_cfg.workers.retry.max_attempts,
|
||||||
|
)
|
||||||
|
yield
|
||||||
|
finally:
|
||||||
|
if self._pollers is not None:
|
||||||
|
await self._pollers.close_sources()
|
||||||
|
self._pollers = None
|
||||||
|
if self._engine is not None:
|
||||||
|
await self._engine.dispose()
|
||||||
|
self._engine = None
|
||||||
|
|
||||||
async def _stop_pool(self) -> None:
|
async def _stop_pool(self) -> None:
|
||||||
"""Stop the worker pool, honouring the shutdown grace, then drain any
|
"""Stop the worker pool, honouring the shutdown grace, then drain any
|
||||||
cancel-cleanup release tasks before the queue connection closes."""
|
cancel-cleanup release tasks before the queue connection closes."""
|
||||||
|
|
@ -149,6 +205,51 @@ class IngesterApp:
|
||||||
if landed:
|
if landed:
|
||||||
logger.info("Drained %d cancel-cleanup release(s) before close", landed)
|
logger.info("Drained %d cancel-cleanup release(s) before close", landed)
|
||||||
|
|
||||||
|
async def _drain_batch(
|
||||||
|
self,
|
||||||
|
started_at: datetime,
|
||||||
|
*,
|
||||||
|
progress_callback: BatchProgressCallback | None = None,
|
||||||
|
) -> BatchReport:
|
||||||
|
assert self._pool is not None and self._jobs is not None
|
||||||
|
total = 0
|
||||||
|
while True:
|
||||||
|
counts = await self._jobs.batch_progress_counts_since(started_at)
|
||||||
|
queued = counts.get("queued", 0)
|
||||||
|
claimed = counts.get("claimed", 0)
|
||||||
|
outstanding = queued + claimed
|
||||||
|
succeeded = counts.get("succeeded", 0)
|
||||||
|
dead = counts.get("dead", 0)
|
||||||
|
completed_count = succeeded + dead
|
||||||
|
total = max(total, outstanding + completed_count)
|
||||||
|
if progress_callback is not None:
|
||||||
|
progress_callback(
|
||||||
|
BatchProgress(
|
||||||
|
total=total,
|
||||||
|
completed=min(completed_count, total),
|
||||||
|
succeeded=succeeded,
|
||||||
|
dead=dead,
|
||||||
|
queued=queued,
|
||||||
|
claimed=claimed,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
if not outstanding:
|
||||||
|
break
|
||||||
|
if self._pool.live_workers == 0:
|
||||||
|
logger.error(
|
||||||
|
"All workers have died with %d outstanding job(s) "
|
||||||
|
"— aborting batch; stranded jobs will be reaped "
|
||||||
|
"on next start",
|
||||||
|
outstanding,
|
||||||
|
)
|
||||||
|
break
|
||||||
|
await asyncio.sleep(0.1)
|
||||||
|
completed = await self._jobs.counts_by_status_since(started_at)
|
||||||
|
return BatchReport(
|
||||||
|
succeeded=completed.get("succeeded", 0),
|
||||||
|
dead=completed.get("dead", 0),
|
||||||
|
)
|
||||||
|
|
||||||
async def serve(self, *, api: bool = True) -> None:
|
async def serve(self, *, api: bool = True) -> None:
|
||||||
"""Run pollers + workers (and the HTTP API when enabled) until a
|
"""Run pollers + workers (and the HTTP API when enabled) until a
|
||||||
SIGINT/SIGTERM is received. Drains in-flight work on shutdown."""
|
SIGINT/SIGTERM is received. Drains in-flight work on shutdown."""
|
||||||
|
|
@ -164,7 +265,7 @@ class IngesterApp:
|
||||||
uses_docling_serve = (
|
uses_docling_serve = (
|
||||||
proc.converter == "docling-serve" or proc.chunker == "docling-serve"
|
proc.converter == "docling-serve" or proc.chunker == "docling-serve"
|
||||||
)
|
)
|
||||||
if uses_docling_serve:
|
if uses_docling_serve: # pragma: no cover
|
||||||
logger.info(
|
logger.info(
|
||||||
"Ingester running: %d worker(s), %d source(s), "
|
"Ingester running: %d worker(s), %d source(s), "
|
||||||
"%d docling-serve instance(s)",
|
"%d docling-serve instance(s)",
|
||||||
|
|
@ -202,7 +303,9 @@ class IngesterApp:
|
||||||
await self._stop_pool()
|
await self._stop_pool()
|
||||||
await self._pollers.close_sources()
|
await self._pollers.close_sources()
|
||||||
|
|
||||||
async def run_batch(self) -> BatchReport:
|
async def run_batch(
|
||||||
|
self, *, progress_callback: BatchProgressCallback | None = None
|
||||||
|
) -> BatchReport:
|
||||||
"""Run one discover() sweep across every configured source, drain the
|
"""Run one discover() sweep across every configured source, drain the
|
||||||
queue to completion, then stop. Unlike `serve`, the periodic poller
|
queue to completion, then stop. Unlike `serve`, the periodic poller
|
||||||
loops never start — discovery is driven explicitly, so the run is
|
loops never start — discovery is driven explicitly, so the run is
|
||||||
|
|
@ -220,25 +323,128 @@ class IngesterApp:
|
||||||
await self._pool.start()
|
await self._pool.start()
|
||||||
try:
|
try:
|
||||||
failed_sweeps = await self._pollers.sweep_all()
|
failed_sweeps = await self._pollers.sweep_all()
|
||||||
while True:
|
report = await self._drain_batch(
|
||||||
counts = await self._jobs.counts_by_status()
|
started_at, progress_callback=progress_callback
|
||||||
if not counts.get("queued") and not counts.get("claimed"):
|
)
|
||||||
break
|
report.failed_sweeps = failed_sweeps
|
||||||
if self._pool.live_workers == 0:
|
return report
|
||||||
outstanding = counts.get("queued", 0) + counts.get("claimed", 0)
|
finally:
|
||||||
logger.error(
|
await self._stop_pool()
|
||||||
"All workers have died with %d outstanding job(s) "
|
await self._pollers.close_sources()
|
||||||
"— aborting batch; stranded jobs will be reaped "
|
|
||||||
"on next start",
|
async def run_batch_dry_run(self) -> BatchDryRunReport:
|
||||||
outstanding,
|
"""Run one discover() sweep across every configured source and return
|
||||||
)
|
the jobs that would be enqueued, without mutating jobs or sync_state."""
|
||||||
break
|
async with self._discovery_resources():
|
||||||
await asyncio.sleep(0.1)
|
assert self._pollers is not None
|
||||||
completed = await self._jobs.counts_by_status_since(started_at)
|
manifest, failed_sweeps = await self._pollers.dry_run_manifest()
|
||||||
return BatchReport(
|
return BatchDryRunReport(manifest=manifest, failed_sweeps=failed_sweeps)
|
||||||
succeeded=completed.get("succeeded", 0),
|
|
||||||
dead=completed.get("dead", 0),
|
async def run_batch_from_manifest(
|
||||||
failed_sweeps=failed_sweeps,
|
self,
|
||||||
|
manifest: BatchManifest,
|
||||||
|
*,
|
||||||
|
progress_callback: BatchProgressCallback | None = None,
|
||||||
|
) -> BatchReport:
|
||||||
|
"""Enqueue and drain a dry-run manifest without running a fresh
|
||||||
|
discovery sweep."""
|
||||||
|
if manifest.version != 1: # pragma: no cover
|
||||||
|
raise ValueError(f"Unsupported manifest version: {manifest.version}")
|
||||||
|
async with self._resources():
|
||||||
|
assert (
|
||||||
|
self._pollers is not None
|
||||||
|
and self._pool is not None
|
||||||
|
and self._jobs is not None
|
||||||
|
)
|
||||||
|
configured = {source.source_id for source in self._pollers.sources}
|
||||||
|
manifest_sources = {change.source_id for change in manifest.changes}
|
||||||
|
missing = sorted(manifest_sources - configured)
|
||||||
|
if missing: # pragma: no cover
|
||||||
|
await self._pollers.close_sources()
|
||||||
|
raise ValueError(
|
||||||
|
"Manifest references unconfigured source(s): " + ", ".join(missing)
|
||||||
|
)
|
||||||
|
|
||||||
|
seen: set[tuple[str, str]] = set()
|
||||||
|
duplicates: set[tuple[str, str]] = set()
|
||||||
|
for change in manifest.changes:
|
||||||
|
key = (change.source_id, change.uri)
|
||||||
|
if key in seen:
|
||||||
|
duplicates.add(key)
|
||||||
|
seen.add(key)
|
||||||
|
if duplicates:
|
||||||
|
await self._pollers.close_sources()
|
||||||
|
rendered = ", ".join(
|
||||||
|
f"{source_id}:{uri}" for source_id, uri in duplicates
|
||||||
|
)
|
||||||
|
raise ValueError(f"Manifest contains duplicate change(s): {rendered}")
|
||||||
|
|
||||||
|
manifest_key = manifest.generated_at.isoformat()
|
||||||
|
manifest_change_keys = {
|
||||||
|
_manifest_change_key(change) for change in manifest.changes
|
||||||
|
}
|
||||||
|
live_jobs = [
|
||||||
|
*await self._jobs.list_jobs(status=JobStatus.QUEUED, limit=10_000),
|
||||||
|
*await self._jobs.list_jobs(status=JobStatus.CLAIMED, limit=10_000),
|
||||||
|
]
|
||||||
|
stale_jobs: list[Job] = []
|
||||||
|
live_manifest_keys: set[tuple[str, str, str, str | None]] = set()
|
||||||
|
for job in live_jobs:
|
||||||
|
extra = job.extra or {}
|
||||||
|
job_manifest = extra.get(_MANIFEST_EXTRA_KEY) or {}
|
||||||
|
key = _manifest_job_key(job)
|
||||||
|
if (
|
||||||
|
job_manifest.get("generated_at") != manifest_key
|
||||||
|
or key not in manifest_change_keys
|
||||||
|
):
|
||||||
|
stale_jobs.append(job)
|
||||||
|
continue
|
||||||
|
live_manifest_keys.add(key)
|
||||||
|
if stale_jobs:
|
||||||
|
await self._pollers.close_sources()
|
||||||
|
raise ValueError(
|
||||||
|
"Cannot replay manifest while the queue has non-manifest "
|
||||||
|
f"pending work: {len(stale_jobs)} queued/claimed job(s)"
|
||||||
|
)
|
||||||
|
|
||||||
|
default_max_attempts = self._config.ingester.workers.retry.max_attempts
|
||||||
|
max_attempts_by_source = {
|
||||||
|
poller.source_id: (
|
||||||
|
poller.config.retry.max_attempts
|
||||||
|
if poller.config.retry is not None
|
||||||
|
else default_max_attempts
|
||||||
|
)
|
||||||
|
for poller in self._pollers.pollers
|
||||||
|
}
|
||||||
|
for change in manifest.changes:
|
||||||
|
if _manifest_change_key(change) in live_manifest_keys:
|
||||||
|
continue
|
||||||
|
job = await self._jobs.enqueue(
|
||||||
|
change.source_id,
|
||||||
|
change.uri,
|
||||||
|
op=change.op,
|
||||||
|
revision=change.revision,
|
||||||
|
max_attempts=max_attempts_by_source[change.source_id],
|
||||||
|
extra={
|
||||||
|
_MANIFEST_EXTRA_KEY: {
|
||||||
|
"version": manifest.version,
|
||||||
|
"generated_at": manifest_key,
|
||||||
|
"discovered_at": change.discovered_at.isoformat(),
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
if job is None: # pragma: no cover
|
||||||
|
await self._pollers.close_sources()
|
||||||
|
raise ValueError(
|
||||||
|
"Cannot replay manifest because a live job already exists "
|
||||||
|
f"for {change.source_id}:{change.uri}"
|
||||||
|
)
|
||||||
|
|
||||||
|
started_at = datetime.now(UTC)
|
||||||
|
await self._pool.start()
|
||||||
|
try:
|
||||||
|
return await self._drain_batch(
|
||||||
|
started_at, progress_callback=progress_callback
|
||||||
)
|
)
|
||||||
finally:
|
finally:
|
||||||
await self._stop_pool()
|
await self._stop_pool()
|
||||||
|
|
|
||||||
33
haiku_rag_slim/haiku/rag/ingester/batch.py
Normal file
33
haiku_rag_slim/haiku/rag/ingester/batch.py
Normal file
|
|
@ -0,0 +1,33 @@
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
from haiku.rag.ingester.queue.models import JobOp
|
||||||
|
|
||||||
|
|
||||||
|
class BatchChange(BaseModel):
|
||||||
|
op: JobOp
|
||||||
|
source_id: str
|
||||||
|
uri: str
|
||||||
|
revision: str | None = None
|
||||||
|
discovered_at: datetime
|
||||||
|
|
||||||
|
|
||||||
|
class BatchSourceSummary(BaseModel):
|
||||||
|
source_id: str
|
||||||
|
upsert_count: int = 0
|
||||||
|
delete_count: int = 0
|
||||||
|
unchanged_count: int = 0
|
||||||
|
ignored_delete_count: int = 0
|
||||||
|
|
||||||
|
|
||||||
|
class BatchManifest(BaseModel):
|
||||||
|
version: int = 1
|
||||||
|
generated_at: datetime
|
||||||
|
sources: list[BatchSourceSummary] = []
|
||||||
|
changes: list[BatchChange] = []
|
||||||
|
|
||||||
|
|
||||||
|
class BatchDryRunReport(BaseModel):
|
||||||
|
manifest: BatchManifest
|
||||||
|
failed_sweeps: list[str] = []
|
||||||
|
|
@ -1,9 +1,21 @@
|
||||||
import asyncio
|
import asyncio
|
||||||
import sys
|
import sys
|
||||||
|
from collections.abc import Iterator
|
||||||
|
from contextlib import contextmanager
|
||||||
|
from datetime import UTC, datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import typer
|
import typer
|
||||||
|
import yaml
|
||||||
from dotenv import find_dotenv, load_dotenv
|
from dotenv import find_dotenv, load_dotenv
|
||||||
|
from rich.console import Console
|
||||||
|
from rich.progress import (
|
||||||
|
BarColumn,
|
||||||
|
Progress,
|
||||||
|
TextColumn,
|
||||||
|
TimeElapsedColumn,
|
||||||
|
TimeRemainingColumn,
|
||||||
|
)
|
||||||
from sqlalchemy import make_url
|
from sqlalchemy import make_url
|
||||||
|
|
||||||
load_dotenv(find_dotenv(usecwd=True))
|
load_dotenv(find_dotenv(usecwd=True))
|
||||||
|
|
@ -16,7 +28,12 @@ from haiku.rag.config import ( # noqa: E402
|
||||||
load_yaml_config,
|
load_yaml_config,
|
||||||
set_config,
|
set_config,
|
||||||
)
|
)
|
||||||
from haiku.rag.ingester.app import IngesterApp # noqa: E402
|
from haiku.rag.ingester.app import ( # noqa: E402
|
||||||
|
BatchProgress,
|
||||||
|
BatchProgressCallback,
|
||||||
|
IngesterApp,
|
||||||
|
)
|
||||||
|
from haiku.rag.ingester.batch import BatchManifest # noqa: E402
|
||||||
from haiku.rag.ingester.queue.migrations import open_queue # noqa: E402
|
from haiku.rag.ingester.queue.migrations import open_queue # noqa: E402
|
||||||
from haiku.rag.logging import configure_cli_logging # noqa: E402
|
from haiku.rag.logging import configure_cli_logging # noqa: E402
|
||||||
from haiku.rag.store.exceptions import ( # noqa: E402
|
from haiku.rag.store.exceptions import ( # noqa: E402
|
||||||
|
|
@ -137,6 +154,64 @@ def _resolve_db_path(config: AppConfig, override: Path | None) -> Path:
|
||||||
return override or (config.storage.data_dir / "haiku.rag.lancedb")
|
return override or (config.storage.data_dir / "haiku.rag.lancedb")
|
||||||
|
|
||||||
|
|
||||||
|
def _default_manifest_path() -> Path:
|
||||||
|
datestamp = datetime.now(UTC).strftime("%Y%m%d-%H%M%SZ")
|
||||||
|
return Path(f"manifest-{datestamp}.yaml")
|
||||||
|
|
||||||
|
|
||||||
|
def _write_manifest(manifest: BatchManifest, path: Path) -> None:
|
||||||
|
data = manifest.model_dump(mode="json")
|
||||||
|
path.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8")
|
||||||
|
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def _batch_progress(
|
||||||
|
description: str,
|
||||||
|
) -> Iterator[BatchProgressCallback | None]: # pragma: no cover
|
||||||
|
console = Console(file=sys.stdout)
|
||||||
|
if not console.is_terminal:
|
||||||
|
yield None
|
||||||
|
return
|
||||||
|
|
||||||
|
progress = Progress(
|
||||||
|
TextColumn("[progress.description]{task.description}"),
|
||||||
|
BarColumn(),
|
||||||
|
TextColumn("{task.completed}/{task.total}"),
|
||||||
|
TimeRemainingColumn(),
|
||||||
|
TimeElapsedColumn(),
|
||||||
|
console=console,
|
||||||
|
transient=True,
|
||||||
|
)
|
||||||
|
task_id = None
|
||||||
|
|
||||||
|
def _update(snapshot: BatchProgress) -> None:
|
||||||
|
nonlocal task_id
|
||||||
|
task_description = (
|
||||||
|
f"{description} ({snapshot.succeeded} ok, {snapshot.dead} dead)"
|
||||||
|
)
|
||||||
|
if task_id is None:
|
||||||
|
task_id = progress.add_task(
|
||||||
|
task_description,
|
||||||
|
total=snapshot.total,
|
||||||
|
completed=snapshot.completed,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
progress.update(
|
||||||
|
task_id,
|
||||||
|
description=task_description,
|
||||||
|
total=snapshot.total,
|
||||||
|
completed=snapshot.completed,
|
||||||
|
)
|
||||||
|
|
||||||
|
with progress:
|
||||||
|
yield _update
|
||||||
|
|
||||||
|
|
||||||
|
def _load_manifest(path: Path) -> BatchManifest:
|
||||||
|
data = yaml.safe_load(path.read_text(encoding="utf-8"))
|
||||||
|
return BatchManifest.model_validate(data)
|
||||||
|
|
||||||
|
|
||||||
@_cli.command("serve")
|
@_cli.command("serve")
|
||||||
def serve(
|
def serve(
|
||||||
db: Path | None = typer.Option(
|
db: Path | None = typer.Option(
|
||||||
|
|
@ -188,18 +263,98 @@ def run_batch(
|
||||||
"--db",
|
"--db",
|
||||||
help="LanceDB path (overrides config.storage.data_dir).",
|
help="LanceDB path (overrides config.storage.data_dir).",
|
||||||
),
|
),
|
||||||
|
dry_run: bool = typer.Option(
|
||||||
|
False,
|
||||||
|
"--dry-run",
|
||||||
|
help="Discover planned changes and write a YAML manifest without ingesting.",
|
||||||
|
),
|
||||||
|
output: Path | None = typer.Option(
|
||||||
|
None,
|
||||||
|
"--output",
|
||||||
|
"-o",
|
||||||
|
help="Dry-run manifest path (defaults to manifest-<datestamp>.yaml).",
|
||||||
|
),
|
||||||
|
manifest: Path | None = typer.Option(
|
||||||
|
None,
|
||||||
|
"--manifest",
|
||||||
|
help="Replay a dry-run manifest instead of running discovery.",
|
||||||
|
),
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Run one discover sweep across every configured source, drain the queue,
|
"""Run one discover sweep across every configured source, drain the queue,
|
||||||
then exit. New and changed resources are ingested, resources that vanished
|
then exit. New and changed resources are ingested, resources that vanished
|
||||||
from a source are deleted. Exits non-zero if any job dead-letters or a
|
from a source are deleted. Exits non-zero if any job dead-letters or a
|
||||||
source's sweep does not complete."""
|
source's sweep does not complete."""
|
||||||
asyncio.run(_run_batch(get_config(), db))
|
if manifest is not None and dry_run:
|
||||||
|
typer.echo("Error: --manifest cannot be combined with --dry-run")
|
||||||
|
raise typer.Exit(2)
|
||||||
|
if output is not None and not dry_run:
|
||||||
|
typer.echo("Error: --output is only valid with --dry-run")
|
||||||
|
raise typer.Exit(2)
|
||||||
|
asyncio.run(
|
||||||
|
_run_batch(
|
||||||
|
get_config(),
|
||||||
|
db,
|
||||||
|
dry_run=dry_run,
|
||||||
|
output=output,
|
||||||
|
manifest_path=manifest,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
async def _run_batch(app_config: AppConfig, db_path: Path | None) -> None:
|
async def _run_batch(
|
||||||
|
app_config: AppConfig,
|
||||||
|
db_path: Path | None,
|
||||||
|
*,
|
||||||
|
dry_run: bool = False,
|
||||||
|
output: Path | None = None,
|
||||||
|
manifest_path: Path | None = None,
|
||||||
|
) -> None:
|
||||||
db = _resolve_db_path(app_config, db_path)
|
db = _resolve_db_path(app_config, db_path)
|
||||||
app = IngesterApp(config=app_config, db_path=db)
|
app = IngesterApp(config=app_config, db_path=db)
|
||||||
report = await app.run_batch()
|
if dry_run:
|
||||||
|
report = await app.run_batch_dry_run()
|
||||||
|
if report.failed_sweeps:
|
||||||
|
typer.echo(
|
||||||
|
f"Sources that failed to sweep: {', '.join(report.failed_sweeps)}"
|
||||||
|
)
|
||||||
|
raise typer.Exit(1)
|
||||||
|
manifest_path = output or _default_manifest_path()
|
||||||
|
_write_manifest(report.manifest, manifest_path)
|
||||||
|
upserts = sum(source.upsert_count for source in report.manifest.sources)
|
||||||
|
deletes = sum(source.delete_count for source in report.manifest.sources)
|
||||||
|
unchanged = sum(source.unchanged_count for source in report.manifest.sources)
|
||||||
|
typer.echo(
|
||||||
|
"Dry run complete: "
|
||||||
|
f"{upserts} upsert, {deletes} delete, {unchanged} unchanged "
|
||||||
|
f"-> {manifest_path}"
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
if manifest_path is not None:
|
||||||
|
try:
|
||||||
|
manifest = _load_manifest(manifest_path)
|
||||||
|
with _batch_progress("Replaying manifest") as progress_callback:
|
||||||
|
if progress_callback is None:
|
||||||
|
report = await app.run_batch_from_manifest(manifest)
|
||||||
|
else:
|
||||||
|
report = await app.run_batch_from_manifest(
|
||||||
|
manifest, progress_callback=progress_callback
|
||||||
|
)
|
||||||
|
except ValueError as exc:
|
||||||
|
typer.echo(f"Error: {exc}")
|
||||||
|
raise typer.Exit(1) from exc
|
||||||
|
typer.echo(
|
||||||
|
f"Manifest batch complete: {report.succeeded} succeeded, {report.dead} dead"
|
||||||
|
)
|
||||||
|
if report.dead:
|
||||||
|
raise typer.Exit(1)
|
||||||
|
return
|
||||||
|
|
||||||
|
with _batch_progress("Running batch") as progress_callback:
|
||||||
|
if progress_callback is None:
|
||||||
|
report = await app.run_batch()
|
||||||
|
else:
|
||||||
|
report = await app.run_batch(progress_callback=progress_callback)
|
||||||
typer.echo(f"Batch complete: {report.succeeded} succeeded, {report.dead} dead")
|
typer.echo(f"Batch complete: {report.succeeded} succeeded, {report.dead} dead")
|
||||||
if report.failed_sweeps:
|
if report.failed_sweeps:
|
||||||
typer.echo(
|
typer.echo(
|
||||||
|
|
|
||||||
|
|
@ -4,6 +4,7 @@ import random
|
||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
|
|
||||||
from haiku.rag.config import SourceConfig
|
from haiku.rag.config import SourceConfig
|
||||||
|
from haiku.rag.ingester.batch import BatchChange, BatchSourceSummary
|
||||||
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
|
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
|
||||||
from haiku.rag.ingester.queue.models import JobOp, SyncRow
|
from haiku.rag.ingester.queue.models import JobOp, SyncRow
|
||||||
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
|
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
|
||||||
|
|
@ -165,6 +166,82 @@ class BasePoller:
|
||||||
)
|
)
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
async def _dry_run_once(self) -> tuple[bool, BatchSourceSummary, list[BatchChange]]:
|
||||||
|
"""Collect what one discover() sweep would enqueue without writing
|
||||||
|
jobs or sync_state."""
|
||||||
|
summary = BatchSourceSummary(source_id=self.source_id)
|
||||||
|
changes: list[BatchChange] = []
|
||||||
|
if self._breaker.is_open:
|
||||||
|
self._last_skip_reason = "circuit_open"
|
||||||
|
logger.debug(
|
||||||
|
"Skipping dry-run discover() — circuit breaker open for %s",
|
||||||
|
self.source_id,
|
||||||
|
)
|
||||||
|
return False, summary, changes
|
||||||
|
with logfire.span("ingester.poller.dry_run", source_id=self.source_id) as span:
|
||||||
|
if await self._jobs.has_pending(self.source_id):
|
||||||
|
self._last_skip_reason = "pending_work"
|
||||||
|
span.set_attribute("skipped", True)
|
||||||
|
span.set_attribute("skip_reason", "pending_work")
|
||||||
|
logger.debug(
|
||||||
|
"Skipping dry-run discover() — %s has pending work in the queue",
|
||||||
|
self.source_id,
|
||||||
|
)
|
||||||
|
return False, summary, changes
|
||||||
|
try:
|
||||||
|
revisions = await self._sync.get_revision_snapshot(self.source_id)
|
||||||
|
known = await self._sync.list_known_uris(self.source_id)
|
||||||
|
async for event in self.source.discover(
|
||||||
|
since=revisions, known_uris=known
|
||||||
|
):
|
||||||
|
if event.kind is SourceEventKind.UPSERT:
|
||||||
|
summary.upsert_count += 1
|
||||||
|
changes.append(
|
||||||
|
BatchChange(
|
||||||
|
op=JobOp.UPSERT,
|
||||||
|
source_id=event.source_id,
|
||||||
|
uri=event.uri,
|
||||||
|
revision=event.revision,
|
||||||
|
discovered_at=event.discovered_at,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
elif event.kind is SourceEventKind.UNCHANGED:
|
||||||
|
summary.unchanged_count += 1
|
||||||
|
elif event.kind is SourceEventKind.DELETE:
|
||||||
|
if self.config.delete_orphans:
|
||||||
|
summary.delete_count += 1
|
||||||
|
changes.append(
|
||||||
|
BatchChange(
|
||||||
|
op=JobOp.DELETE,
|
||||||
|
source_id=event.source_id,
|
||||||
|
uri=event.uri,
|
||||||
|
revision=None,
|
||||||
|
discovered_at=event.discovered_at,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
summary.ignored_delete_count += 1
|
||||||
|
self._breaker.record_success()
|
||||||
|
self._last_polled_at = datetime.now(UTC)
|
||||||
|
self._last_skip_reason = None
|
||||||
|
span.set_attribute("upsert", summary.upsert_count)
|
||||||
|
span.set_attribute("delete", summary.delete_count)
|
||||||
|
span.set_attribute("unchanged", summary.unchanged_count)
|
||||||
|
return True, summary, changes
|
||||||
|
except Exception as exc:
|
||||||
|
self._breaker.record_failure()
|
||||||
|
span.set_attribute(
|
||||||
|
"consecutive_failures", self._breaker.consecutive_failures
|
||||||
|
)
|
||||||
|
span.record_exception(exc)
|
||||||
|
logger.exception(
|
||||||
|
"dry-run discover() failed for %s (consecutive=%d): %s",
|
||||||
|
self.source_id,
|
||||||
|
self._breaker.consecutive_failures,
|
||||||
|
exc,
|
||||||
|
)
|
||||||
|
return False, summary, changes
|
||||||
|
|
||||||
async def _handle_event(
|
async def _handle_event(
|
||||||
self,
|
self,
|
||||||
event: SourceEvent,
|
event: SourceEvent,
|
||||||
|
|
|
||||||
|
|
@ -1,9 +1,11 @@
|
||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
from collections.abc import Sequence
|
from collections.abc import Sequence
|
||||||
|
from datetime import UTC, datetime
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from haiku.rag.config import FSSourceConfig, SourceConfig
|
from haiku.rag.config import FSSourceConfig, SourceConfig
|
||||||
|
from haiku.rag.ingester.batch import BatchManifest
|
||||||
from haiku.rag.ingester.pollers.base import BasePoller
|
from haiku.rag.ingester.pollers.base import BasePoller
|
||||||
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
|
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
|
||||||
from haiku.rag.ingester.pollers.factory import build_source
|
from haiku.rag.ingester.pollers.factory import build_source
|
||||||
|
|
@ -90,6 +92,27 @@ class PollerManager:
|
||||||
failed.append(poller.source_id)
|
failed.append(poller.source_id)
|
||||||
return failed
|
return failed
|
||||||
|
|
||||||
|
async def dry_run_manifest(self) -> tuple[BatchManifest, list[str]]:
|
||||||
|
"""Collect what one sweep across every source would enqueue without
|
||||||
|
mutating queue jobs or sync_state."""
|
||||||
|
failed: list[str] = []
|
||||||
|
summaries = []
|
||||||
|
changes = []
|
||||||
|
for poller in self._pollers:
|
||||||
|
ok, summary, source_changes = await poller._dry_run_once()
|
||||||
|
summaries.append(summary)
|
||||||
|
changes.extend(source_changes)
|
||||||
|
if not ok:
|
||||||
|
failed.append(poller.source_id)
|
||||||
|
return (
|
||||||
|
BatchManifest(
|
||||||
|
generated_at=datetime.now(UTC),
|
||||||
|
sources=summaries,
|
||||||
|
changes=changes,
|
||||||
|
),
|
||||||
|
failed,
|
||||||
|
)
|
||||||
|
|
||||||
async def stop(self) -> None:
|
async def stop(self) -> None:
|
||||||
for poller in self._pollers:
|
for poller in self._pollers:
|
||||||
await poller.stop()
|
await poller.stop()
|
||||||
|
|
|
||||||
|
|
@ -371,6 +371,27 @@ class JobRepo:
|
||||||
rows = (await conn.execute(query)).all()
|
rows = (await conn.execute(query)).all()
|
||||||
return {status: n for status, n in rows}
|
return {status: n for status, n in rows}
|
||||||
|
|
||||||
|
async def batch_progress_counts_since(self, since: datetime) -> dict[str, int]:
|
||||||
|
"""Counts for a one-shot batch progress snapshot.
|
||||||
|
|
||||||
|
Live rows are counted regardless of enqueue time because run-batch
|
||||||
|
drains the whole pending queue. Terminal rows are counted only when
|
||||||
|
this run completed them, matching BatchReport semantics.
|
||||||
|
"""
|
||||||
|
query = (
|
||||||
|
sa.select(jobs.c.status, sa.func.count().label("n"))
|
||||||
|
.where(
|
||||||
|
sa.or_(
|
||||||
|
jobs.c.status.in_(["queued", "claimed"]),
|
||||||
|
jobs.c.completed_at >= since.isoformat(),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
.group_by(jobs.c.status)
|
||||||
|
)
|
||||||
|
async with self._engine.connect() as conn:
|
||||||
|
rows = (await conn.execute(query)).all()
|
||||||
|
return {status: n for status, n in rows}
|
||||||
|
|
||||||
async def count_succeeded_since(self, seconds: int) -> int:
|
async def count_succeeded_since(self, seconds: int) -> int:
|
||||||
"""How many jobs reached `succeeded` in the last `seconds` seconds.
|
"""How many jobs reached `succeeded` in the last `seconds` seconds.
|
||||||
Drives the dashboard's rolling-throughput chips."""
|
Drives the dashboard's rolling-throughput chips."""
|
||||||
|
|
|
||||||
|
|
@ -110,23 +110,38 @@ async def run_job(
|
||||||
),
|
),
|
||||||
):
|
):
|
||||||
try:
|
try:
|
||||||
|
manifest_context = extra.get("_manifest")
|
||||||
if job.op is JobOp.DELETE:
|
if job.op is JobOp.DELETE:
|
||||||
# An atomic-rename save can let a spurious DELETE win the
|
# An atomic-rename save can let a spurious DELETE win the
|
||||||
# enqueue race while the file is mid-rewrite. If the resource
|
# enqueue race while the file is mid-rewrite. If the resource
|
||||||
# is already back, skip the delete (it would blackhole a live
|
# is already back, skip the delete (it would blackhole a live
|
||||||
# document) and let the next sweep re-ingest it.
|
# document) and let the next sweep re-ingest it. Manifest
|
||||||
try:
|
# replay intentionally follows the frozen dry-run changeset.
|
||||||
source = resolve_configured_source(job.uri, job.source_id, sources)
|
if manifest_context is None:
|
||||||
restored = await source.head(job.uri) is not None
|
try:
|
||||||
except Exception:
|
source = resolve_configured_source(
|
||||||
restored = False
|
job.uri, job.source_id, sources
|
||||||
if restored:
|
)
|
||||||
return JobResult(deleted=False)
|
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)
|
doc = await client.get_document_by_uri(job.uri)
|
||||||
if doc is not None and doc.id is not None:
|
if doc is not None and doc.id is not None:
|
||||||
await client.delete_document(doc.id)
|
await client.delete_document(doc.id)
|
||||||
return JobResult(deleted=True)
|
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(
|
result = await client.create_document_from_source(
|
||||||
job.uri,
|
job.uri,
|
||||||
sources=sources,
|
sources=sources,
|
||||||
|
|
|
||||||
|
|
@ -1,13 +1,25 @@
|
||||||
"""haiku-ingester CLI: exercises every subcommand via CliRunner with
|
"""haiku-ingester CLI: exercises every subcommand via CliRunner with
|
||||||
IngesterApp / open_queue patched out so no real ingestion runs."""
|
IngesterApp / open_queue patched out so no real ingestion runs."""
|
||||||
|
|
||||||
|
from contextlib import contextmanager
|
||||||
|
from datetime import UTC, datetime
|
||||||
from unittest.mock import AsyncMock, MagicMock
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
import yaml
|
||||||
from typer.testing import CliRunner
|
from typer.testing import CliRunner
|
||||||
|
|
||||||
from haiku.rag.ingester.app import BatchReport
|
from haiku.rag.config import QueueConfig
|
||||||
|
from haiku.rag.ingester.app import BatchProgress, BatchReport
|
||||||
|
from haiku.rag.ingester.batch import (
|
||||||
|
BatchChange,
|
||||||
|
BatchDryRunReport,
|
||||||
|
BatchManifest,
|
||||||
|
BatchSourceSummary,
|
||||||
|
)
|
||||||
from haiku.rag.ingester.cli import _cli as cli
|
from haiku.rag.ingester.cli import _cli as cli
|
||||||
|
from haiku.rag.ingester.cli import _resolve_queue_config
|
||||||
|
from haiku.rag.ingester.queue.models import JobOp
|
||||||
|
|
||||||
runner = CliRunner()
|
runner = CliRunner()
|
||||||
|
|
||||||
|
|
@ -15,6 +27,17 @@ runner = CliRunner()
|
||||||
# --- helpers ---
|
# --- helpers ---
|
||||||
|
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def _progress_context(callback):
|
||||||
|
yield callback
|
||||||
|
|
||||||
|
|
||||||
|
def _config_with_queue(queue: QueueConfig):
|
||||||
|
config = MagicMock()
|
||||||
|
config.ingester.queue = queue
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
def _fake_app(report: BatchReport, monkeypatch) -> AsyncMock:
|
def _fake_app(report: BatchReport, monkeypatch) -> AsyncMock:
|
||||||
fake = AsyncMock()
|
fake = AsyncMock()
|
||||||
fake.run_batch.return_value = report
|
fake.run_batch.return_value = report
|
||||||
|
|
@ -22,6 +45,54 @@ def _fake_app(report: BatchReport, monkeypatch) -> AsyncMock:
|
||||||
return fake
|
return fake
|
||||||
|
|
||||||
|
|
||||||
|
def _manifest() -> BatchManifest:
|
||||||
|
now = datetime(2026, 6, 22, 10, 30, tzinfo=UTC)
|
||||||
|
return BatchManifest(
|
||||||
|
generated_at=now,
|
||||||
|
sources=[
|
||||||
|
BatchSourceSummary(
|
||||||
|
source_id="docs",
|
||||||
|
upsert_count=1,
|
||||||
|
delete_count=1,
|
||||||
|
unchanged_count=2,
|
||||||
|
)
|
||||||
|
],
|
||||||
|
changes=[
|
||||||
|
BatchChange(
|
||||||
|
op=JobOp.UPSERT,
|
||||||
|
source_id="docs",
|
||||||
|
uri="file:///a.md",
|
||||||
|
revision="r1",
|
||||||
|
discovered_at=now,
|
||||||
|
),
|
||||||
|
BatchChange(
|
||||||
|
op=JobOp.DELETE,
|
||||||
|
source_id="docs",
|
||||||
|
uri="file:///gone.md",
|
||||||
|
discovered_at=now,
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _fake_dry_run_app(report: BatchDryRunReport, monkeypatch) -> AsyncMock:
|
||||||
|
fake = AsyncMock()
|
||||||
|
fake.run_batch_dry_run.return_value = report
|
||||||
|
monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
|
||||||
|
return fake
|
||||||
|
|
||||||
|
|
||||||
|
def _fake_manifest_app(report: BatchReport, monkeypatch) -> AsyncMock:
|
||||||
|
fake = AsyncMock()
|
||||||
|
fake.run_batch_from_manifest.return_value = report
|
||||||
|
monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
|
||||||
|
return fake
|
||||||
|
|
||||||
|
|
||||||
|
def _write_manifest(path) -> None:
|
||||||
|
path.write_text(yaml.safe_dump(_manifest().model_dump(mode="json")))
|
||||||
|
|
||||||
|
|
||||||
def test_run_batch_reports_and_exits_zero(monkeypatch):
|
def test_run_batch_reports_and_exits_zero(monkeypatch):
|
||||||
fake = _fake_app(BatchReport(succeeded=3, dead=0), monkeypatch)
|
fake = _fake_app(BatchReport(succeeded=3, dead=0), monkeypatch)
|
||||||
|
|
||||||
|
|
@ -32,6 +103,27 @@ def test_run_batch_reports_and_exits_zero(monkeypatch):
|
||||||
fake.run_batch.assert_awaited_once()
|
fake.run_batch.assert_awaited_once()
|
||||||
|
|
||||||
|
|
||||||
|
def test_run_batch_passes_progress_callback_when_enabled(monkeypatch):
|
||||||
|
fake = AsyncMock()
|
||||||
|
|
||||||
|
async def _run_batch(*, progress_callback):
|
||||||
|
progress_callback(BatchProgress(total=1, completed=1, succeeded=1))
|
||||||
|
return BatchReport(succeeded=1, dead=0)
|
||||||
|
|
||||||
|
fake.run_batch.side_effect = _run_batch
|
||||||
|
monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"haiku.rag.ingester.cli._batch_progress",
|
||||||
|
lambda _: _progress_context(lambda snapshot: None),
|
||||||
|
)
|
||||||
|
|
||||||
|
result = runner.invoke(cli, ["run-batch", "--db", "x.lancedb"])
|
||||||
|
|
||||||
|
assert result.exit_code == 0
|
||||||
|
fake.run_batch.assert_awaited_once()
|
||||||
|
assert "1 succeeded, 0 dead" in result.output
|
||||||
|
|
||||||
|
|
||||||
def test_run_batch_exits_nonzero_when_dead(monkeypatch):
|
def test_run_batch_exits_nonzero_when_dead(monkeypatch):
|
||||||
_fake_app(BatchReport(succeeded=1, dead=2), monkeypatch)
|
_fake_app(BatchReport(succeeded=1, dead=2), monkeypatch)
|
||||||
|
|
||||||
|
|
@ -50,6 +142,181 @@ def test_run_batch_exits_nonzero_when_sweep_fails(monkeypatch):
|
||||||
assert "failed to sweep: docs" in result.output
|
assert "failed to sweep: docs" in result.output
|
||||||
|
|
||||||
|
|
||||||
|
def test_run_batch_dry_run_writes_default_manifest(monkeypatch, tmp_path):
|
||||||
|
fake = _fake_dry_run_app(
|
||||||
|
BatchDryRunReport(manifest=_manifest()),
|
||||||
|
monkeypatch,
|
||||||
|
)
|
||||||
|
|
||||||
|
with runner.isolated_filesystem(temp_dir=tmp_path):
|
||||||
|
result = runner.invoke(cli, ["run-batch", "--dry-run", "--db", "x.lancedb"])
|
||||||
|
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "Dry run complete: 1 upsert, 1 delete, 2 unchanged -> manifest-" in (
|
||||||
|
result.output
|
||||||
|
)
|
||||||
|
written = list(tmp_path.glob("*/manifest-*.yaml"))
|
||||||
|
assert len(written) == 1
|
||||||
|
data = yaml.safe_load(written[0].read_text())
|
||||||
|
|
||||||
|
assert data["version"] == 1
|
||||||
|
assert data["sources"][0]["source_id"] == "docs"
|
||||||
|
assert [change["op"] for change in data["changes"]] == ["upsert", "delete"]
|
||||||
|
fake.run_batch_dry_run.assert_awaited_once()
|
||||||
|
fake.run_batch.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
|
def test_run_batch_dry_run_writes_explicit_output(monkeypatch, tmp_path):
|
||||||
|
output = tmp_path / "custom.yaml"
|
||||||
|
_fake_dry_run_app(BatchDryRunReport(manifest=_manifest()), monkeypatch)
|
||||||
|
|
||||||
|
result = runner.invoke(
|
||||||
|
cli,
|
||||||
|
["run-batch", "--dry-run", "--output", str(output), "--db", "x.lancedb"],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert f"-> {output}" in result.output
|
||||||
|
data = yaml.safe_load(output.read_text())
|
||||||
|
assert data["changes"][0]["uri"] == "file:///a.md"
|
||||||
|
|
||||||
|
|
||||||
|
def test_run_batch_dry_run_exits_nonzero_when_sweep_fails(monkeypatch, tmp_path):
|
||||||
|
output = tmp_path / "failed.yaml"
|
||||||
|
_fake_dry_run_app(
|
||||||
|
BatchDryRunReport(manifest=_manifest(), failed_sweeps=["docs"]),
|
||||||
|
monkeypatch,
|
||||||
|
)
|
||||||
|
|
||||||
|
result = runner.invoke(
|
||||||
|
cli,
|
||||||
|
["run-batch", "--dry-run", "--output", str(output), "--db", "x.lancedb"],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.exit_code == 1
|
||||||
|
assert "failed to sweep: docs" in result.output
|
||||||
|
assert not output.exists()
|
||||||
|
|
||||||
|
|
||||||
|
def test_run_batch_manifest_replays_manifest(monkeypatch, tmp_path):
|
||||||
|
manifest_path = tmp_path / "manifest.yaml"
|
||||||
|
_write_manifest(manifest_path)
|
||||||
|
fake = _fake_manifest_app(BatchReport(succeeded=2, dead=0), monkeypatch)
|
||||||
|
|
||||||
|
result = runner.invoke(
|
||||||
|
cli, ["run-batch", "--manifest", str(manifest_path), "--db", "x.lancedb"]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "Manifest batch complete: 2 succeeded, 0 dead" in result.output
|
||||||
|
fake.run_batch_from_manifest.assert_awaited_once()
|
||||||
|
loaded = fake.run_batch_from_manifest.await_args.args[0]
|
||||||
|
assert isinstance(loaded, BatchManifest)
|
||||||
|
assert loaded.changes[0].uri == "file:///a.md"
|
||||||
|
fake.run_batch.assert_not_awaited()
|
||||||
|
fake.run_batch_dry_run.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
|
def test_run_batch_manifest_passes_progress_callback(monkeypatch, tmp_path):
|
||||||
|
manifest_path = tmp_path / "manifest.yaml"
|
||||||
|
_write_manifest(manifest_path)
|
||||||
|
fake = AsyncMock()
|
||||||
|
|
||||||
|
async def _run_manifest(_manifest, *, progress_callback):
|
||||||
|
progress_callback(BatchProgress(total=1, completed=1, succeeded=1))
|
||||||
|
return BatchReport(succeeded=1, dead=0)
|
||||||
|
|
||||||
|
fake.run_batch_from_manifest.side_effect = _run_manifest
|
||||||
|
monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"haiku.rag.ingester.cli._batch_progress",
|
||||||
|
lambda _: _progress_context(lambda snapshot: None),
|
||||||
|
)
|
||||||
|
|
||||||
|
result = runner.invoke(cli, ["run-batch", "--manifest", str(manifest_path)])
|
||||||
|
|
||||||
|
assert result.exit_code == 0
|
||||||
|
fake.run_batch_from_manifest.assert_awaited_once()
|
||||||
|
assert "1 succeeded, 0 dead" in result.output
|
||||||
|
|
||||||
|
|
||||||
|
def test_run_batch_manifest_exits_nonzero_when_dead(monkeypatch, tmp_path):
|
||||||
|
manifest_path = tmp_path / "manifest.yaml"
|
||||||
|
_write_manifest(manifest_path)
|
||||||
|
_fake_manifest_app(BatchReport(succeeded=1, dead=1), monkeypatch)
|
||||||
|
|
||||||
|
result = runner.invoke(cli, ["run-batch", "--manifest", str(manifest_path)])
|
||||||
|
|
||||||
|
assert result.exit_code == 1
|
||||||
|
assert "1 dead" in result.output
|
||||||
|
|
||||||
|
|
||||||
|
def test_run_batch_manifest_reports_validation_error(monkeypatch, tmp_path):
|
||||||
|
manifest_path = tmp_path / "manifest.yaml"
|
||||||
|
_write_manifest(manifest_path)
|
||||||
|
fake = AsyncMock()
|
||||||
|
fake.run_batch_from_manifest.side_effect = ValueError("bad manifest")
|
||||||
|
monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
|
||||||
|
|
||||||
|
result = runner.invoke(cli, ["run-batch", "--manifest", str(manifest_path)])
|
||||||
|
|
||||||
|
assert result.exit_code == 1
|
||||||
|
assert "Error: bad manifest" in result.output
|
||||||
|
|
||||||
|
|
||||||
|
def test_run_batch_manifest_conflicts_with_dry_run(tmp_path):
|
||||||
|
manifest_path = tmp_path / "manifest.yaml"
|
||||||
|
_write_manifest(manifest_path)
|
||||||
|
|
||||||
|
result = runner.invoke(
|
||||||
|
cli, ["run-batch", "--manifest", str(manifest_path), "--dry-run"]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.exit_code != 0
|
||||||
|
assert "--manifest cannot be combined with --dry-run" in result.output
|
||||||
|
|
||||||
|
|
||||||
|
def test_run_batch_manifest_conflicts_with_output(tmp_path):
|
||||||
|
manifest_path = tmp_path / "manifest.yaml"
|
||||||
|
_write_manifest(manifest_path)
|
||||||
|
|
||||||
|
result = runner.invoke(
|
||||||
|
cli,
|
||||||
|
[
|
||||||
|
"run-batch",
|
||||||
|
"--manifest",
|
||||||
|
str(manifest_path),
|
||||||
|
"--output",
|
||||||
|
str(tmp_path / "out.yaml"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.exit_code != 0
|
||||||
|
assert "--output is only valid with --dry-run" in result.output
|
||||||
|
|
||||||
|
|
||||||
|
def test_run_batch_output_requires_dry_run(tmp_path):
|
||||||
|
result = runner.invoke(
|
||||||
|
cli,
|
||||||
|
["run-batch", "--output", str(tmp_path / "out.yaml")],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.exit_code != 0
|
||||||
|
assert "--output is only valid with --dry-run" in result.output
|
||||||
|
|
||||||
|
|
||||||
|
def test_resolve_queue_config_keeps_dburi_when_path_override_present(tmp_path):
|
||||||
|
queue = QueueConfig(
|
||||||
|
dburi="postgresql+asyncpg://user:pass@example.test/db",
|
||||||
|
path=tmp_path / "configured.db",
|
||||||
|
)
|
||||||
|
config = _config_with_queue(queue)
|
||||||
|
|
||||||
|
resolved = _resolve_queue_config(config, tmp_path / "override.db")
|
||||||
|
|
||||||
|
assert resolved is queue
|
||||||
|
|
||||||
|
|
||||||
# --- serve ---
|
# --- serve ---
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -16,6 +16,7 @@ def _job(
|
||||||
*,
|
*,
|
||||||
op: JobOp = JobOp.UPSERT,
|
op: JobOp = JobOp.UPSERT,
|
||||||
uri: str = "https://example.com/a.pdf",
|
uri: str = "https://example.com/a.pdf",
|
||||||
|
revision: str | None = None,
|
||||||
extra: dict | None = None,
|
extra: dict | None = None,
|
||||||
attempts: int = 0,
|
attempts: int = 0,
|
||||||
) -> Job:
|
) -> Job:
|
||||||
|
|
@ -25,6 +26,7 @@ def _job(
|
||||||
source_id="src",
|
source_id="src",
|
||||||
uri=uri,
|
uri=uri,
|
||||||
op=op,
|
op=op,
|
||||||
|
revision=revision,
|
||||||
status=JobStatus.CLAIMED,
|
status=JobStatus.CLAIMED,
|
||||||
attempts=attempts,
|
attempts=attempts,
|
||||||
max_attempts=5,
|
max_attempts=5,
|
||||||
|
|
@ -255,6 +257,22 @@ async def test_delete_skipped_when_resource_restored_on_source():
|
||||||
client.delete_document.assert_not_awaited()
|
client.delete_document.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_manifest_delete_proceeds_when_resource_restored_on_source():
|
||||||
|
client = _mock_client()
|
||||||
|
client.get_document_by_uri.return_value = Document(id="doc-9", content="", uri="u")
|
||||||
|
sources: list[Source] = [_StubSource("src", "12345")]
|
||||||
|
|
||||||
|
result = await run_job(
|
||||||
|
client,
|
||||||
|
_job(op=JobOp.DELETE, extra={"_manifest": {"version": 1}}),
|
||||||
|
sources=sources,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.deleted is True
|
||||||
|
client.delete_document.assert_awaited_once_with("doc-9")
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_delete_proceeds_when_resource_absent_on_source():
|
async def test_delete_proceeds_when_resource_absent_on_source():
|
||||||
client = _mock_client()
|
client = _mock_client()
|
||||||
|
|
@ -267,6 +285,42 @@ async def test_delete_proceeds_when_resource_absent_on_source():
|
||||||
client.delete_document.assert_awaited_once_with("doc-9")
|
client.delete_document.assert_awaited_once_with("doc-9")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_manifest_upsert_rejects_stale_revision_before_fetch():
|
||||||
|
client = _mock_client()
|
||||||
|
sources: list[Source] = [_StubSource("src", "r2")]
|
||||||
|
|
||||||
|
with pytest.raises(PermanentError, match="manifest revision is stale"):
|
||||||
|
await run_job(
|
||||||
|
client,
|
||||||
|
_job(revision="r1", extra={"_manifest": {"version": 1}}),
|
||||||
|
sources=sources,
|
||||||
|
)
|
||||||
|
|
||||||
|
client.create_document_from_source.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_manifest_upsert_ingests_when_revision_matches():
|
||||||
|
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", "source_revision": "r1"},
|
||||||
|
)
|
||||||
|
sources: list[Source] = [_StubSource("src", "r1")]
|
||||||
|
|
||||||
|
result = await run_job(
|
||||||
|
client,
|
||||||
|
_job(revision="r1", extra={"_manifest": {"version": 1}}),
|
||||||
|
sources=sources,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.document_id == "doc-42"
|
||||||
|
client.create_document_from_source.assert_awaited_once()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_delete_proceeds_when_source_unresolvable():
|
async def test_delete_proceeds_when_source_unresolvable():
|
||||||
"""No configured source for the job: the probe can't run, so the delete
|
"""No configured source for the job: the probe can't run, so the delete
|
||||||
|
|
|
||||||
|
|
@ -309,6 +309,78 @@ async def test_per_source_retry_policy_overrides_default(jobs, sync, tmp_path):
|
||||||
assert queued[0].max_attempts == 9
|
assert queued[0].max_attempts == 9
|
||||||
|
|
||||||
|
|
||||||
|
# --- dry-run collection ---
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_dry_run_reports_changes_without_mutating_queue_or_sync(
|
||||||
|
fs_config, jobs, sync
|
||||||
|
):
|
||||||
|
source = _StubSource(
|
||||||
|
"src",
|
||||||
|
[
|
||||||
|
[
|
||||||
|
_event("file:///a.md", revision="r1"),
|
||||||
|
_event(
|
||||||
|
"file:///b.md",
|
||||||
|
kind=SourceEventKind.UNCHANGED,
|
||||||
|
revision="r2",
|
||||||
|
),
|
||||||
|
_event("file:///gone.md", kind=SourceEventKind.DELETE),
|
||||||
|
]
|
||||||
|
],
|
||||||
|
)
|
||||||
|
poller = _periodic(source, fs_config, jobs, sync)
|
||||||
|
|
||||||
|
ok, summary, changes = await poller._dry_run_once()
|
||||||
|
|
||||||
|
assert ok is True
|
||||||
|
assert summary.upsert_count == 1
|
||||||
|
assert summary.delete_count == 1
|
||||||
|
assert summary.unchanged_count == 1
|
||||||
|
assert summary.ignored_delete_count == 0
|
||||||
|
assert [(c.op, c.uri, c.revision) for c in changes] == [
|
||||||
|
(JobOp.UPSERT, "file:///a.md", "r1"),
|
||||||
|
(JobOp.DELETE, "file:///gone.md", None),
|
||||||
|
]
|
||||||
|
assert await jobs.list_jobs(source_id="src") == []
|
||||||
|
assert await sync.list_known_uris("src") == set()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_dry_run_counts_ignored_deletes_when_orphan_delete_disabled(
|
||||||
|
fs_config, jobs, sync
|
||||||
|
):
|
||||||
|
cfg = fs_config.model_copy(update={"delete_orphans": False})
|
||||||
|
source = _StubSource(
|
||||||
|
"src", [[_event("file:///gone.md", kind=SourceEventKind.DELETE)]]
|
||||||
|
)
|
||||||
|
poller = _periodic(source, cfg, jobs, sync)
|
||||||
|
|
||||||
|
ok, summary, changes = await poller._dry_run_once()
|
||||||
|
|
||||||
|
assert ok is True
|
||||||
|
assert summary.delete_count == 0
|
||||||
|
assert summary.ignored_delete_count == 1
|
||||||
|
assert changes == []
|
||||||
|
assert await jobs.list_jobs(source_id="src") == []
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_dry_run_skips_when_queue_has_pending_work(fs_config, jobs, sync):
|
||||||
|
await jobs.enqueue("src", "file:///already.md", op=JobOp.UPSERT)
|
||||||
|
source = _StubSource("src", [[_event("file:///a.md")]])
|
||||||
|
poller = _periodic(source, fs_config, jobs, sync)
|
||||||
|
|
||||||
|
ok, summary, changes = await poller._dry_run_once()
|
||||||
|
|
||||||
|
assert ok is False
|
||||||
|
assert summary.source_id == "src"
|
||||||
|
assert changes == []
|
||||||
|
assert source.discover_calls == 0
|
||||||
|
assert poller.last_skip_reason == "pending_work"
|
||||||
|
|
||||||
|
|
||||||
# --- PollerManager lifecycle ---
|
# --- PollerManager lifecycle ---
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -5,6 +5,7 @@ pruning, not embedding."""
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
from contextlib import asynccontextmanager
|
from contextlib import asynccontextmanager
|
||||||
|
from datetime import UTC, datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from unittest.mock import AsyncMock
|
from unittest.mock import AsyncMock
|
||||||
from urllib.parse import unquote, urlparse
|
from urllib.parse import unquote, urlparse
|
||||||
|
|
@ -23,7 +24,11 @@ from haiku.rag.config import (
|
||||||
WorkerConfig,
|
WorkerConfig,
|
||||||
)
|
)
|
||||||
from haiku.rag.ingester.app import IngesterApp
|
from haiku.rag.ingester.app import IngesterApp
|
||||||
|
from haiku.rag.ingester.batch import BatchChange, BatchManifest
|
||||||
from haiku.rag.ingester.pollers.manager import PollerManager
|
from haiku.rag.ingester.pollers.manager import PollerManager
|
||||||
|
from haiku.rag.ingester.queue.migrations import open_queue
|
||||||
|
from haiku.rag.ingester.queue.models import JobOp
|
||||||
|
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
|
||||||
from haiku.rag.ingester.workers.pool import WorkerPool
|
from haiku.rag.ingester.workers.pool import WorkerPool
|
||||||
from haiku.rag.store.models.document import Document
|
from haiku.rag.store.models.document import Document
|
||||||
|
|
||||||
|
|
@ -83,6 +88,10 @@ def _mock_client() -> AsyncMock:
|
||||||
return client
|
return client
|
||||||
|
|
||||||
|
|
||||||
|
def _manifest(*changes: BatchChange) -> BatchManifest:
|
||||||
|
return BatchManifest(generated_at=datetime.now(UTC), changes=list(changes))
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def use_client(monkeypatch):
|
def use_client(monkeypatch):
|
||||||
"""Make IngesterApp's internally-created HaikuRAG resolve to the given
|
"""Make IngesterApp's internally-created HaikuRAG resolve to the given
|
||||||
|
|
@ -123,6 +132,30 @@ async def test_run_batch_drains_upserts(tmp_path, use_client):
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_batch_reports_progress(tmp_path, use_client):
|
||||||
|
(tmp_path / "a.md").write_text("hello")
|
||||||
|
(tmp_path / "b.md").write_text("world")
|
||||||
|
|
||||||
|
client = _mock_client()
|
||||||
|
use_client(client)
|
||||||
|
progress = []
|
||||||
|
|
||||||
|
report = await IngesterApp(
|
||||||
|
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
|
||||||
|
).run_batch(progress_callback=progress.append)
|
||||||
|
|
||||||
|
assert report.succeeded == 2
|
||||||
|
assert report.dead == 0
|
||||||
|
assert progress
|
||||||
|
assert progress[-1].total == 2
|
||||||
|
assert progress[-1].completed == 2
|
||||||
|
assert progress[-1].succeeded == 2
|
||||||
|
assert progress[-1].dead == 0
|
||||||
|
assert progress[-1].queued == 0
|
||||||
|
assert progress[-1].claimed == 0
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_run_batch_prunes_orphans(tmp_path, use_client):
|
async def test_run_batch_prunes_orphans(tmp_path, use_client):
|
||||||
(tmp_path / "a.md").write_text("hello")
|
(tmp_path / "a.md").write_text("hello")
|
||||||
|
|
@ -239,6 +272,290 @@ async def test_run_batch_empty_source_returns_immediately(tmp_path, use_client):
|
||||||
client.create_document_from_source.assert_not_awaited()
|
client.create_document_from_source.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_batch_dry_run_reports_manifest_without_mutating_queue(tmp_path):
|
||||||
|
(tmp_path / "a.md").write_text("hello")
|
||||||
|
config = _config(tmp_path)
|
||||||
|
db_path = tmp_path / "db.lancedb"
|
||||||
|
|
||||||
|
engine = await open_queue(config.ingester.queue)
|
||||||
|
try:
|
||||||
|
sync = SyncStateRepo(engine)
|
||||||
|
await sync.upsert("local", (tmp_path / "gone.md").as_uri(), revision="old")
|
||||||
|
finally:
|
||||||
|
await engine.dispose()
|
||||||
|
|
||||||
|
report = await IngesterApp(config=config, db_path=db_path).run_batch_dry_run()
|
||||||
|
|
||||||
|
assert report.failed_sweeps == []
|
||||||
|
assert report.manifest.version == 1
|
||||||
|
assert [(change.op, change.uri) for change in report.manifest.changes] == [
|
||||||
|
(JobOp.UPSERT, (tmp_path / "a.md").as_uri()),
|
||||||
|
(JobOp.DELETE, (tmp_path / "gone.md").as_uri()),
|
||||||
|
]
|
||||||
|
source_summary = report.manifest.sources[0]
|
||||||
|
assert source_summary.source_id == "local"
|
||||||
|
assert source_summary.upsert_count == 1
|
||||||
|
assert source_summary.delete_count == 1
|
||||||
|
|
||||||
|
engine = await open_queue(config.ingester.queue)
|
||||||
|
try:
|
||||||
|
jobs = JobRepo(engine)
|
||||||
|
sync = SyncStateRepo(engine)
|
||||||
|
assert await jobs.list_jobs(source_id="local") == []
|
||||||
|
assert await sync.list_known_uris("local") == {(tmp_path / "gone.md").as_uri()}
|
||||||
|
finally:
|
||||||
|
await engine.dispose()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_batch_from_manifest_drains_changes_without_sweeping(
|
||||||
|
tmp_path, use_client, monkeypatch
|
||||||
|
):
|
||||||
|
(tmp_path / "a.md").write_text("hello")
|
||||||
|
revision = str((tmp_path / "a.md").stat().st_mtime_ns)
|
||||||
|
client = _mock_client()
|
||||||
|
use_client(client)
|
||||||
|
sweep_all = AsyncMock(side_effect=AssertionError("manifest replay must not sweep"))
|
||||||
|
monkeypatch.setattr(PollerManager, "sweep_all", sweep_all)
|
||||||
|
|
||||||
|
report = await IngesterApp(
|
||||||
|
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
|
||||||
|
).run_batch_from_manifest(
|
||||||
|
_manifest(
|
||||||
|
BatchChange(
|
||||||
|
op=JobOp.UPSERT,
|
||||||
|
source_id="local",
|
||||||
|
uri=(tmp_path / "a.md").as_uri(),
|
||||||
|
revision=revision,
|
||||||
|
discovered_at=datetime.now(UTC),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert report.succeeded == 1
|
||||||
|
assert report.dead == 0
|
||||||
|
client.create_document_from_source.assert_awaited_once()
|
||||||
|
sweep_all.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_batch_from_manifest_rejects_stale_upsert_revision(
|
||||||
|
tmp_path, use_client
|
||||||
|
):
|
||||||
|
(tmp_path / "a.md").write_text("hello")
|
||||||
|
client = _mock_client()
|
||||||
|
use_client(client)
|
||||||
|
|
||||||
|
report = await IngesterApp(
|
||||||
|
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
|
||||||
|
).run_batch_from_manifest(
|
||||||
|
_manifest(
|
||||||
|
BatchChange(
|
||||||
|
op=JobOp.UPSERT,
|
||||||
|
source_id="local",
|
||||||
|
uri=(tmp_path / "a.md").as_uri(),
|
||||||
|
revision="stale",
|
||||||
|
discovered_at=datetime.now(UTC),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert report.succeeded == 0
|
||||||
|
assert report.dead == 1
|
||||||
|
client.create_document_from_source.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_batch_from_manifest_delete_uses_manifest_even_if_file_reappears(
|
||||||
|
tmp_path, use_client
|
||||||
|
):
|
||||||
|
path = tmp_path / "gone.md"
|
||||||
|
path.write_text("back")
|
||||||
|
client = _mock_client()
|
||||||
|
use_client(client)
|
||||||
|
|
||||||
|
report = await IngesterApp(
|
||||||
|
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
|
||||||
|
).run_batch_from_manifest(
|
||||||
|
_manifest(
|
||||||
|
BatchChange(
|
||||||
|
op=JobOp.DELETE,
|
||||||
|
source_id="local",
|
||||||
|
uri=path.as_uri(),
|
||||||
|
discovered_at=datetime.now(UTC),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert report.succeeded == 1
|
||||||
|
assert report.dead == 0
|
||||||
|
client.delete_document.assert_awaited_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_batch_from_manifest_resumes_same_manifest_work(tmp_path, use_client):
|
||||||
|
(tmp_path / "a.md").write_text("hello")
|
||||||
|
revision = str((tmp_path / "a.md").stat().st_mtime_ns)
|
||||||
|
config = _config(tmp_path)
|
||||||
|
client = _mock_client()
|
||||||
|
use_client(client)
|
||||||
|
manifest = _manifest(
|
||||||
|
BatchChange(
|
||||||
|
op=JobOp.UPSERT,
|
||||||
|
source_id="local",
|
||||||
|
uri=(tmp_path / "a.md").as_uri(),
|
||||||
|
revision=revision,
|
||||||
|
discovered_at=datetime.now(UTC),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
engine = await open_queue(config.ingester.queue)
|
||||||
|
try:
|
||||||
|
jobs = JobRepo(engine)
|
||||||
|
await jobs.enqueue(
|
||||||
|
"local",
|
||||||
|
(tmp_path / "a.md").as_uri(),
|
||||||
|
op=JobOp.UPSERT,
|
||||||
|
revision=revision,
|
||||||
|
extra={
|
||||||
|
"_manifest": {
|
||||||
|
"version": manifest.version,
|
||||||
|
"generated_at": manifest.generated_at.isoformat(),
|
||||||
|
"discovered_at": manifest.changes[0].discovered_at.isoformat(),
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
await engine.dispose()
|
||||||
|
|
||||||
|
report = await IngesterApp(
|
||||||
|
config=config, db_path=tmp_path / "db.lancedb"
|
||||||
|
).run_batch_from_manifest(manifest)
|
||||||
|
|
||||||
|
assert report.succeeded == 1
|
||||||
|
assert report.dead == 0
|
||||||
|
client.create_document_from_source.assert_awaited_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_batch_from_manifest_rejects_non_manifest_pending_work(
|
||||||
|
tmp_path, use_client
|
||||||
|
):
|
||||||
|
(tmp_path / "a.md").write_text("hello")
|
||||||
|
config = _config(tmp_path)
|
||||||
|
client = _mock_client()
|
||||||
|
use_client(client)
|
||||||
|
engine = await open_queue(config.ingester.queue)
|
||||||
|
try:
|
||||||
|
jobs = JobRepo(engine)
|
||||||
|
await jobs.enqueue("local", (tmp_path / "a.md").as_uri(), op=JobOp.UPSERT)
|
||||||
|
finally:
|
||||||
|
await engine.dispose()
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="non-manifest pending work"):
|
||||||
|
await IngesterApp(
|
||||||
|
config=config, db_path=tmp_path / "db.lancedb"
|
||||||
|
).run_batch_from_manifest(
|
||||||
|
_manifest(
|
||||||
|
BatchChange(
|
||||||
|
op=JobOp.UPSERT,
|
||||||
|
source_id="local",
|
||||||
|
uri=(tmp_path / "a.md").as_uri(),
|
||||||
|
discovered_at=datetime.now(UTC),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_batch_from_manifest_rejects_different_manifest_pending_work(
|
||||||
|
tmp_path, use_client
|
||||||
|
):
|
||||||
|
(tmp_path / "a.md").write_text("hello")
|
||||||
|
config = _config(tmp_path)
|
||||||
|
client = _mock_client()
|
||||||
|
use_client(client)
|
||||||
|
manifest = _manifest(
|
||||||
|
BatchChange(
|
||||||
|
op=JobOp.UPSERT,
|
||||||
|
source_id="local",
|
||||||
|
uri=(tmp_path / "a.md").as_uri(),
|
||||||
|
discovered_at=datetime.now(UTC),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
engine = await open_queue(config.ingester.queue)
|
||||||
|
try:
|
||||||
|
jobs = JobRepo(engine)
|
||||||
|
await jobs.enqueue(
|
||||||
|
"local",
|
||||||
|
(tmp_path / "a.md").as_uri(),
|
||||||
|
op=JobOp.UPSERT,
|
||||||
|
extra={
|
||||||
|
"_manifest": {
|
||||||
|
"version": 1,
|
||||||
|
"generated_at": "2026-01-01T00:00:00+00:00",
|
||||||
|
"discovered_at": manifest.changes[0].discovered_at.isoformat(),
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
await engine.dispose()
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="non-manifest pending work"):
|
||||||
|
await IngesterApp(
|
||||||
|
config=config, db_path=tmp_path / "db.lancedb"
|
||||||
|
).run_batch_from_manifest(manifest)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_batch_from_manifest_rejects_unrelated_pending_work(
|
||||||
|
tmp_path, use_client
|
||||||
|
):
|
||||||
|
(tmp_path / "a.md").write_text("hello")
|
||||||
|
config = _config(tmp_path)
|
||||||
|
client = _mock_client()
|
||||||
|
use_client(client)
|
||||||
|
engine = await open_queue(config.ingester.queue)
|
||||||
|
try:
|
||||||
|
jobs = JobRepo(engine)
|
||||||
|
await jobs.enqueue("other", "file:///outside.md", op=JobOp.UPSERT)
|
||||||
|
finally:
|
||||||
|
await engine.dispose()
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="non-manifest pending work"):
|
||||||
|
await IngesterApp(
|
||||||
|
config=config, db_path=tmp_path / "db.lancedb"
|
||||||
|
).run_batch_from_manifest(
|
||||||
|
_manifest(
|
||||||
|
BatchChange(
|
||||||
|
op=JobOp.UPSERT,
|
||||||
|
source_id="local",
|
||||||
|
uri=(tmp_path / "a.md").as_uri(),
|
||||||
|
discovered_at=datetime.now(UTC),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_batch_from_manifest_rejects_duplicate_changes(tmp_path, use_client):
|
||||||
|
path = tmp_path / "a.md"
|
||||||
|
path.write_text("hello")
|
||||||
|
client = _mock_client()
|
||||||
|
use_client(client)
|
||||||
|
change = BatchChange(
|
||||||
|
op=JobOp.UPSERT,
|
||||||
|
source_id="local",
|
||||||
|
uri=path.as_uri(),
|
||||||
|
discovered_at=datetime.now(UTC),
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="duplicate change"):
|
||||||
|
await IngesterApp(
|
||||||
|
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
|
||||||
|
).run_batch_from_manifest(_manifest(change, change))
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_run_batch_aborts_when_all_workers_die(
|
async def test_run_batch_aborts_when_all_workers_die(
|
||||||
tmp_path, use_client, monkeypatch, caplog
|
tmp_path, use_client, monkeypatch, caplog
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue