""" Celery tasks for document processing in DocuVault. extract_and_classify — called via .delay(document_id) by the upload handler. The task is a plain sync def (Celery workers have no asyncio event loop); it bridges into the async service layer via asyncio.run(). Flow: 1. Open a fresh AsyncSession (one per task invocation — never share sessions) 2. Look up the Document row to get the MinIO object_key 3. Retrieve file bytes from MinIO via the storage backend 4. Extract text from bytes using services.extractor 5. Persist extracted_text back to the Document row 6. Call services.classifier.classify_document to assign topics 7. Return a result dict; classification failures are retried with exponential backoff Celery retry harness (D-09, D-10 — Pitfall 3 from RESEARCH.md): CRITICAL: self.retry() MUST be raised from the outer sync task, NOT inside asyncio.run(). _ClassificationError is a sentinel raised by _run() to signal a retryable classification failure. The outer extract_and_classify catches it and calls self.retry(countdown=...) in the sync layer. """ import asyncio from celery.exceptions import MaxRetriesExceededError from celery_app import celery_app class _ClassificationError(Exception): """Sentinel exception raised by _run() to signal a retryable classification failure. This exception escapes asyncio.run() and is caught by the outer sync task, which then calls self.retry() in the Celery sync layer (not inside asyncio.run). Non-classification failures (extract_failed, invalid_id) are NOT retried — they return a dict directly from _run() without raising _ClassificationError. """ pass async def _mark_classification_failed(document_id: str) -> None: """Write doc.status = 'classification_failed' after all retries are exhausted. D-10: called by extract_and_classify's MaxRetriesExceededError handler via asyncio.run(_mark_classification_failed(document_id)) so the final status is written to DB regardless of Celery task context. """ import uuid as _uuid from db.session import AsyncSessionLocal from db.models import Document async with AsyncSessionLocal() as session: try: doc_uuid = _uuid.UUID(document_id) except ValueError: return # Silently ignore invalid IDs — nothing to update doc = await session.get(Document, doc_uuid) if doc is not None: doc.status = "classification_failed" await session.commit() @celery_app.task( name="tasks.document_tasks.extract_and_classify", bind=True, max_retries=3, ) def extract_and_classify(self, document_id: str) -> dict: """Synchronous Celery entry-point — delegates to async _run via asyncio.run. Retry harness (D-09): classification failures are retried up to 3 times with exponential backoff: 30s / 90s / 270s. Pitfall 3 guard: self.retry() is called HERE in the sync layer, never inside asyncio.run(). _ClassificationError is the sentinel that escapes asyncio.run() to trigger the retry. """ try: return asyncio.run(_run(document_id)) except _ClassificationError as exc: # Exponential backoff: 30s on first retry, 90s on second, 270s on third countdowns = [30, 90, 270] countdown = countdowns[min(self.request.retries, 2)] try: raise self.retry(exc=exc, countdown=countdown) except MaxRetriesExceededError: # All retries exhausted — write final failure status to DB asyncio.run(_mark_classification_failed(document_id)) return {"document_id": document_id, "status": "classification_failed"} async def _run(document_id: str) -> dict: """Async body of extract_and_classify. Opens its own AsyncSession (not shared with the upload request) to avoid cross-thread session contamination. Cloud-aware: when doc.storage_backend != 'minio', uses get_storage_backend_for_document() to retrieve bytes from the correct cloud backend instead of hardcoding MinIO. Classification failures raise _ClassificationError (D-09 retryable sentinel). Non-classification failures return a status dict (not retried). """ import uuid as _uuid from db.session import AsyncSessionLocal from db.models import Document from services import extractor, classifier from storage import get_storage_backend, get_storage_backend_for_document async with AsyncSessionLocal() as session: # ── Step 1: fetch Document row ───────────────────────────────────────── try: doc_uuid = _uuid.UUID(document_id) except ValueError: return {"document_id": document_id, "status": "invalid_id"} doc = await session.get(Document, doc_uuid) if doc is None: return {"document_id": document_id, "status": "not_found"} if not doc.object_key: return {"document_id": document_id, "status": "missing_object"} # ── Resolve per-user AI config (D-14, D-15) ──────────────────────────── from db.models import User from config import settings as app_settings user = await session.get(User, doc.user_id) if doc.user_id else None ai_provider = (user.ai_provider if user else None) or app_settings.default_ai_provider ai_model = (user.ai_model if user else None) or app_settings.default_ai_model # ── Step 2: retrieve bytes from the correct backend ──────────────────── # Cloud-aware: routes to cloud backend for non-MinIO documents (Plan 09). # T-05-09-03: cloud credentials are loaded from DB inside this task's own # session — no credentials travel through the Celery broker message. try: if doc.storage_backend is None or doc.storage_backend == "minio": backend = get_storage_backend() file_bytes = await backend.get_object(doc.object_key) else: # Cloud path: user must be present (doc.user_id set at upload time) if user is None: return {"document_id": document_id, "status": "missing_user"} from storage.exceptions import CloudConnectionError try: backend = await get_storage_backend_for_document(doc, user, session) file_bytes = await backend.get_object(doc.object_key) except CloudConnectionError: return { "document_id": document_id, "status": "extract_failed", "error": "cloud backend error", } except Exception as e: return { "document_id": document_id, "status": "extract_failed", "error": f"retrieval failed: {e}", } # ── Step 3: extract text from bytes ──────────────────────────────────── try: text = extractor.extract_text_from_bytes(file_bytes, doc.content_type) doc.extracted_text = text await session.commit() except Exception as e: return { "document_id": document_id, "status": "extract_failed", "error": f"Text extraction failed: {e}", } # ── Step 4: classify document (retryable via _ClassificationError) ───── try: topics = await classifier.classify_document(session, document_id, ai_provider=ai_provider, ai_model=ai_model) return { "document_id": document_id, "status": "classified", "topics": topics, } except Exception as e: # Raise sentinel to allow the outer sync layer to call self.retry() # (D-09 Pitfall 3: self.retry() cannot be called inside asyncio.run) raise _ClassificationError(str(e)) from e @celery_app.task(name="tasks.document_tasks.cleanup_abandoned_uploads") def cleanup_abandoned_uploads() -> dict: """Periodic Celery beat task — deletes Document rows with status='pending' older than 1 hour and their MinIO objects (D-06). Enqueued by Celery beat every 30 minutes (celery_app.py beat_schedule). Quota is never reserved for pending rows — no quota cleanup needed. """ return asyncio.run(_cleanup_abandoned()) async def _cleanup_abandoned() -> dict: """Async body for cleanup_abandoned_uploads. Selects Document rows with status='pending' older than 1 hour, removes their MinIO objects (best-effort), then deletes the DB rows. Returns {"cleaned": N} count. """ from datetime import datetime, timezone, timedelta from sqlalchemy import select from db.session import AsyncSessionLocal from db.models import Document from storage import get_storage_backend cutoff = datetime.now(timezone.utc) - timedelta(hours=1) async with AsyncSessionLocal() as session: result = await session.execute( select(Document).where( Document.status == "pending", Document.created_at < cutoff, ) ) docs = result.scalars().all() backend = get_storage_backend() cleaned = 0 for doc in docs: try: if doc.object_key: await backend.delete_object(doc.object_key) except Exception: pass # MinIO object may not exist yet — safe to ignore await session.delete(doc) cleaned += 1 await session.commit() return {"cleaned": cleaned}