Corp/backend/eusrr_backend/celery.py

"""Celery configuration and task lifecycle logging for EUSRR."""

import logging
import os
from inspect import signature
from threading import Lock
from time import monotonic

from celery import Celery
from celery.signals import before_task_publish, task_failure, task_postrun, task_prerun

from .logging_utils import (
    bind_log_context,
    get_log_context,
    new_request_id,
    reset_log_context,
    resolve_client_ip,
)


logger = logging.getLogger("eusrr.celery")
_task_state = {}
_task_state_lock = Lock()
_TASK_LOG_CONTEXT_HEADER = "eusrr_log_context"
_PROPAGATED_LOG_FIELDS = ("request_id", "user_id", "client_ip")

# Устанавливаем настройки Django по умолчанию
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "eusrr_backend.settings")

app = Celery("eusrr_backend")

# Загружаем конфигурацию из Django settings
# namespace='CELERY' означает что все настройки Celery в settings.py
# должны начинаться с префикса CELERY_
app.config_from_object("django.conf:settings", namespace="CELERY")

# Автоматически находит tasks.py в каждом приложении
app.autodiscover_tasks()


@before_task_publish.connect
def inject_task_log_context(headers=None, **kwargs):
    """Copy safe request correlation fields into Celery message headers."""
    if headers is None:
        return
    context = get_log_context()
    propagated = {
        key: context[key]
        for key in _PROPAGATED_LOG_FIELDS
        if context.get(key) is not None
    }
    if propagated:
        headers[_TASK_LOG_CONTEXT_HEADER] = propagated


def _task_headers(sender=None, task=None) -> dict:
    task_object = task or sender
    request = getattr(task_object, "request", None)
    headers = getattr(request, "headers", None)
    return headers if isinstance(headers, dict) else {}


def _task_user_id(sender=None, task=None, task_args=None, task_kwargs=None):
    if isinstance(task_kwargs, dict) and task_kwargs.get("user_id") is not None:
        return task_kwargs["user_id"]

    task_object = task or sender
    run = getattr(task_object, "run", None)
    if not callable(run):
        return None
    try:
        arguments = (
            signature(run)
            .bind_partial(
                *(task_args or ()),
                **(task_kwargs or {}),
            )
            .arguments
        )
    except (TypeError, ValueError):
        return None
    return arguments.get("user_id")


def _propagated_task_context(
    sender=None,
    task=None,
    task_args=None,
    task_kwargs=None,
) -> dict:
    headers = _task_headers(sender=sender, task=task)
    raw_context = headers.get(_TASK_LOG_CONTEXT_HEADER, {})
    context = raw_context if isinstance(raw_context, dict) else {}
    propagated = {}

    request_id = context.get("request_id")
    if request_id is not None:
        propagated["request_id"] = new_request_id(str(request_id))

    client_ip = resolve_client_ip(context.get("client_ip"))
    if client_ip:
        propagated["client_ip"] = client_ip

    user_id = context.get("user_id")
    if user_id is None:
        user_id = _task_user_id(
            sender=sender,
            task=task,
            task_args=task_args,
            task_kwargs=task_kwargs,
        )
    if isinstance(user_id, (int, str)) and not isinstance(user_id, bool):
        propagated["user_id"] = user_id

    return propagated


@task_prerun.connect
def log_task_started(
    sender=None,
    task_id=None,
    task=None,
    args=None,
    kwargs=None,
    **signal_kwargs,
):
    task_name = getattr(sender, "name", None) or getattr(task, "name", None)
    token = bind_log_context(
        **_propagated_task_context(
            sender=sender,
            task=task,
            task_args=args,
            task_kwargs=kwargs,
        ),
        task_id=task_id,
        task_name=task_name,
    )
    with _task_state_lock:
        _task_state[task_id] = (token, monotonic())
    logger.info(
        "Celery task started",
        extra={
            "event": "celery.task.started",
            "task_id": task_id,
            "task_name": task_name,
        },
    )


@task_failure.connect
def log_task_failed(
    sender=None, task_id=None, exception=None, traceback=None, **kwargs
):
    task_name = getattr(sender, "name", None)
    logger.error(
        "Celery task failed",
        extra={
            "event": "celery.task.failed",
            "task_id": task_id,
            "task_name": task_name,
            "error_type": type(exception).__name__ if exception else None,
        },
        exc_info=(type(exception), exception, traceback) if exception else None,
    )


@task_postrun.connect
def log_task_completed(sender=None, task_id=None, state=None, **kwargs):
    with _task_state_lock:
        task_context = _task_state.pop(task_id, None)
    token, started = task_context or (None, monotonic())
    duration_ms = round((monotonic() - started) * 1000, 2)
    task_name = getattr(sender, "name", None)
    level = logging.INFO if state == "SUCCESS" else logging.WARNING
    logger.log(
        level,
        "Celery task completed",
        extra={
            "event": "celery.task.completed",
            "task_id": task_id,
            "task_name": task_name,
            "task_state": state,
            "duration_ms": duration_ms,
        },
    )
    if token is not None:
        reset_log_context(token)


@app.task(bind=True, ignore_result=True)
def debug_task(self):
    """Тестовая задача для проверки работоспособности Celery"""
    logger.debug(
        "Celery debug task executed",
        extra={
            "event": "celery.debug_task.executed",
            "task_id": self.request.id,
        },
    )