mirror of
https://github.com/langgenius/dify.git
synced 2026-11-12 06:11:12 +08:00
133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
"""
|
|
Celery SQL comment context for OpenTelemetry SQLCommenter.
|
|
|
|
Injects Celery-specific metadata (framework, task_name, traceparent, celery_retries,
|
|
routing_key) into SQL comments for queries executed by Celery workers. This improves
|
|
trace-to-SQL correlation and debugging in production.
|
|
|
|
Uses the OpenTelemetry context key SQLCOMMENTER_ORM_TAGS_AND_VALUES, which is read
|
|
by opentelemetry.instrumentation.sqlcommenter_utils._add_framework_tags() when the
|
|
SQLAlchemy instrumentor appends comments to SQL statements.
|
|
"""
|
|
|
|
import logging
|
|
from typing import Any, TypedDict
|
|
|
|
from celery.app.task import Task
|
|
from celery.signals import task_postrun, task_prerun
|
|
from opentelemetry import context
|
|
from opentelemetry.instrumentation.celery import CeleryInstrumentor
|
|
from opentelemetry.metrics import MeterProvider
|
|
from opentelemetry.trace import TracerProvider
|
|
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_TRACE_PROPAGATOR = TraceContextTextMapPropagator()
|
|
|
|
_SQLCOMMENTER_CONTEXT_KEY = "SQLCOMMENTER_ORM_TAGS_AND_VALUES"
|
|
_TOKEN_ATTR = "_dify_sqlcommenter_context_token"
|
|
|
|
|
|
class CelerySqlcommenterTagsDict(TypedDict, total=False):
|
|
framework: str
|
|
task_name: str
|
|
traceparent: str
|
|
celery_retries: int
|
|
routing_key: str
|
|
|
|
|
|
def _build_celery_sqlcommenter_tags(task: Any) -> CelerySqlcommenterTagsDict:
|
|
"""Build SQL commenter tags from the current Celery task and OpenTelemetry context."""
|
|
tags: CelerySqlcommenterTagsDict = {}
|
|
|
|
try:
|
|
tags["framework"] = f"celery:{_get_celery_version()}"
|
|
except Exception:
|
|
tags["framework"] = "celery:unknown"
|
|
|
|
if task and getattr(task, "name", None):
|
|
tags["task_name"] = str(task.name)
|
|
|
|
traceparent = _get_traceparent()
|
|
if traceparent:
|
|
tags["traceparent"] = traceparent
|
|
|
|
if task and hasattr(task, "request"):
|
|
request = task.request
|
|
retries = getattr(request, "retries", None)
|
|
if retries is not None and retries > 0:
|
|
tags["celery_retries"] = int(retries)
|
|
|
|
delivery_info = getattr(request, "delivery_info", None) or {}
|
|
if isinstance(delivery_info, dict):
|
|
routing_key = delivery_info.get("routing_key")
|
|
if routing_key:
|
|
tags["routing_key"] = str(routing_key)
|
|
|
|
return tags
|
|
|
|
|
|
def _get_celery_version() -> str:
|
|
import celery
|
|
|
|
return getattr(celery, "__version__", "unknown")
|
|
|
|
|
|
def _get_traceparent() -> str | None:
|
|
"""Extract traceparent from the current OpenTelemetry context."""
|
|
carrier: dict[str, str] = {}
|
|
_TRACE_PROPAGATOR.inject(carrier)
|
|
return carrier.get("traceparent")
|
|
|
|
|
|
def _on_task_prerun(*args: object, **kwargs: object) -> None:
|
|
task = kwargs.get("task")
|
|
if not isinstance(task, Task):
|
|
return
|
|
|
|
tags = _build_celery_sqlcommenter_tags(task)
|
|
if not tags:
|
|
return
|
|
|
|
current = context.get_current()
|
|
new_ctx = context.set_value(_SQLCOMMENTER_CONTEXT_KEY, tags, current)
|
|
token = context.attach(new_ctx)
|
|
# Celery reuses one Task instance across concurrent invocations. Its request
|
|
# is execution-local, so another invocation cannot overwrite this token.
|
|
setattr(task.request, _TOKEN_ATTR, token)
|
|
|
|
|
|
def _on_task_postrun(*args: object, **kwargs: object) -> None:
|
|
task = kwargs.get("task")
|
|
if not isinstance(task, Task):
|
|
return
|
|
|
|
request = task.request
|
|
token = getattr(request, _TOKEN_ATTR, None)
|
|
if token is None:
|
|
return
|
|
|
|
delattr(request, _TOKEN_ATTR)
|
|
context.detach(token)
|
|
|
|
|
|
def setup_celery_sqlcommenter(*, tracer_provider: TracerProvider, meter_provider: MeterProvider) -> None:
|
|
"""
|
|
Install Celery instrumentation with SQL comments inside its activation scope.
|
|
|
|
Celery invokes signal receivers in registration order. The SQL context must
|
|
attach after OTel's task span is active, then detach before OTel restores its
|
|
parent. Registering both handlers after instrument() would restore an ended
|
|
task span during postrun instead of leaving the caller's context intact.
|
|
"""
|
|
instrumentor = CeleryInstrumentor()
|
|
if instrumentor.is_instrumented_by_opentelemetry:
|
|
# Auto-instrumentation may already own the signal order and providers.
|
|
# Keep its spans intact rather than add a context we cannot unwind first.
|
|
logger.warning("Celery is already instrumented; skipping SQL commenter signal registration")
|
|
return
|
|
|
|
task_postrun.connect(_on_task_postrun, weak=False)
|
|
instrumentor.instrument(tracer_provider=tracer_provider, meter_provider=meter_provider)
|
|
task_prerun.connect(_on_task_prerun, weak=False)
|