copilot-swe-agent[bot]
651b48658c
Merge origin/main into feature branch, resolving 3 conflicts: - app/api/__init__.py: add classification_rules_router alongside new routers from main (audit_logs, i18n, mobile, compliance, translation) - app/models.py: keep ClassificationRuleModel alongside new models from main (MobileDevice, ComplianceTemplate, PipelineRoutingRule) - tests/conftest.py: import both ClassificationRuleModel and new models from main (AuditLog, ComplianceTemplate) Also renumber migration from 027 to 037 to chain from the latest migration on main (036_add_document_translation_fields). Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
121 lines
4.4 KiB
Python
121 lines
4.4 KiB
Python
# app/celery_app.py
|
|
|
|
import logging
|
|
import os
|
|
|
|
from celery import Celery
|
|
from celery.signals import task_failure, worker_ready
|
|
|
|
from app.config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
celery = Celery(
|
|
"document_processor",
|
|
broker=settings.redis_url,
|
|
backend=settings.redis_url,
|
|
)
|
|
|
|
|
|
# Optionally add this line to retain connection retry behavior at startup:
|
|
celery.conf.broker_connection_retry_on_startup = True
|
|
|
|
# Set the default queue and routing so that tasks are enqueued on "document_processor"
|
|
celery.conf.task_default_queue = "document_processor"
|
|
celery.conf.task_routes = {
|
|
"app.tasks.*": {"queue": "document_processor"},
|
|
}
|
|
|
|
# Mapping of document pipeline task names to the positional index of ``file_id``
|
|
# in their ``args`` tuple. These indices correspond to the task signatures:
|
|
# process_with_ocr(filename, file_id, ...) → index 1
|
|
# extract_metadata_with_gpt(filename, text, file_id) → index 2
|
|
# embed_metadata_into_pdf(path, text, metadata, file_id) → index 3
|
|
# Tasks that always pass ``file_id`` as a keyword argument
|
|
# (e.g. ``process_document``, ``finalize_document_storage``) are not listed
|
|
# here — their ``file_id`` is found via ``kwargs`` instead.
|
|
_FILE_ID_ARG_INDEX: dict[str, int] = {
|
|
"app.tasks.process_with_ocr.process_with_ocr": 1,
|
|
"app.tasks.extract_metadata_with_gpt.extract_metadata_with_gpt": 2,
|
|
"app.tasks.embed_metadata_into_pdf.embed_metadata_into_pdf": 3,
|
|
}
|
|
|
|
|
|
def _dispatch_user_failure_notification(sender, exception, args: list | None, kwargs: dict | None) -> None:
|
|
"""Best-effort per-user failure notification for document pipeline tasks.
|
|
|
|
Extracts ``file_id`` from the failed task's arguments, looks up the owning
|
|
user from the database, and dispatches a ``document.failed`` notification.
|
|
"""
|
|
from app.database import SessionLocal
|
|
from app.models import FileRecord
|
|
from app.utils.user_notification import notify_user_document_failed
|
|
|
|
task_name = sender.name if sender else ""
|
|
if not task_name.startswith("app.tasks."):
|
|
return
|
|
|
|
# 1. Resolve file_id from kwargs or positional args
|
|
file_id = (kwargs or {}).get("file_id")
|
|
if file_id is None:
|
|
idx = _FILE_ID_ARG_INDEX.get(task_name)
|
|
if idx is not None and args and len(args) > idx:
|
|
val = args[idx]
|
|
if isinstance(val, int):
|
|
file_id = val
|
|
|
|
if file_id is None:
|
|
return
|
|
|
|
# 2. Look up owner from the database
|
|
with SessionLocal() as db:
|
|
record = db.query(FileRecord).filter(FileRecord.id == file_id).first()
|
|
if not record or not record.owner_id:
|
|
return
|
|
owner_id = record.owner_id
|
|
filename = record.original_filename or record.local_filename or "unknown"
|
|
|
|
# 3. Dispatch per-user notification
|
|
error_msg = f"{type(exception).__name__}: {exception}" if exception else "Unknown error"
|
|
notify_user_document_failed(
|
|
owner_id=owner_id,
|
|
filename=os.path.basename(filename),
|
|
error=error_msg,
|
|
file_id=file_id,
|
|
)
|
|
|
|
|
|
@worker_ready.connect
|
|
def init_sentry_on_worker_ready(**kwargs):
|
|
"""Initialise Sentry SDK in the Celery worker process."""
|
|
from app.utils.sentry import init_sentry
|
|
|
|
init_sentry(integrations_extra=["celery"])
|
|
|
|
|
|
@task_failure.connect
|
|
def task_failure_handler(
|
|
sender=None, task_id=None, exception=None, args=None, kwargs=None, traceback=None, einfo=None, **kw
|
|
):
|
|
"""Handler for Celery task failures to send notifications"""
|
|
if getattr(settings, "notify_on_task_failure", True):
|
|
try:
|
|
# Import here to avoid circular imports
|
|
from app.utils.notification import notify_celery_failure
|
|
|
|
notify_celery_failure(
|
|
task_name=sender.name if sender else "Unknown",
|
|
task_id=task_id or "N/A",
|
|
exc=exception,
|
|
args=args or [],
|
|
kwargs=kwargs or {},
|
|
)
|
|
except Exception as e:
|
|
logger.exception(f"Failed to send task failure notification: {e}")
|
|
|
|
# Also dispatch a per-user failure notification for document pipeline tasks
|
|
try:
|
|
_dispatch_user_failure_notification(sender, exception, args, kwargs)
|
|
except Exception:
|
|
logger.warning("Could not dispatch per-user failure notification", exc_info=True)
|