Add run-batch manifest replay
This commit is contained in:
parent
a3b542dba8
commit
7251d104c4
6 changed files with 449 additions and 30 deletions
|
|
@ -9,7 +9,7 @@ from typing import TYPE_CHECKING
|
|||
from pydantic import BaseModel
|
||||
|
||||
from haiku.rag.config import AppConfig
|
||||
from haiku.rag.ingester.batch import BatchDryRunReport
|
||||
from haiku.rag.ingester.batch import BatchDryRunReport, BatchManifest
|
||||
from haiku.rag.ingester.metadata import build_providers, load_metadata_providers
|
||||
from haiku.rag.ingester.pollers.manager import PollerManager
|
||||
from haiku.rag.ingester.queue.migrations import open_queue
|
||||
|
|
@ -179,6 +179,28 @@ class IngesterApp:
|
|||
if landed:
|
||||
logger.info("Drained %d cancel-cleanup release(s) before close", landed)
|
||||
|
||||
async def _drain_batch(self, started_at: datetime) -> BatchReport:
|
||||
assert self._pool is not None and self._jobs is not None
|
||||
while True:
|
||||
counts = await self._jobs.counts_by_status()
|
||||
if not counts.get("queued") and not counts.get("claimed"):
|
||||
break
|
||||
if self._pool.live_workers == 0:
|
||||
outstanding = counts.get("queued", 0) + counts.get("claimed", 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:
|
||||
"""Run pollers + workers (and the HTTP API when enabled) until a
|
||||
SIGINT/SIGTERM is received. Drains in-flight work on shutdown."""
|
||||
|
|
@ -250,26 +272,9 @@ class IngesterApp:
|
|||
await self._pool.start()
|
||||
try:
|
||||
failed_sweeps = 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
|
||||
if self._pool.live_workers == 0:
|
||||
outstanding = counts.get("queued", 0) + counts.get("claimed", 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),
|
||||
failed_sweeps=failed_sweeps,
|
||||
)
|
||||
report = await self._drain_batch(started_at)
|
||||
report.failed_sweeps = failed_sweeps
|
||||
return report
|
||||
finally:
|
||||
await self._stop_pool()
|
||||
await self._pollers.close_sources()
|
||||
|
|
@ -282,6 +287,91 @@ class IngesterApp:
|
|||
manifest, failed_sweeps = await self._pollers.dry_run_manifest()
|
||||
return BatchDryRunReport(manifest=manifest, failed_sweeps=failed_sweeps)
|
||||
|
||||
async def run_batch_from_manifest(self, manifest: BatchManifest) -> BatchReport:
|
||||
"""Enqueue and drain a dry-run manifest without running a fresh
|
||||
discovery sweep."""
|
||||
if manifest.version != 1:
|
||||
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:
|
||||
await self._pollers.close_sources()
|
||||
raise ValueError(
|
||||
"Manifest references unconfigured source(s): " + ", ".join(missing)
|
||||
)
|
||||
|
||||
pending = [
|
||||
source_id
|
||||
for source_id in sorted(manifest_sources)
|
||||
if await self._jobs.has_pending(source_id)
|
||||
]
|
||||
if pending:
|
||||
await self._pollers.close_sources()
|
||||
raise ValueError(
|
||||
"Cannot replay manifest while source(s) have pending work: "
|
||||
+ ", ".join(pending)
|
||||
)
|
||||
|
||||
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}")
|
||||
|
||||
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:
|
||||
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": {
|
||||
"version": manifest.version,
|
||||
"generated_at": manifest.generated_at.isoformat(),
|
||||
"discovered_at": change.discovered_at.isoformat(),
|
||||
}
|
||||
},
|
||||
)
|
||||
if job is None:
|
||||
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)
|
||||
finally:
|
||||
await self._stop_pool()
|
||||
await self._pollers.close_sources()
|
||||
|
||||
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."""
|
||||
|
|
|
|||
|
|
@ -150,6 +150,11 @@ def _write_manifest(manifest: BatchManifest, path: Path) -> None:
|
|||
path.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8")
|
||||
|
||||
|
||||
def _load_manifest(path: Path) -> BatchManifest:
|
||||
data = yaml.safe_load(path.read_text(encoding="utf-8"))
|
||||
return BatchManifest.model_validate(data)
|
||||
|
||||
|
||||
@_cli.command("serve")
|
||||
def serve(
|
||||
db: Path | None = typer.Option(
|
||||
|
|
@ -212,12 +217,29 @@ def run_batch(
|
|||
"-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:
|
||||
"""Run one discover sweep across every configured source, drain the queue,
|
||||
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
|
||||
source's sweep does not complete."""
|
||||
asyncio.run(_run_batch(get_config(), db, dry_run=dry_run, output=output))
|
||||
if manifest is not None and dry_run:
|
||||
raise typer.BadParameter("--manifest cannot be combined with --dry-run")
|
||||
if manifest is not None and output is not None:
|
||||
raise typer.BadParameter("--output is only valid with --dry-run")
|
||||
asyncio.run(
|
||||
_run_batch(
|
||||
get_config(),
|
||||
db,
|
||||
dry_run=dry_run,
|
||||
output=output,
|
||||
manifest_path=manifest,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _run_batch(
|
||||
|
|
@ -226,6 +248,7 @@ async def _run_batch(
|
|||
*,
|
||||
dry_run: bool = False,
|
||||
output: Path | None = None,
|
||||
manifest_path: Path | None = None,
|
||||
) -> None:
|
||||
db = _resolve_db_path(app_config, db_path)
|
||||
app = IngesterApp(config=app_config, db_path=db)
|
||||
|
|
@ -248,6 +271,20 @@ async def _run_batch(
|
|||
)
|
||||
return
|
||||
|
||||
if manifest_path is not None:
|
||||
try:
|
||||
manifest = _load_manifest(manifest_path)
|
||||
report = await app.run_batch_from_manifest(manifest)
|
||||
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
|
||||
|
||||
report = await app.run_batch()
|
||||
typer.echo(f"Batch complete: {report.succeeded} succeeded, {report.dead} dead")
|
||||
if report.failed_sweeps:
|
||||
|
|
|
|||
|
|
@ -110,23 +110,38 @@ async def run_job(
|
|||
),
|
||||
):
|
||||
try:
|
||||
manifest_context = extra.get("_manifest")
|
||||
if job.op is JobOp.DELETE:
|
||||
# An atomic-rename save can let a spurious DELETE win the
|
||||
# enqueue race while the file is mid-rewrite. If the resource
|
||||
# is already back, skip the delete (it would blackhole a live
|
||||
# document) and let the next sweep re-ingest it.
|
||||
try:
|
||||
source = resolve_configured_source(job.uri, job.source_id, sources)
|
||||
restored = await source.head(job.uri) is not None
|
||||
except Exception:
|
||||
restored = False
|
||||
if restored:
|
||||
return JobResult(deleted=False)
|
||||
# document) and let the next sweep re-ingest it. Manifest
|
||||
# replay intentionally follows the frozen dry-run changeset.
|
||||
if manifest_context is None:
|
||||
try:
|
||||
source = resolve_configured_source(
|
||||
job.uri, job.source_id, sources
|
||||
)
|
||||
restored = await source.head(job.uri) is not None
|
||||
except Exception:
|
||||
restored = False
|
||||
if restored:
|
||||
return JobResult(deleted=False)
|
||||
doc = await client.get_document_by_uri(job.uri)
|
||||
if doc is not None and doc.id is not None:
|
||||
await client.delete_document(doc.id)
|
||||
return JobResult(deleted=True)
|
||||
|
||||
if manifest_context is not None and job.revision is not None:
|
||||
source = resolve_configured_source(job.uri, job.source_id, sources)
|
||||
current_revision = await source.head(job.uri)
|
||||
if current_revision != job.revision:
|
||||
raise PermanentError(
|
||||
"manifest revision is stale for "
|
||||
f"{job.uri}: expected {job.revision!r}, "
|
||||
f"current {current_revision!r}"
|
||||
)
|
||||
|
||||
result = await client.create_document_from_source(
|
||||
job.uri,
|
||||
sources=sources,
|
||||
|
|
|
|||
|
|
@ -68,6 +68,17 @@ def _fake_dry_run_app(report: BatchDryRunReport, monkeypatch) -> AsyncMock:
|
|||
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):
|
||||
fake = _fake_app(BatchReport(succeeded=3, dead=0), monkeypatch)
|
||||
|
||||
|
|
@ -152,6 +163,80 @@ def test_run_batch_dry_run_exits_nonzero_when_sweep_fails(monkeypatch, tmp_path)
|
|||
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_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
|
||||
|
||||
|
||||
# --- serve ---
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ def _job(
|
|||
*,
|
||||
op: JobOp = JobOp.UPSERT,
|
||||
uri: str = "https://example.com/a.pdf",
|
||||
revision: str | None = None,
|
||||
extra: dict | None = None,
|
||||
attempts: int = 0,
|
||||
) -> Job:
|
||||
|
|
@ -25,6 +26,7 @@ def _job(
|
|||
source_id="src",
|
||||
uri=uri,
|
||||
op=op,
|
||||
revision=revision,
|
||||
status=JobStatus.CLAIMED,
|
||||
attempts=attempts,
|
||||
max_attempts=5,
|
||||
|
|
@ -255,6 +257,22 @@ async def test_delete_skipped_when_resource_restored_on_source():
|
|||
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
|
||||
async def test_delete_proceeds_when_resource_absent_on_source():
|
||||
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")
|
||||
|
||||
|
||||
@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
|
||||
async def test_delete_proceeds_when_source_unresolvable():
|
||||
"""No configured source for the job: the probe can't run, so the delete
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ pruning, not embedding."""
|
|||
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock
|
||||
from urllib.parse import unquote, urlparse
|
||||
|
|
@ -23,6 +24,7 @@ from haiku.rag.config import (
|
|||
WorkerConfig,
|
||||
)
|
||||
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.queue.migrations import open_queue
|
||||
from haiku.rag.ingester.queue.models import JobOp
|
||||
|
|
@ -86,6 +88,10 @@ def _mock_client() -> AsyncMock:
|
|||
return client
|
||||
|
||||
|
||||
def _manifest(*changes: BatchChange) -> BatchManifest:
|
||||
return BatchManifest(generated_at=datetime.now(UTC), changes=list(changes))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def use_client(monkeypatch):
|
||||
"""Make IngesterApp's internally-created HaikuRAG resolve to the given
|
||||
|
|
@ -278,6 +284,138 @@ async def test_run_batch_dry_run_reports_manifest_without_mutating_queue(tmp_pat
|
|||
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_rejects_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="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
|
||||
async def test_run_batch_aborts_when_all_workers_die(
|
||||
tmp_path, use_client, monkeypatch, caplog
|
||||
|
|
|
|||
Loading…
Reference in a new issue