Add haiku-ingester run-batch, remove run-once

CHANGELOG.md

@@ -1,6 +1,14 @@
 # Changelog
 ## [Unreleased]
 
+### Added
+
+- `haiku-ingester run-batch`: one discover sweep across every configured source, drains the queue, then exits. Orphan deletion requires a persisted `ingester.db`.
+
+### Removed
+
+- `haiku-ingester run-once`. Use `run-batch` to ingest configured sources, or `serve` for continuous operation.
+
 ## [0.51.0] - 2026-05-29
 
 ### Added

docs/ingester.md

@@ -307,19 +307,23 @@ ingester:
 
 ## Operating
 
-### Smoke-test a single URI
+### One-shot batch build
 
-`run-once` bypasses the queue and runs a single Job through the
+`run-batch` runs a single discover sweep across every configured source,
-pipeline. Useful for sanity-checking a source before starting the
+drains the queue, then exits. New and changed resources are ingested,
-service.
+resources that vanished from a source are deleted. The periodic poller
+loops never start, so the run is deterministic and finishes as soon as the
+queue is empty. This is the mode for building a database in CI or on a
+schedule rather than running the service continuously.
 
 ```bash
-haiku-ingester run-once /path/to/test.pdf
+haiku-ingester run-batch
-haiku-ingester run-once https://example.com/spec.pdf
+haiku-ingester run-batch --db rag.lancedb
-haiku-ingester run-once s3://my-bucket/key.pdf
 ```
 
-Exit codes: `0` success, `1` transient error, `2` permanent error.
+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 non-zero if any job dead-letters.
 
 ### The queue
 

haiku_rag_slim/haiku/rag/ingester/app.py

@@ -1,9 +1,13 @@
 import asyncio
 import logging
 import signal
+from contextlib import asynccontextmanager
+from datetime import UTC, datetime
 from pathlib import Path
 from typing import TYPE_CHECKING
 
+from pydantic import BaseModel
+
 from haiku.rag.config import AppConfig
 from haiku.rag.ingester.pollers.manager import PollerManager
 from haiku.rag.ingester.queue.migrations import open_queue
