diff --git a/haiku_rag_slim/haiku/rag/ingester/cli.py b/haiku_rag_slim/haiku/rag/ingester/cli.py index 059b3e08..96b961bf 100644 --- a/haiku_rag_slim/haiku/rag/ingester/cli.py +++ b/haiku_rag_slim/haiku/rag/ingester/cli.py @@ -22,6 +22,7 @@ from haiku.rag.ingester.exceptions import PermanentError, TransientError # noqa 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, ReadOnlyError, @@ -37,14 +38,15 @@ _cli = typer.Typer( def _configure_logfire() -> None: """Logfire emits spans only when LOGFIRE_TOKEN is set; otherwise it - stays silent (no warning either way). Matches the haiku-rag CLI.""" + stays silent. Console output is disabled in either case so span lines + don't interleave with the ingester's own RichHandler logs — telemetry + lives in the Logfire UI.""" try: import logfire - is_production = get_config().environment != "development" logfire.configure( send_to_logfire="if-token-present", - console=False if is_production else None, + console=False, ) logfire.instrument_pydantic_ai() except Exception: # pragma: no cover @@ -53,6 +55,7 @@ def _configure_logfire() -> None: def cli() -> None: """Entry point that translates store-state errors into a clean exit.""" + configure_cli_logging() _configure_logfire() try: _cli() diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/base.py b/haiku_rag_slim/haiku/rag/ingester/pollers/base.py index 6d371998..5c39dd52 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/base.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/base.py @@ -84,10 +84,24 @@ class BasePoller: return False try: snapshot = await self._sync.get_snapshot(self.source_id) + counts = { + SourceEventKind.UPSERT: 0, + SourceEventKind.DELETE: 0, + SourceEventKind.UNCHANGED: 0, + } async for event in self.source.discover(since=snapshot): + counts[event.kind] += 1 await self._handle_event(event) self._breaker.record_success() self._last_polled_at = datetime.now(UTC) + if counts[SourceEventKind.UPSERT] or counts[SourceEventKind.DELETE]: + logger.info( + "Swept %s: %d upsert, %d delete, %d unchanged", + self.source_id, + counts[SourceEventKind.UPSERT], + counts[SourceEventKind.DELETE], + counts[SourceEventKind.UNCHANGED], + ) return True except Exception as exc: self._breaker.record_failure() diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/migrations.py b/haiku_rag_slim/haiku/rag/ingester/queue/migrations.py index b7ff944c..87be765a 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/migrations.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/migrations.py @@ -7,33 +7,39 @@ from haiku.rag.ingester.queue.schema import ALL_DDL, SCHEMA_VERSION __all__ = ["SCHEMA_VERSION", "apply_migrations", "open_queue"] +async def _exec(conn: aiosqlite.Connection, sql: str, *params) -> None: + """Execute a statement and finalize the cursor. aiosqlite cursors stay + attached until closed, blocking subsequent commits with 'SQL statements + in progress'.""" + async with conn.execute(sql, params): + pass + + async def apply_migrations(conn: aiosqlite.Connection) -> int: """Idempotently create tables/indexes/views and pin schema_version. Returns the schema version after the call. Safe to call on a fresh DB or on one already at the latest version. """ - await conn.execute("PRAGMA journal_mode=WAL") - await conn.execute("PRAGMA synchronous=NORMAL") - await conn.execute("PRAGMA foreign_keys=ON") + await _exec(conn, "PRAGMA journal_mode=WAL") + await _exec(conn, "PRAGMA synchronous=NORMAL") + await _exec(conn, "PRAGMA foreign_keys=ON") for stmt in ALL_DDL: - await conn.execute(stmt) + await _exec(conn, stmt) - cursor = await conn.execute("SELECT version FROM schema_version LIMIT 1") - row = await cursor.fetchone() + async with conn.execute("SELECT version FROM schema_version LIMIT 1") as cursor: + row = await cursor.fetchone() if row is None: - await conn.execute( - "INSERT INTO schema_version (version) VALUES (?)", (SCHEMA_VERSION,) + await _exec( + conn, "INSERT INTO schema_version (version) VALUES (?)", SCHEMA_VERSION ) else: current = row[0] if current < SCHEMA_VERSION: # No diff migrations exist yet — future versions add UPDATE/ALTER # statements between here and the version bump. - await conn.execute( - "UPDATE schema_version SET version = ?", (SCHEMA_VERSION,) - ) + await _exec(conn, "UPDATE schema_version SET version = ?", SCHEMA_VERSION) await conn.commit() return SCHEMA_VERSION diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index 25aead7e..61eaa8da 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py @@ -1,3 +1,4 @@ +import asyncio import json import uuid from datetime import UTC, datetime, timedelta @@ -53,6 +54,12 @@ class JobRepo: # Row access by name in helpers below. conn.row_factory = aiosqlite.Row self._conn = conn + # Serialize repo calls on the shared connection so cursors from one + # coroutine don't sit "in progress" when another tries to commit. + # aiosqlite executes statements on a single worker thread, but + # individual cursors don't finalize until closed or GC'd — SQLite + # then refuses commit() with "SQL statements in progress". + self._lock = asyncio.Lock() async def enqueue( self, @@ -71,126 +78,139 @@ class JobRepo: job_id = str(uuid.uuid4()) now = _utcnow_iso() extra_json = json.dumps(extra) if extra is not None else None - cursor = await self._conn.execute( - """ - INSERT INTO jobs ( - id, source_id, uri, op, content_hash, revision, status, - attempts, max_attempts, last_error, extra, - enqueued_at, scheduled_at - ) VALUES (?, ?, ?, ?, ?, ?, 'queued', 0, ?, NULL, ?, ?, ?) - ON CONFLICT DO NOTHING - RETURNING * - """, - ( - job_id, - source_id, - uri, - op.value, - content_hash, - revision, - max_attempts, - extra_json, - now, - now, - ), - ) - row = await cursor.fetchone() - await self._conn.commit() + async with self._lock: + async with self._conn.execute( + """ + INSERT INTO jobs ( + id, source_id, uri, op, content_hash, revision, status, + attempts, max_attempts, last_error, extra, + enqueued_at, scheduled_at + ) VALUES (?, ?, ?, ?, ?, ?, 'queued', 0, ?, NULL, ?, ?, ?) + ON CONFLICT DO NOTHING + RETURNING * + """, + ( + job_id, + source_id, + uri, + op.value, + content_hash, + revision, + max_attempts, + extra_json, + now, + now, + ), + ) as cursor: + row = await cursor.fetchone() + await self._conn.commit() return _row_to_job(row) if row else None async def claim_next(self, worker_id: str) -> Job | None: """Atomically claim the oldest queued job whose scheduled_at <= now. Implemented as a single UPDATE ... RETURNING — no SELECT/UPDATE race.""" now = _utcnow_iso() - cursor = await self._conn.execute( - """ - UPDATE jobs - SET status = 'claimed', - claimed_at = ?, - claimed_by = ?, - attempts = attempts + 1 - WHERE id = ( - SELECT id FROM jobs - WHERE status = 'queued' AND scheduled_at <= ? - ORDER BY scheduled_at - LIMIT 1 - ) - RETURNING * - """, - (now, worker_id, now), - ) - row = await cursor.fetchone() - await self._conn.commit() + async with self._lock: + async with self._conn.execute( + """ + UPDATE jobs + SET status = 'claimed', + claimed_at = ?, + claimed_by = ?, + attempts = attempts + 1 + WHERE id = ( + SELECT id FROM jobs + WHERE status = 'queued' AND scheduled_at <= ? + ORDER BY scheduled_at + LIMIT 1 + ) + RETURNING * + """, + (now, worker_id, now), + ) as cursor: + row = await cursor.fetchone() + await self._conn.commit() return _row_to_job(row) if row else None async def get_job(self, job_id: str) -> Job | None: - cursor = await self._conn.execute("SELECT * FROM jobs WHERE id = ?", (job_id,)) - row = await cursor.fetchone() + async with self._lock: + async with self._conn.execute( + "SELECT * FROM jobs WHERE id = ?", (job_id,) + ) as cursor: + row = await cursor.fetchone() return _row_to_job(row) if row else None async def mark_succeeded(self, job_id: str) -> None: - await self._conn.execute( - "UPDATE jobs SET status='succeeded', completed_at=? WHERE id=?", - (_utcnow_iso(), job_id), - ) - await self._conn.commit() + async with self._lock: + async with self._conn.execute( + "UPDATE jobs SET status='succeeded', completed_at=? WHERE id=?", + (_utcnow_iso(), job_id), + ): + pass + await self._conn.commit() async def mark_dead(self, job_id: str, error: str) -> None: - await self._conn.execute( - "UPDATE jobs SET status='dead', completed_at=?, last_error=? WHERE id=?", - (_utcnow_iso(), error, job_id), - ) - await self._conn.commit() + async with self._lock: + async with self._conn.execute( + "UPDATE jobs SET status='dead', completed_at=?, last_error=? WHERE id=?", + (_utcnow_iso(), error, job_id), + ): + pass + await self._conn.commit() async def reschedule(self, job_id: str, delay_seconds: float, error: str) -> None: scheduled = (datetime.now(UTC) + timedelta(seconds=delay_seconds)).isoformat() - await self._conn.execute( - """ - UPDATE jobs - SET status='queued', - scheduled_at=?, - claimed_at=NULL, - claimed_by=NULL, - last_error=? - WHERE id=? - """, - (scheduled, error, job_id), - ) - await self._conn.commit() + async with self._lock: + async with self._conn.execute( + """ + UPDATE jobs + SET status='queued', + scheduled_at=?, + claimed_at=NULL, + claimed_by=NULL, + last_error=? + WHERE id=? + """, + (scheduled, error, job_id), + ): + pass + await self._conn.commit() async def retry(self, job_id: str) -> Job: """Rescue a dead job: status='queued', attempts=0, error cleared. Raises KeyError if the job doesn't exist.""" now = _utcnow_iso() - cursor = await self._conn.execute( - """ - UPDATE jobs - SET status='queued', - attempts=0, - last_error=NULL, - claimed_at=NULL, - claimed_by=NULL, - completed_at=NULL, - scheduled_at=? - WHERE id=? - RETURNING * - """, - (now, job_id), - ) - row = await cursor.fetchone() - await self._conn.commit() + async with self._lock: + async with self._conn.execute( + """ + UPDATE jobs + SET status='queued', + attempts=0, + last_error=NULL, + claimed_at=NULL, + claimed_by=NULL, + completed_at=NULL, + scheduled_at=? + WHERE id=? + RETURNING * + """, + (now, job_id), + ) as cursor: + row = await cursor.fetchone() + await self._conn.commit() if not row: raise KeyError(f"Job {job_id!r} not found") return _row_to_job(row) async def cancel(self, job_id: str) -> bool: """Delete a queued or claimed job. Returns True if a row was removed.""" - cursor = await self._conn.execute( - "DELETE FROM jobs WHERE id=? AND status IN ('queued', 'claimed') RETURNING id", - (job_id,), - ) - row = await cursor.fetchone() - await self._conn.commit() + async with self._lock: + async with self._conn.execute( + "DELETE FROM jobs WHERE id=? AND status IN ('queued', 'claimed') RETURNING id", + (job_id,), + ) as cursor: + row = await cursor.fetchone() + await self._conn.commit() return row is not None async def list_jobs( @@ -215,18 +235,20 @@ class JobRepo: params.append(uri) where = ("WHERE " + " AND ".join(clauses)) if clauses else "" params.extend([limit, offset]) - cursor = await self._conn.execute( - f"SELECT * FROM jobs {where} ORDER BY enqueued_at DESC LIMIT ? OFFSET ?", - params, - ) - rows = await cursor.fetchall() + async with self._lock: + async with self._conn.execute( + f"SELECT * FROM jobs {where} ORDER BY enqueued_at DESC LIMIT ? OFFSET ?", + params, + ) as cursor: + rows = await cursor.fetchall() return [_row_to_job(r) for r in rows] async def counts_by_status(self) -> dict[str, int]: - cursor = await self._conn.execute( - "SELECT status, COUNT(*) AS n FROM jobs GROUP BY status" - ) - rows = await cursor.fetchall() + async with self._lock: + async with self._conn.execute( + "SELECT status, COUNT(*) AS n FROM jobs GROUP BY status" + ) as cursor: + rows = await cursor.fetchall() return {row["status"]: row["n"] for row in rows} async def reap_stale(self, claim_timeout_seconds: int) -> int: @@ -235,39 +257,46 @@ class JobRepo: threshold = ( datetime.now(UTC) - timedelta(seconds=claim_timeout_seconds) ).isoformat() - cursor = await self._conn.execute( - """ - UPDATE jobs - SET status='queued', claimed_at=NULL, claimed_by=NULL - WHERE status='claimed' AND claimed_at < ? - """, - (threshold,), - ) - await self._conn.commit() - return cursor.rowcount or 0 + async with self._lock: + cursor = await self._conn.execute( + """ + UPDATE jobs + SET status='queued', claimed_at=NULL, claimed_by=NULL + WHERE status='claimed' AND claimed_at < ? + """, + (threshold,), + ) + rowcount = cursor.rowcount or 0 + await cursor.close() + await self._conn.commit() + return rowcount class SyncStateRepo: def __init__(self, conn: aiosqlite.Connection): conn.row_factory = aiosqlite.Row self._conn = conn + # See JobRepo for why we serialize on the shared connection. + self._lock = asyncio.Lock() async def get_snapshot(self, source_id: str) -> dict[str, str]: """uri -> revision map for the source. Drops rows where revision is NULL (the poller can't compare against an absent revision).""" - cursor = await self._conn.execute( - "SELECT uri, revision FROM sync_state WHERE source_id=? AND revision IS NOT NULL", - (source_id,), - ) - rows = await cursor.fetchall() + async with self._lock: + async with self._conn.execute( + "SELECT uri, revision FROM sync_state WHERE source_id=? AND revision IS NOT NULL", + (source_id,), + ) as cursor: + rows = await cursor.fetchall() return {row["uri"]: row["revision"] for row in rows} async def get_row(self, source_id: str, uri: str) -> SyncStateRow | None: - cursor = await self._conn.execute( - "SELECT * FROM sync_state WHERE source_id=? AND uri=?", - (source_id, uri), - ) - row = await cursor.fetchone() + async with self._lock: + async with self._conn.execute( + "SELECT * FROM sync_state WHERE source_id=? AND uri=?", + (source_id, uri), + ) as cursor: + row = await cursor.fetchone() return _row_to_sync_state(row) if row else None async def upsert( @@ -283,25 +312,29 @@ class SyncStateRepo: last_ingested_at; otherwise only last_seen_at is bumped.""" now = _utcnow_iso() ingested_at = now if ingested else None - await self._conn.execute( - """ - INSERT INTO sync_state ( - source_id, uri, revision, content_hash, last_seen_at, last_ingested_at - ) - VALUES (?, ?, ?, ?, ?, ?) - ON CONFLICT(source_id, uri) DO UPDATE SET - revision = excluded.revision, - content_hash = excluded.content_hash, - last_seen_at = excluded.last_seen_at, - last_ingested_at = COALESCE(excluded.last_ingested_at, last_ingested_at) - """, - (source_id, uri, revision, content_hash, now, ingested_at), - ) - await self._conn.commit() + async with self._lock: + async with self._conn.execute( + """ + INSERT INTO sync_state ( + source_id, uri, revision, content_hash, last_seen_at, last_ingested_at + ) + VALUES (?, ?, ?, ?, ?, ?) + ON CONFLICT(source_id, uri) DO UPDATE SET + revision = excluded.revision, + content_hash = excluded.content_hash, + last_seen_at = excluded.last_seen_at, + last_ingested_at = COALESCE(excluded.last_ingested_at, last_ingested_at) + """, + (source_id, uri, revision, content_hash, now, ingested_at), + ): + pass + await self._conn.commit() async def delete(self, source_id: str, uri: str) -> None: - await self._conn.execute( - "DELETE FROM sync_state WHERE source_id=? AND uri=?", - (source_id, uri), - ) - await self._conn.commit() + async with self._lock: + async with self._conn.execute( + "DELETE FROM sync_state WHERE source_id=? AND uri=?", + (source_id, uri), + ): + pass + await self._conn.commit() diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index c9d79b87..f7fc5a52 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -1,5 +1,6 @@ import asyncio import logging +import time from typing import TYPE_CHECKING from haiku.rag.ingester.exceptions import PermanentError, TransientError @@ -113,6 +114,8 @@ class WorkerPool: pass async def _process(self, job: Job) -> None: + started = time.monotonic() + logger.info("Processing %s %s (job %s)", job.op.value, job.uri, job.id) try: result = await run_job(self._client, job) except PermanentError as e: @@ -154,3 +157,6 @@ class WorkerPool: content_hash=result.content_hash, ingested=True, ) + logger.info( + "Job %s succeeded in %.2fs: %s", job.id, time.monotonic() - started, job.uri + )