"""Extract worker handler: download blob → extract/OCR → upload text → publish. The handler is intentionally separate from the consumer so it can be tested in-process without spinning up a real RabbitMQ consumer. It owns: - idempotency checks against documents.status - status transitions (queued → extracting) - MinIO download/upload - extraction/OCR via core.extraction - publishing DocumentExtracted to analyze.q - updating the jobs row, recording failure class, and refund-on-DLQ. """ from __future__ import annotations import asyncio import uuid from typing import TYPE_CHECKING from src.contract_check.core.billing.quota import release_document_slot from src.contract_check.core.config import get_settings from src.contract_check.core.credits import refund_credit from src.contract_check.core.db.enums import DOC_TERMINAL, FailureClass from src.contract_check.core.db.repositories import DocumentRepository, JobRepository from src.contract_check.core.extraction import ( ExtractionFailedError, UnsupportedFormatError, detect_format, extract_document, ) from src.contract_check.core.logging import get_logger from src.contract_check.core.metrics import ( extract_duration, extraction_total, mq_failed, mq_published, ) from src.contract_check.core.mq.messages import ( DocumentExtracted, DocumentUploaded, PrescreenRequested, ) from src.contract_check.core.mq.publisher import Publisher from src.contract_check.core.mq.topology import RK_ANALYZE, RK_PRESCREEN from src.contract_check.core.s3 import extracted_key from src.contract_check.core.s3.minio_storage import MinioStorage if TYPE_CHECKING: from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from src.contract_check.core.extraction import ExtractedDocument log = get_logger(__name__) class ExtractHandler: """Business logic for worker-extract.""" def __init__( self, *, session_factory: async_sessionmaker[AsyncSession], publish_routing_key: str = RK_ANALYZE, ) -> None: self._session_factory = session_factory self._publish_routing_key = publish_routing_key self._settings = get_settings() self._storage = MinioStorage.from_endpoint_url( endpoint_url=self._settings.s3_endpoint_url, access_key=self._settings.s3_access_key, secret_key=self._settings.s3_secret_key, bucket=self._settings.s3_bucket, region=self._settings.s3_region, ) self._publisher: Publisher | None = None async def _publisher_instance(self) -> Publisher: if self._publisher is None: self._publisher = Publisher(self._settings.rabbitmq_url, origin="worker-extract") await self._publisher.connect() return self._publisher async def _db_status(self, session: AsyncSession, document_id: uuid.UUID) -> str | None: return await DocumentRepository(session).get_status_for_update(document_id) async def handle(self, payload: DocumentUploaded) -> None: async with self._session_factory() as session: docs = DocumentRepository(session) jobs = JobRepository(session) current = await self._db_status(session, payload.document_id) if current is None: log.warning("document_not_found", document_id=str(payload.document_id)) return if current in DOC_TERMINAL: log.info("document_already_terminal", status=current) return await docs.update_status(payload.document_id, status="extracting", stage="downloading") await jobs.claim_start(payload.document_id, "extract") await session.commit() document = await self._extract_document(payload) extracted_text = document.markdown fmt = document.metadata.get("format", "unknown") ocr_used = document.metadata.get("is_scan", False) is_structured = document.is_structured has_tables = document.metadata.get("has_tables", False) ext_key = extracted_key(str(payload.user_id), str(payload.document_id)) text_bytes = extracted_text.encode("utf-8") await self._storage.put(ext_key, text_bytes, content_type="text/markdown; charset=utf-8") publisher = await self._publisher_instance() if self._settings.prescreen_enabled: msg = PrescreenRequested( correlation_id=payload.correlation_id, document_id=payload.document_id, user_id=payload.user_id, text_s3_key=ext_key, filename=payload.filename, char_count=len(extracted_text), is_structured=is_structured, has_tables=has_tables, attempt=payload.attempt, ) publish_rk = RK_PRESCREEN queue_metric = "prescreen" doc_status = "prescreening" doc_stage = "queued_prescreen" else: msg = DocumentExtracted( correlation_id=payload.correlation_id, document_id=payload.document_id, user_id=payload.user_id, extracted_s3_key=ext_key, char_count=len(extracted_text), ocr_used=ocr_used, is_structured=is_structured, has_tables=has_tables, attempt=payload.attempt, ) publish_rk = RK_ANALYZE queue_metric = "analyze" doc_status = "analyzing" doc_stage = "queued_analyze" await publisher.publish(msg, routing_key=publish_rk) mq_published.labels(queue=queue_metric).inc() async with self._session_factory() as session: docs = DocumentRepository(session) jobs = JobRepository(session) await docs.update_status( payload.document_id, status=doc_status, stage=doc_stage, extracted_s3_key=ext_key, ) await jobs.mark_done(payload.document_id, "extract") await session.commit() log.info( "extract_success", document_id=str(payload.document_id), format=fmt, structured=is_structured, tables=has_tables, char_count=len(extracted_text), ) async def _extract_document(self, payload: DocumentUploaded) -> ExtractedDocument: data = await self._storage.get(payload.s3_key) fmt = detect_format(data, mime=payload.mime, filename=payload.filename) def _extract() -> ExtractedDocument: return extract_document(data, mime=payload.mime, filename=payload.filename) with extract_duration.labels(format=fmt).time(): document = await asyncio.to_thread(_extract) extraction_total.labels(format=fmt, structured=str(document.is_structured).lower()).inc() return document def classify(self, exc: BaseException) -> FailureClass: from src.contract_check.core.analysis.ocr import OCRError if isinstance(exc, (ExtractionFailedError, UnsupportedFormatError)): return "extraction_failed" if isinstance(exc, OCRError): return "ocr_failed" return "infra" async def on_failure( self, payload: DocumentUploaded, failure_class: FailureClass, attempt: int, error: str ) -> None: mq_failed.labels(queue="extract", failure_class=failure_class).inc() async with self._session_factory() as session: await JobRepository(session).mark_retrying( payload.document_id, "extract", attempt=attempt, failure_class=failure_class, error=error, ) await session.commit() async def on_terminal_failure( self, payload: DocumentUploaded, failure_class: FailureClass, error: str ) -> None: mq_failed.labels(queue="extract", failure_class=failure_class).inc() async with self._session_factory() as session: docs = DocumentRepository(session) jobs = JobRepository(session) await jobs.mark_dlq( payload.document_id, "extract", failure_class=failure_class, error=error ) await docs.mark_failed(payload.document_id, stage=failure_class) await release_document_slot(session, payload.document_id) await refund_credit( session, payload.document_id, failure_class, self._settings.refund_policy, ) await session.commit() log.warning( "extract_terminal_failure", document_id=str(payload.document_id), failure_class=failure_class, )