diff --git a/sentry_sdk/integrations/huey.py b/sentry_sdk/integrations/huey.py index 9641160099..43c03936b1 100644 --- a/sentry_sdk/integrations/huey.py +++ b/sentry_sdk/integrations/huey.py @@ -6,10 +6,15 @@ from sentry_sdk._compat import reraise from sentry_sdk._types import TYPE_CHECKING from sentry_sdk import Hub +from sentry_sdk.api import continue_trace, get_baggage, get_traceparent from sentry_sdk.consts import OP from sentry_sdk.hub import _should_send_default_pii from sentry_sdk.integrations import DidNotEnable, Integration -from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_TASK +from sentry_sdk.tracing import ( + BAGGAGE_HEADER_NAME, + SENTRY_TRACE_HEADER_NAME, + TRANSACTION_SOURCE_TASK, +) from sentry_sdk.utils import ( capture_internal_exceptions, event_from_exception, @@ -25,7 +30,7 @@ F = TypeVar("F", bound=Callable[..., Any]) try: - from huey.api import Huey, Result, ResultGroup, Task + from huey.api import Huey, Result, ResultGroup, Task, PeriodicTask from huey.exceptions import CancelExecution, RetryTask, TaskLockedException except ImportError: raise DidNotEnable("Huey is not installed") @@ -56,6 +61,14 @@ def _sentry_enqueue(self, task): return old_enqueue(self, task) with hub.start_span(op=OP.QUEUE_SUBMIT_HUEY, description=task.name): + if not isinstance(task, PeriodicTask): + # Attach trace propagation data to task kwargs. We do + # not do this for periodic tasks, as these don't + # really have an originating transaction. + task.kwargs["sentry_headers"] = { + BAGGAGE_HEADER_NAME: get_baggage(), + SENTRY_TRACE_HEADER_NAME: get_traceparent(), + } return old_enqueue(self, task) Huey.enqueue = _sentry_enqueue @@ -145,12 +158,15 @@ def _sentry_execute(self, task, timestamp=None): scope.clear_breadcrumbs() scope.add_event_processor(_make_event_processor(task)) - transaction = Transaction( + sentry_headers = task.kwargs.pop("sentry_headers", None) + + transaction = continue_trace( + sentry_headers or {}, name=task.name, - status="ok", op=OP.QUEUE_TASK_HUEY, source=TRANSACTION_SOURCE_TASK, ) + transaction.set_status("ok") if not getattr(task, "_sentry_is_patched", False): task.execute = _wrap_task_execute(task.execute) diff --git a/tests/integrations/huey/test_huey.py b/tests/integrations/huey/test_huey.py index 0bebd91b19..48a3da97f4 100644 --- a/tests/integrations/huey/test_huey.py +++ b/tests/integrations/huey/test_huey.py @@ -172,3 +172,21 @@ def dummy_task(): assert len(event["spans"]) assert event["spans"][0]["op"] == "queue.submit.huey" assert event["spans"][0]["description"] == "different_task_name" + + +def test_huey_propagate_trace(init_huey, capture_events): + huey = init_huey() + + events = capture_events() + + @huey.task() + def propagated_trace_task(): + pass + + with start_transaction() as outer_transaction: + execute_huey_task(huey, propagated_trace_task) + + assert ( + events[0]["transaction"] == "propagated_trace_task" + ) # the "inner" transaction + assert events[0]["contexts"]["trace"]["trace_id"] == outer_transaction.trace_id