kudzues/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-04-05 01:39:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(api): type celery sqlcommenter tags with CelerySqlcommenterTagsDict TypedDict (#34526)

Browse Source
This commit is contained in:
YBoy
2026-04-03 14:06:15 +02:00
committed by GitHub
parent e85d9a0d72
commit 0f051d5886
1 changed files with 11 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
api/extensions/otel/celery_sqlcommenter.py
View File

@@ -11,7 +11,7 @@ SQLAlchemy instrumentor appends comments to SQL statements.
"""
import logging
from typing import Any
from typing import Any, TypedDict
from celery.signals import task_postrun, task_prerun
from opentelemetry import context
@@ -24,9 +24,17 @@ _SQLCOMMENTER_CONTEXT_KEY = "SQLCOMMENTER_ORM_TAGS_AND_VALUES"
_TOKEN_ATTR = "_dify_sqlcommenter_context_token"
def _build_celery_sqlcommenter_tags(task: Any) -> dict[str, str | int]:
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: dict[str, str | int] = {}
tags: CelerySqlcommenterTagsDict = {}
try:
tags["framework"] = f"celery:{_get_celery_version()}"