src/knowledge_ingest/store.py

"""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:
        # `skipped` is deliberately NOT persisted: it is a property of the
        # request (which documents were declined and why), not of the batch's
        # progress, and the column does not exist. It survives in the 202
        # response and in the in-process job for the life of the request.
        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