curo1305
e9ee5d4ba5
- Add _ClassificationError sentinel raised by _run() for classification failures - Add _mark_classification_failed() async helper for final status writeback - Change decorator to bind=True, max_retries=3 - Outer sync task catches _ClassificationError and calls self.retry(countdown=[30,90,270]) - MaxRetriesExceededError caught in nested try/except to avoid re-raise from exc block - Promote test_retry_backoff and test_exhaustion_sets_failed_status from xfail to passing - Tests use push_request(retries=N) + patch.object(task, "retry") to verify countdown values - test_exhaustion_sets_failed_status uses assert_called_once_with (asyncio.run calling convention)
236 lines
9.8 KiB
Python
236 lines
9.8 KiB
Python
"""
|
|
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}
|