mcp-nextcloud/nextcloud_mcp_server/vector/batch_ocr_store.py

"""Tracking store for in-flight async batch OCR jobs (Deck #332).

When ``DOCUMENT_OCR_MODE=batch`` the OCR tier submits a document to the gateway's
async batch route and then re-polls across procrastinate retries. procrastinate
job args are immutable, so the gateway ``job_id`` (and submit time, for the poll
deadline) live in the ``batch_ocr_jobs`` app-DB table, keyed on the document +
its content version (``etag``).

Engine reuse mirrors :class:`~nextcloud_mcp_server.usage.store.UsageEventStore`:
rather than open its own engine this store borrows the process-wide
:class:`RefreshTokenStorage` singleton (``get_shared_storage()``) — same app DB,
dialect handling, ``?``-placeholder shim, and the guarantee that Alembic
migrations (incl. ``batch_ocr_jobs``) already ran.
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass

import anyio

from nextcloud_mcp_server.auth.storage import RefreshTokenStorage, get_shared_storage

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class BatchOcrJob:
    """A tracked in-flight batch OCR job."""

    job_id: str
    status: str
    submitted_at: int


class BatchOcrJobStore:
    """CRUD for the ``batch_ocr_jobs`` table (one row per in-flight job)."""

    _shared_instance: "BatchOcrJobStore | None" = None
    _shared_lock: anyio.Lock = anyio.Lock()

    def __init__(self, storage: RefreshTokenStorage) -> None:
        self._storage = storage

    @classmethod
    async def shared(cls) -> "BatchOcrJobStore":
        """Process-wide store backed by the storage singleton (mirrors
        ``UsageEventStore.shared``). Tests should construct
        ``BatchOcrJobStore(storage)`` directly — the cache is a process global
        with no teardown hook."""
        async with cls._shared_lock:
            if cls._shared_instance is None:
                cls._shared_instance = cls(await get_shared_storage())
        return cls._shared_instance

    async def get(
        self, *, user_id: str, doc_id: str, doc_type: str, etag: str
    ) -> BatchOcrJob | None:
        """The in-flight job for this document+version, or ``None``."""
        async with self._storage.acquire() as db:
            async with db.execute(
                "SELECT job_id, status, submitted_at FROM batch_ocr_jobs "
                "WHERE user_id = ? AND doc_id = ? AND doc_type = ? AND etag = ?",
                (user_id, doc_id, doc_type, etag),
            ) as cursor:
                row = await cursor.fetchone()
        if row is None:
            return None
        return BatchOcrJob(job_id=row[0], status=row[1], submitted_at=int(row[2]))

    async def insert_pending(
        self,
        *,
        user_id: str,
        doc_id: str,
        doc_type: str,
        etag: str,
        job_id: str,
        submitted_at: int | None = None,
    ) -> None:
        """Record a freshly-submitted job. ``ON CONFLICT DO NOTHING`` makes a
        racing double-submit harmless (the first row wins; the loser's job id is
        abandoned and reaped by the gateway-side file purge)."""
        now = submitted_at if submitted_at is not None else int(time.time())
        async with self._storage.acquire() as db:
            await db.execute(
                "INSERT INTO batch_ocr_jobs "
                "(user_id, doc_id, doc_type, etag, job_id, status, submitted_at, updated_at) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?) "
                "ON CONFLICT (user_id, doc_id, doc_type, etag) DO NOTHING",
                (user_id, doc_id, doc_type, etag, job_id, "pending", now, now),
            )
            await db.commit()

    async def delete(
        self, *, user_id: str, doc_id: str, doc_type: str, etag: str
    ) -> None:
        """Drop the row once the job is terminal (succeeded or failed)."""
        async with self._storage.acquire() as db:
            await db.execute(
                "DELETE FROM batch_ocr_jobs "
                "WHERE user_id = ? AND doc_id = ? AND doc_type = ? AND etag = ?",
                (user_id, doc_id, doc_type, etag),
            )
            await db.commit()

    async def delete_stale_for_doc(
        self, *, user_id: str, doc_id: str, doc_type: str, keep_etag: str
    ) -> None:
        """Remove superseded-version rows for a document (any etag other than the
        current one) before a resubmit, so a re-edited file doesn't leave its
        old in-flight job tracked forever."""
        async with self._storage.acquire() as db:
            await db.execute(
                "DELETE FROM batch_ocr_jobs "
                "WHERE user_id = ? AND doc_id = ? AND doc_type = ? AND etag != ?",
                (user_id, doc_id, doc_type, keep_etag),
            )
            await db.commit()