| """Batch queue worker.""" |
|
|
| from __future__ import annotations |
|
|
| import asyncio |
| import logging |
|
|
| from sqlalchemy import select |
|
|
| from app.db import get_session_factory |
| from app.models import Batch, Document |
| from app.schemas import DocStatus |
| from app.services.pipeline import process_document |
| from app.timeutil import utc_now |
|
|
| logger = logging.getLogger(__name__) |
|
|
| _running: set[str] = set() |
|
|
|
|
| async def run_batch(batch_id: str) -> None: |
| """Process all queued documents in a batch.""" |
| if batch_id in _running: |
| return |
| _running.add(batch_id) |
| try: |
| sessions = get_session_factory() |
| async with sessions() as session: |
| batch = await session.get(Batch, batch_id) |
| if not batch: |
| return |
| batch.status = "processing" |
| await session.commit() |
|
|
| result = await session.execute( |
| select(Document).where( |
| Document.batch_id == batch_id, |
| Document.status == DocStatus.QUEUED.value, |
| ) |
| ) |
| docs = list(result.scalars().all()) |
|
|
| for doc in docs: |
| await process_document(session, doc) |
| batch = await session.get(Batch, batch_id) |
| if not batch: |
| return |
| batch.processed += 1 |
| if doc.status == DocStatus.FAILED.value: |
| batch.failed += 1 |
| await session.commit() |
|
|
| batch = await session.get(Batch, batch_id) |
| if batch: |
| batch.status = "completed" |
| batch.finished_at = utc_now() |
| await session.commit() |
| except Exception as exc: |
| logger.exception("batch %s failed: %s", batch_id, exc) |
| try: |
| sessions = get_session_factory() |
| except RuntimeError: |
| return |
| async with sessions() as session: |
| batch = await session.get(Batch, batch_id) |
| if batch: |
| batch.status = "failed" |
| batch.finished_at = utc_now() |
| await session.commit() |
| finally: |
| _running.discard(batch_id) |
|
|
|
|
| def enqueue_batch(batch_id: str) -> None: |
| """Fire-and-forget batch processing task.""" |
| asyncio.create_task(run_batch(batch_id)) |
|
|