@@ -17,6 +21,14 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+class BatchReport(BaseModel):
+    """Outcome of a one-shot batch run: terminal job counts after the queue
+    drained."""
+
+    succeeded: int = 0
+    dead: int = 0
+
+
 class IngesterApp:
     """Top-level lifecycle for the production ingester.
 
@@ -32,11 +44,13 @@ class IngesterApp:
         self._sync: SyncStateRepo | None = None
         self._pool: WorkerPool | None = None
         self._pollers: PollerManager | None = None
-        self._client = None
 
-    async def serve(self, *, api: bool = True) -> None:
+    @asynccontextmanager
-        """Run pollers + workers (and the HTTP API when enabled) until a
+    async def _resources(self):
-        SIGINT/SIGTERM is received. Drains in-flight work on shutdown."""
+        """Open the queue connection and construct the repos, client, pollers
+        and worker pool. Yields with everything built but nothing started —
+        callers own the start/stop lifecycle. Closes the client and queue
+        connection on exit."""
         from haiku.rag.client import HaikuRAG
         from haiku.rag.converters import get_converter
 
@@ -58,13 +72,12 @@ class IngesterApp:
                 jitter=ingester_cfg.workers.retry.jitter,
             )
 
-            # The ingester is the sole writer for its LanceDB target;
+            # The ingester is the sole writer for its LanceDB target; create on
-            # create on first start so docker-compose / fresh deployments
+            # first start so docker-compose / fresh deployments don't require a
-            # don't require a manual `haiku-rag init`.
+            # manual `haiku-rag init`.
             async with HaikuRAG(
                 self._db_path, config=self._config, create=True
             ) as client:
-                self._client = client
                 self._pollers = PollerManager(
                     configs=ingester_cfg.sources,
                     job_repo=self._jobs,
@@ -86,73 +99,7 @@ class IngesterApp:
                     # HTTP / WebDAV / S3 fetches reuse credentials.
                     sources=self._pollers.sources,
                 )
-
+                yield
-                stop_event = asyncio.Event()
-                loop = asyncio.get_running_loop()
-                for sig in (signal.SIGINT, signal.SIGTERM):
-                    try:
-                        loop.add_signal_handler(sig, stop_event.set)
-                    except NotImplementedError:
-                        # Windows; signal handlers unavailable in asyncio.
-                        pass
-
-                await self._pollers.start()
-                await self._pool.start()
-                # Log the docling-serve fleet size when relevant so the
-                # operator can eyeball the worker/instance ratio. The convert
-                # phase is usually the throughput ceiling.
-                proc = self._config.processing
-                uses_docling_serve = (
-                    proc.converter == "docling-serve" or proc.chunker == "docling-serve"
-                )
-                if uses_docling_serve:
-                    logger.info(
-                        "Ingester running: %d worker(s), %d source(s), "
-                        "%d docling-serve instance(s)",
-                        ingester_cfg.workers.worker_count,
-                        len(ingester_cfg.sources),
-                        len(self._config.providers.docling_serve.base_urls),
-                    )
-                else:
-                    logger.info(
-                        "Ingester running: %d worker(s), %d source(s)",
-                        ingester_cfg.workers.worker_count,
-                        len(ingester_cfg.sources),
-                    )
-
-                api_task, api_server = await self._maybe_start_api(api)
-
-                try:
-                    await stop_event.wait()
-                finally:
-                    logger.info("Shutting down ingester")
-                    if api_server is not None:
-                        api_server.should_exit = True
-                    if api_task is not None:
-                        await asyncio.gather(api_task, return_exceptions=True)
-                    await self._pollers.stop()
-                    grace_s = ingester_cfg.workers.shutdown_grace_s
-                    try:
-                        await asyncio.wait_for(self._pool.stop(), timeout=grace_s)
-                    except TimeoutError:
-                        # In-flight jobs stay 'claimed'; the reaper resets
-                        # them after claim_timeout_s on the next start.
-                        logger.warning(
-                            "Shutdown grace of %.1fs elapsed with jobs still "
-                            "in flight; cancelling — they'll be reclaimed after "
-                            "claim_timeout_s on next start",
-                            grace_s,
-                        )
-                    # Drain any cancel-cleanup release Tasks the worker pool
-                    # spawned but didn't get to await (timeout path). They
-                    # need the queue connection that the outer finally is
-                    # about to close.
-                    landed = await self._pool.drain_pending_releases(timeout=2.0)
-                    if landed:
-                        logger.info(
-                            "Drained %d cancel-cleanup release(s) before close",
-                            landed,
-                        )
         finally:
             # Close the queue connection unconditionally. aiosqlite runs the
             # underlying sqlite3 in a background thread; leaving it open holds
@@ -162,6 +109,109 @@ class IngesterApp:
                 await self._queue_conn.close()
                 self._queue_conn = None
 
+    async def _stop_pool(self) -> None:
+        """Stop the worker pool, honouring the shutdown grace, then drain any
+        cancel-cleanup release tasks before the queue connection closes."""
+        assert self._pool is not None
+        grace_s = self._config.ingester.workers.shutdown_grace_s
+        try:
+            await asyncio.wait_for(self._pool.stop(), timeout=grace_s)
+        except TimeoutError:
+            # In-flight jobs stay 'claimed'; the reaper resets them after
+            # claim_timeout_s on the next start.
+            logger.warning(
+                "Shutdown grace of %.1fs elapsed with jobs still in flight; "
+                "cancelling — they'll be reclaimed after claim_timeout_s on "
+                "next start",
+                grace_s,
+            )
+        landed = await self._pool.drain_pending_releases(timeout=2.0)
+        if landed:
+            logger.info("Drained %d cancel-cleanup release(s) before close", landed)
+
+    async def serve(self, *, api: bool = True) -> None:
+        """Run pollers + workers (and the HTTP API when enabled) until a
+        SIGINT/SIGTERM is received. Drains in-flight work on shutdown."""
+        ingester_cfg = self._config.ingester
+        async with self._resources():
+            assert self._pollers is not None and self._pool is not None
+            await self._pollers.start()
+            await self._pool.start()
+            # Log the docling-serve fleet size when relevant so the operator
+            # can eyeball the worker/instance ratio. The convert phase is
+            # usually the throughput ceiling.
+            proc = self._config.processing
+            uses_docling_serve = (
+                proc.converter == "docling-serve" or proc.chunker == "docling-serve"
+            )
+            if uses_docling_serve:
+                logger.info(
+                    "Ingester running: %d worker(s), %d source(s), "
+                    "%d docling-serve instance(s)",
+                    ingester_cfg.workers.worker_count,
+                    len(ingester_cfg.sources),
+                    len(self._config.providers.docling_serve.base_urls),
+                )
+            else:
+                logger.info(
+                    "Ingester running: %d worker(s), %d source(s)",
+                    ingester_cfg.workers.worker_count,
+                    len(ingester_cfg.sources),
+                )
+
+            api_task, api_server = await self._maybe_start_api(api)
+
+            stop_event = asyncio.Event()
+            loop = asyncio.get_running_loop()
+            for sig in (signal.SIGINT, signal.SIGTERM):
+                try:
+                    loop.add_signal_handler(sig, stop_event.set)
+                except NotImplementedError:
+                    # Windows; signal handlers unavailable in asyncio.
+                    pass
+
+            try:
+                await stop_event.wait()
+            finally:
+                logger.info("Shutting down ingester")
+                if api_server is not None:
+                    api_server.should_exit = True
+                if api_task is not None:
+                    await asyncio.gather(api_task, return_exceptions=True)
+                await self._pollers.stop()
+                await self._stop_pool()
+
+    async def run_batch(self) -> BatchReport:
+        """Run one discover() sweep across every configured source, drain the
+        queue to completion, then stop. Unlike `serve`, the periodic poller
+        loops never start — discovery is driven explicitly, so the run is
+        deterministic and exits as soon as the queue is empty."""
+        async with self._resources():
+            assert (
+                self._pollers is not None
+                and self._pool is not None
+                and self._jobs is not None
+            )
+            # A persisted queue carries terminal rows from previous runs, and
+            # a recovered URI's dead row gets pruned mid-run, so the report
+            # counts only jobs that completed at or after this start instant.
+            started_at = datetime.now(UTC)
+            await self._pool.start()
+            try:
+                await self._pollers.sweep_all()
+                while True:
+                    counts = await self._jobs.counts_by_status()
+                    if not counts.get("queued") and not counts.get("claimed"):
+                        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),
+                )
+            finally:
+                await self._stop_pool()
+
     async def _maybe_start_api(self, api: bool):
         """Spin up the FastAPI control plane on an asyncio task. Returns
         (task, server) or (None, None) when the API is disabled."""

haiku_rag_slim/haiku/rag/ingester/cli.py

@@ -1,7 +1,5 @@
 import asyncio
 import sys
-import uuid
-from datetime import UTC, datetime
 from pathlib import Path
 
 import typer
@@ -9,7 +7,6 @@ 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,
@@ -18,10 +15,7 @@ from haiku.rag.config import (  # noqa: E402
     set_config,
 )
 from haiku.rag.ingester.app import IngesterApp  # noqa: E402
-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
 from haiku.rag.logging import configure_cli_logging  # noqa: E402
 from haiku.rag.store.exceptions import (  # noqa: E402
     MigrationRequiredError,
@@ -163,59 +157,24 @@ def serve(
     asyncio.run(app.serve(api=not no_api))
 
 
-@_cli.command("run-once")
+@_cli.command("run-batch")
-def run_once(
+def run_batch(
-    uri: str = typer.Argument(..., help="URI to ingest (file://, http(s)://, s3://)."),
     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.
+    """Run one discover sweep across every configured source, drain the queue,
-
+    then exit. New and changed resources are ingested, resources that vanished
-    Bypasses the queue — does NOT enqueue. Useful for smoke-testing the
+    from a source are deleted. Exits non-zero if any job dead-letters."""
-    Source adapter + pipeline path without spinning up the full pool.
+    asyncio.run(_run_batch(get_config(), db))
-    """
-    asyncio.run(_run_once(get_config(), uri, db, delete))
 
 
-async def _run_once(
+async def _run_batch(app_config: AppConfig, db_path: Path | None) -> None:
-    app_config: AppConfig, uri: str, db_path: Path | None, delete: bool
-) -> None:
     db = _resolve_db_path(app_config, db_path)
