| """JobStore — the seam ingestion job state is read and written through. |
| |
| Two implementations. `PostgresJobStore` is the real one and is what the service |
| uses: the `knowledge_jobs` table was created 2026-09-07 and jobs must survive a |
| restart, because a batch runs for minutes and an in-memory job lost to a redeploy |
| leaves the caller polling an id that will never resolve. `InMemoryJobStore` |
| remains for tests and for running the service without a database. |
| |
| **This store deliberately does NOT follow the never-throw convention.** The |
| seams listed in CLAUDE.md §5.4 exist so a *chat turn* degrades to soft output |
| rather than 500ing at the user. An ingestion job is the opposite case: it is |
| explicitly asynchronous, the caller polls for its outcome, and a job that |
| silently lost its rows while reporting success is strictly worse than one that |
| reports `failed`. Persistence failures here propagate and settle the job. |
| |
| ⚠️ `InMemoryJobStore` state does not survive a restart, which is why it is not |
| the default. A lost job is re-triggerable and content-addressing means the retry |
| re-spends nothing, but the caller is left polling an id that no longer exists. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import asyncio |
| from typing import Protocol, runtime_checkable |
|
|
| from sqlalchemy import select |
|
|
| from src.db.postgres.connection import AsyncSessionLocal |
| from src.db.postgres.models import KnowledgeJobRow |
| from src.knowledge_ingest.models import DocumentResult, IngestJob, SkippedDocument |
| from src.middlewares.logging import get_logger |
|
|
| logger = get_logger("knowledge_job_store") |
|
|
|
|
| @runtime_checkable |
| class JobStore(Protocol): |
| """Create, read and update ingestion jobs. |
| |
| `get` returns None for an unknown id — the caller turns that into a 404 |
| rather than the store guessing at an HTTP shape. |
| """ |
|
|
| async def create(self, job: IngestJob) -> IngestJob: ... |
|
|
| async def get(self, job_id: str) -> IngestJob | None: ... |
|
|
| async def save(self, job: IngestJob) -> None: ... |
|
|
| async def list_for_user(self, user_id: str, limit: int = 20) -> list[IngestJob]: ... |
|
|
| async def find_active(self, scope_id: str) -> IngestJob | None: ... |
|
|
|
|
| class InMemoryJobStore: |
| """Process-local job state, guarded by a lock. |
| |
| The lock matters for `find_active`: two clicks arriving together must not |
| both conclude that no job is running and start one each. Single-flight is a |
| hard requirement — two concurrent runs mean two 1.2 GB GLiNER models |
| resident at once (KNOWLEDGE_SURFACE_PLAN.md §4, N10d). |
| """ |
|
|
| def __init__(self) -> None: |
| self._jobs: dict[str, IngestJob] = {} |
| self._lock = asyncio.Lock() |
|
|
| async def create(self, job: IngestJob) -> IngestJob: |
| async with self._lock: |
| self._jobs[job.id] = job |
| logger.info( |
| "ingest_job_created", |
| extra={"job_id": job.id, "scope_id": job.scope_id, "n": job.n_requested}, |
| ) |
| return job |
|
|
| async def get(self, job_id: str) -> IngestJob | None: |
| return self._jobs.get(job_id) |
|
|
| async def save(self, job: IngestJob) -> None: |
| self._jobs[job.id] = job |
|
|
| async def list_for_user(self, user_id: str, limit: int = 20) -> list[IngestJob]: |
| jobs = [j for j in self._jobs.values() if j.user_id == user_id] |
| jobs.sort(key=lambda j: j.created_at, reverse=True) |
| return jobs[:limit] |
|
|
| async def find_active(self, scope_id: str) -> IngestJob | None: |
| """The in-flight job for a scope, if any — the single-flight check. |
| |
| Scope-wide rather than per-user on purpose: the GLiNER model and the |
| executor are process resources, so two users in one company clicking at |
| once is exactly the case that must serialise. |
| """ |
| for job in self._jobs.values(): |
| if job.scope_id == scope_id and not job.is_terminal: |
| return job |
| return None |
|
|
|
|
| class PostgresJobStore: |
| """Job state in `knowledge_jobs` - the durable implementation (N4). |
| |
| Replaces `InMemoryJobStore` now that the table exists. The distinction that |
| matters is restart survival: a batch runs for minutes, and an in-memory job |
| lost to a redeploy leaves the caller polling an id that will never resolve. |
| |
| `IngestJob` mirrors the table field for field, so the mapping here is |
| mechanical in both directions and there is no second definition of a job to |
| keep in step. |
| """ |
|
|
| async def _to_row(self, session, job: IngestJob) -> KnowledgeJobRow | None: |
| return await session.get(KnowledgeJobRow, job.id) |
|
|
| @staticmethod |
| def _from_row(row: KnowledgeJobRow) -> IngestJob: |
| """Row to model. `getattr` throughout: Harry's migration may add columns, |
| and nothing here may assume it has not (CLAUDE.md 2.2).""" |
| return IngestJob( |
| id=row.id, |
| scope_id=row.scope_id, |
| user_id=row.user_id, |
| status=row.status, |
| stage=row.stage, |
| review_mode=row.review_mode or "expert", |
| document_ids=list(row.document_ids or []), |
| n_requested=row.n_requested or 0, |
| n_succeeded=row.n_succeeded or 0, |
| n_failed=row.n_failed or 0, |
| results=[DocumentResult(**r) for r in (row.results or [])], |
| skipped=[SkippedDocument(**s) for s in (getattr(row, "skipped", None) or [])], |
| estimated_cost=row.estimated_cost, |
| error_message=row.error_message, |
| created_at=row.created_at, |
| started_at=row.started_at, |
| completed_at=row.completed_at, |
| ) |
|
|
| @staticmethod |
| def _apply(row: KnowledgeJobRow, job: IngestJob) -> None: |
| row.status = job.status |
| row.stage = job.stage |
| row.document_ids = list(job.document_ids) |
| row.n_requested = job.n_requested |
| row.n_succeeded = job.n_succeeded |
| row.n_failed = job.n_failed |
| row.results = [r.model_dump() for r in job.results] |
| row.estimated_cost = job.estimated_cost |
| row.error_message = job.error_message |
| row.started_at = job.started_at |
| row.completed_at = job.completed_at |
|
|
| async def create(self, job: IngestJob) -> IngestJob: |
| |
| |
| |
| |
| row = KnowledgeJobRow( |
| id=job.id, |
| scope_id=job.scope_id, |
| user_id=job.user_id, |
| status=job.status, |
| stage=job.stage, |
| review_mode=job.review_mode, |
| document_ids=list(job.document_ids), |
| n_requested=job.n_requested, |
| n_succeeded=job.n_succeeded, |
| n_failed=job.n_failed, |
| results=[r.model_dump() for r in job.results], |
| estimated_cost=job.estimated_cost, |
| error_message=job.error_message, |
| started_at=job.started_at, |
| completed_at=job.completed_at, |
| ) |
| async with AsyncSessionLocal() as session: |
| session.add(row) |
| await session.commit() |
| logger.info( |
| "ingest_job_created", |
| extra={"job_id": job.id, "scope_id": job.scope_id, "n": job.n_requested}, |
| ) |
| return job |
|
|
| async def get(self, job_id: str) -> IngestJob | None: |
| async with AsyncSessionLocal() as session: |
| row = await session.get(KnowledgeJobRow, job_id) |
| return self._from_row(row) if row is not None else None |
|
|
| async def save(self, job: IngestJob) -> None: |
| async with AsyncSessionLocal() as session: |
| row = await session.get(KnowledgeJobRow, job.id) |
| if row is None: |
| raise LookupError(f"no ingestion job {job.id}") |
| self._apply(row, job) |
| await session.commit() |
|
|
| async def list_for_user(self, user_id: str, limit: int = 20) -> list[IngestJob]: |
| async with AsyncSessionLocal() as session: |
| result = await session.execute( |
| select(KnowledgeJobRow) |
| .where(KnowledgeJobRow.user_id == user_id) |
| .order_by(KnowledgeJobRow.created_at.desc()) |
| .limit(limit) |
| ) |
| return [self._from_row(r) for r in result.scalars().all()] |
|
|
| async def find_active(self, scope_id: str) -> IngestJob | None: |
| """The in-flight job for a scope - the single-flight check. |
| |
| Scope-wide rather than per-user: the GLiNER model and the executor are |
| process resources, so two colleagues clicking at once is exactly the |
| case that must serialise. |
| |
| ⚠️ A job left `running` by a hard crash blocks the scope until it is |
| settled, because nothing else can distinguish it from one still working. |
| Acceptable while ingestion is admin-frequency; a lease or a heartbeat is |
| the fix if that ever stops being true. |
| """ |
| async with AsyncSessionLocal() as session: |
| result = await session.execute( |
| select(KnowledgeJobRow) |
| .where( |
| KnowledgeJobRow.scope_id == scope_id, |
| KnowledgeJobRow.status.in_(("queued", "running")), |
| ) |
| .order_by(KnowledgeJobRow.created_at.desc()) |
| .limit(1) |
| ) |
| row = result.scalars().first() |
| return self._from_row(row) if row is not None else None |
|
|
