feat(integrations): Propagate sentry-trace and baggage to huey tasks

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)

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