-    now = datetime.now(UTC)
+    app = IngesterApp(config=app_config, db_path=db)
-    job = Job(
+    report = await app.run_batch()
-        id=f"adhoc-{uuid.uuid4()}",
+    typer.echo(f"Batch complete: {report.succeeded} succeeded, {report.dead} dead")
-        source_id="adhoc",
+    if report.dead:
-        uri=uri,
+        raise typer.Exit(1)
-        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}"
-        )

haiku_rag_slim/haiku/rag/ingester/pollers/manager.py

@@ -77,6 +77,13 @@ class PollerManager:
             poller._stop.clear()
             self._tasks.append(asyncio.create_task(poller.run()))
 
+    async def sweep_all(self) -> None:
+        """Run one discover() sweep on every poller, sequentially. Used by
+        one-shot batch runs that drive discovery explicitly rather than
+        through the periodic loop."""
+        for poller in self._pollers:
+            await poller._sweep_once()
+
     async def stop(self) -> None:
         for poller in self._pollers:
             await poller.stop()

haiku_rag_slim/haiku/rag/ingester/queue/repository.py

@@ -294,6 +294,20 @@ class JobRepo:
                 rows = await cursor.fetchall()
         return {row["status"]: row["n"] for row in rows}
 
+    async def counts_by_status_since(self, since: datetime) -> dict[str, int]:
+        """status -> count of jobs that reached a terminal state at or after
+        `since` (by completed_at). Only succeeded/dead set completed_at, so
+        those are the only keys returned. Lets a one-shot batch report the
+        work it finished, independent of terminal rows from earlier runs."""
+        async with self._lock:
+            async with self._conn.execute(
+                "SELECT status, COUNT(*) AS n FROM jobs "
+                "WHERE completed_at >= ? GROUP BY status",
+                (since.isoformat(),),
+            ) as cursor:
+                rows = await cursor.fetchall()
+        return {row["status"]: row["n"] for row in rows}
+
     async def count_succeeded_since(self, seconds: int) -> int:
         """How many jobs reached `succeeded` in the last `seconds` seconds.
         Drives the dashboard's rolling-throughput chips."""

haiku_rag_slim/haiku/rag/ingester/workers/pool.py

@@ -114,7 +114,7 @@ class WorkerPool:
 
     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()`."""
+        coroutine. Used by tests; not by `start()`."""
         processed = 0
         while True:
             job = await self._jobs.claim_next(worker_id)

tests/ingester/test_run_batch.py

@@ -0,0 +1,205 @@
+"""IngesterApp.run_batch: one discover sweep across configured sources,
+drain the queue, then exit. The document store (HaikuRAG) is patched out —
+the behavior under test is the sweep -> queue -> worker -> drain
+orchestration and orphan pruning, not embedding."""
+
+import asyncio
+from contextlib import asynccontextmanager
+from pathlib import Path
+from unittest.mock import AsyncMock
+from urllib.parse import unquote, urlparse
+
+import pytest
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.client.exceptions import UnsupportedSourceError
+from haiku.rag.config import (
+    APIConfig,
+    AppConfig,
+    FSSourceConfig,
+    IngesterConfig,
+    QueueConfig,
+    RetryPolicyConfig,
+    WorkerConfig,
+)
+from haiku.rag.ingester.app import IngesterApp
+from haiku.rag.store.models.document import Document
+
+
+def _config(tmp_path, **worker_kwargs) -> AppConfig:
+    workers = WorkerConfig(
+        worker_count=2,
+        poll_idle_interval_s=0.05,
+        retry=RetryPolicyConfig(max_attempts=1),
+        **worker_kwargs,
+    )
+    return AppConfig(
+        ingester=IngesterConfig(
+            queue=QueueConfig(path=tmp_path / "queue.db"),
+            sources=[
+                FSSourceConfig(
+                    type="fs",
+                    id="local",
+                    root=tmp_path,
+                    poll_interval_s=3600.0,
+                )
+            ],
+            workers=workers,
+            api=APIConfig(enabled=False),
+        ),
+    )
+
+
+def _mock_client() -> AsyncMock:
+    """A HaikuRAG mock whose upserts return a fresh Document (with an md5 so
+    the worker records sync_state) and whose lookups return a Document with a
+    deterministic id so the DELETE path resolves."""
+    client = AsyncMock(spec=HaikuRAG)
+    counter = {"n": 0}
+
+    async def _create(uri, *_, **__):
+        counter["n"] += 1
+        # Mirror the real client: persist the FS revision (mtime_ns) so the
+        # poller's change-detection skips unchanged files on the next sweep.
+        path = Path(unquote(urlparse(uri).path))
+        return Document(
+            id=f"doc-{counter['n']}",
+            content="x",
+            uri=uri,
+            metadata={
+                "content_type": "text/markdown",
+                "md5": f"md5-{counter['n']}",
+                "source_revision": str(path.stat().st_mtime_ns),
+            },
+        )
+
+    async def _get_by_uri(uri):
+        return Document(id=f"doc-for-{uri}", content="x", uri=uri, metadata={})
+
+    client.create_document_from_source.side_effect = _create
+    client.get_document_by_uri.side_effect = _get_by_uri
+    return client
+
+
+@pytest.fixture
+def use_client(monkeypatch):
+    """Make IngesterApp's internally-created HaikuRAG resolve to the given
+    mock client. Call again to swap the client for a later run in the same
+    test."""
+
+    def _install(client):
+        @asynccontextmanager
+        async def _cm(*_, **__):
+            yield client
+
+        monkeypatch.setattr("haiku.rag.client.HaikuRAG", lambda *a, **k: _cm())
+
+    return _install
+
+
+@pytest.mark.asyncio
+async def test_run_batch_drains_upserts(tmp_path, use_client):
+    (tmp_path / "a.md").write_text("hello")
+    (tmp_path / "b.md").write_text("world")
+
+    client = _mock_client()
+    use_client(client)
+
+    report = await IngesterApp(
+        config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
+    ).run_batch()
+
+    assert report.succeeded == 2
+    assert report.dead == 0
+    assert client.create_document_from_source.await_count == 2
+    ingested = {
+        call.args[0] for call in client.create_document_from_source.await_args_list
+    }
+    assert ingested == {
+        (tmp_path / "a.md").as_uri(),
+        (tmp_path / "b.md").as_uri(),
+    }
+
+
+@pytest.mark.asyncio
+async def test_run_batch_prunes_orphans(tmp_path, use_client):
+    (tmp_path / "a.md").write_text("hello")
+    (tmp_path / "b.md").write_text("world")
+
+    client = _mock_client()
+    use_client(client)
+    config = _config(tmp_path)
+    db_path = tmp_path / "db.lancedb"
+
+    # First batch ingests both files and records sync_state for each.
+    first = await IngesterApp(config=config, db_path=db_path).run_batch()
+    assert first.succeeded == 2
+    client.delete_document.assert_not_awaited()
+
+    # b.md disappears from the source tree. The next sweep sees it in
+    # sync_state but not on disk -> enqueues a DELETE for it.
+    (tmp_path / "b.md").unlink()
+
+    second = await IngesterApp(config=config, db_path=db_path).run_batch()
+
+    # a.md is unchanged (same mtime) so it's not re-ingested; only the orphan
+    # delete runs.
+    assert client.create_document_from_source.await_count == 2
+    client.delete_document.assert_awaited_once()
+    assert second.succeeded == 1
+    assert second.dead == 0
+
+
+@pytest.mark.asyncio
+async def test_run_batch_reports_dead_on_permanent_failure(tmp_path, use_client):
+    (tmp_path / "a.md").write_text("hello")
+
+    client = _mock_client()
+    client.create_document_from_source.side_effect = UnsupportedSourceError("nope")
+    use_client(client)
+
+    report = await IngesterApp(
+        config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
+    ).run_batch()
+
+    assert report.succeeded == 0
+    assert report.dead == 1
+
+
+@pytest.mark.asyncio
+async def test_run_batch_recovered_doc_is_not_counted_as_dead(tmp_path, use_client):
+    """A doc that dead-lettered in an earlier run and succeeds in this one
+    prunes its dead row. The report counts only this run's terminal jobs, so
+    the recovery shows as a success, never a negative or stale dead count."""
+    (tmp_path / "a.md").write_text("hello")
+    config = _config(tmp_path)
+    db_path = tmp_path / "db.lancedb"
+
+    failing = _mock_client()
+    failing.create_document_from_source.side_effect = UnsupportedSourceError("nope")
+    use_client(failing)
+    first = await IngesterApp(config=config, db_path=db_path).run_batch()
+    assert first.dead == 1
+
+    healthy = _mock_client()
+    use_client(healthy)
+    second = await IngesterApp(config=config, db_path=db_path).run_batch()
+    assert second.dead == 0
+    assert second.succeeded == 1
+
+
+@pytest.mark.asyncio
+async def test_run_batch_empty_source_returns_immediately(tmp_path, use_client):
+    client = _mock_client()
+    use_client(client)
+
+    report = await asyncio.wait_for(
+        IngesterApp(
+            config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
+        ).run_batch(),
+        timeout=5.0,
+    )
+
+    assert report.succeeded == 0
+    assert report.dead == 0
+    client.create_document_from_source.assert_not_awaited()