feat(integrations): Add tracing to DramatiqIntegration

sentry_sdk/consts.py

@@ -695,6 +695,7 @@ class OP:
     QUEUE_TASK_HUEY = "queue.task.huey"
     QUEUE_SUBMIT_RAY = "queue.submit.ray"
     QUEUE_TASK_RAY = "queue.task.ray"
+    QUEUE_TASK_DRAMATIQ = "queue.task.dramatiq"
     SUBPROCESS = "subprocess"
     SUBPROCESS_WAIT = "subprocess.wait"
     SUBPROCESS_COMMUNICATE = "subprocess.communicate"

sentry_sdk/integrations/dramatiq.py

@@ -1,23 +1,35 @@
 import json
+import contextvars
 
 import sentry_sdk
-from sentry_sdk.integrations import Integration
+from sentry_sdk.consts import OP, SPANSTATUS
+from sentry_sdk.api import continue_trace, get_baggage, get_traceparent
+from sentry_sdk.integrations import Integration, DidNotEnable
 from sentry_sdk.integrations._wsgi_common import request_body_within_bounds
+from sentry_sdk.tracing import (
+    BAGGAGE_HEADER_NAME,
+    SENTRY_TRACE_HEADER_NAME,
+    TransactionSource,
+)
 from sentry_sdk.utils import (
     AnnotatedValue,
     capture_internal_exceptions,
     event_from_exception,
 )
 
-from dramatiq.broker import Broker  # type: ignore
-from dramatiq.message import Message  # type: ignore
-from dramatiq.middleware import Middleware, default_middleware  # type: ignore
-from dramatiq.errors import Retry  # type: ignore
+try:
+    from dramatiq.broker import Broker
+    from dramatiq.message import Message, R
+    from dramatiq.middleware import Middleware, default_middleware
+    from dramatiq.errors import Retry
+except ImportError:
+    raise DidNotEnable("Dramatiq is not installed")
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Any, Callable, Dict, Optional, Union
+    from sentry_sdk.tracing import Transaction
     from sentry_sdk._types import Event, Hint
 
 
@@ -73,10 +85,10 @@ def sentry_patched_broker__init__(self, *args, **kw):
         kw["middleware"] = middleware
         original_broker__init__(self, *args, **kw)
 
-    Broker.__init__ = sentry_patched_broker__init__
+    Broker.__init__ = sentry_patched_broker__init__  # type: ignore[method-assign]
 
 
-class SentryMiddleware(Middleware):  # type: ignore[misc]
+class SentryMiddleware(Middleware):
     """
     A Dramatiq middleware that automatically captures and sends
     exceptions to Sentry.
@@ -85,50 +97,78 @@ class SentryMiddleware(Middleware):  # type: ignore[misc]
     DramatiqIntegration.
     """
 
+    _transaction = contextvars.ContextVar(
+        "_transaction"
+    )  # type: contextvars.ContextVar[Transaction]
+
+    def before_enqueue(self, broker, message, delay):
+        # type: (Broker, Message[R], int) -> None
+        message.options["sentry_headers"] = {
+            BAGGAGE_HEADER_NAME: get_baggage(),
+            SENTRY_TRACE_HEADER_NAME: get_traceparent(),
+        }
+
     def before_process_message(self, broker, message):
-        # type: (Broker, Message) -> None
+        # type: (Broker, Message[R]) -> None
         integration = sentry_sdk.get_client().get_integration(DramatiqIntegration)
         if integration is None:
             return
 
-        message._scope_manager = sentry_sdk.new_scope()
-        message._scope_manager.__enter__()
-
         scope = sentry_sdk.get_current_scope()
-        scope.set_transaction_name(message.actor_name)
-        scope.set_extra("dramatiq_message_id", message.message_id)
+        scope.set_tag("dramatiq_message_id", message.message_id)
+        scope.clear_breadcrumbs()
         scope.add_event_processor(_make_message_event_processor(message, integration))
 
+        transaction = continue_trace(
+            message.options.get("sentry_headers") or {},
+            name=message.actor_name,
+            op=OP.QUEUE_TASK_DRAMATIQ,
+            source=TransactionSource.TASK,
+            # origin=DramatiqIntegration.origin,
+        )
+        transaction.set_status(SPANSTATUS.OK)
+        sentry_sdk.start_transaction(
+            transaction,
+            name=message.actor_name,
+            op=OP.QUEUE_PROCESS,
+            source=TransactionSource.TASK,
+        )
+        transaction.__enter__()
+        self._transaction.set(transaction)
+
     def after_process_message(self, broker, message, *, result=None, exception=None):
-        # type: (Broker, Message, Any, Optional[Any], Optional[Exception]) -> None
-        integration = sentry_sdk.get_client().get_integration(DramatiqIntegration)
-        if integration is None:
-            return
+        # type: (Broker, Message[R], Optional[Any], Optional[Exception]) -> None
 
         actor = broker.get_actor(message.actor_name)
         throws = message.options.get("throws") or actor.options.get("throws")
 
-        try:
-            if (
-                exception is not None
-                and not (throws and isinstance(exception, throws))
-                and not isinstance(exception, Retry)
-            ):
-                event, hint = event_from_exception(
-                    exception,
-                    client_options=sentry_sdk.get_client().options,
-                    mechanism={
-                        "type": DramatiqIntegration.identifier,
-                        "handled": False,
-                    },
-                )
-                sentry_sdk.capture_event(event, hint=hint)
-        finally:
-            message._scope_manager.__exit__(None, None, None)
+        transaction = self._transaction.get()
+
+        is_event_capture_required = (
+            exception is not None
+            and not (throws and isinstance(exception, throws))
+            and not isinstance(exception, Retry)
+        )
+        if not is_event_capture_required:
+            # normal transaction finish
+            transaction.__exit__(None, None, None)
+            return
+
+        event, hint = event_from_exception(
+            exception,  # type: ignore[arg-type]
+            client_options=sentry_sdk.get_client().options,
+            mechanism={
+                "type": DramatiqIntegration.identifier,
+                "handled": False,
+            },
+        )
+        sentry_sdk.capture_event(event, hint=hint)
+        # transaction error
+        transaction.__exit__(type(exception), exception, None)
 
 
 def _make_message_event_processor(message, integration):
-    # type: (Message, DramatiqIntegration) -> Callable[[Event, Hint], Optional[Event]]
+    # type: (Message[R], DramatiqIntegration) -> Callable[[Event, Hint], Optional[Event]]
 
     def inner(event, hint):
         # type: (Event, Hint) -> Optional[Event]
@@ -142,7 +182,7 @@ def inner(event, hint):
 
 class DramatiqMessageExtractor:
     def __init__(self, message):
-        # type: (Message) -> None
+        # type: (Message[R]) -> None
         self.message_data = dict(message.asdict())
 
     def content_length(self):

tests/integrations/dramatiq/test_dramatiq.py

@@ -5,12 +5,21 @@
 from dramatiq.brokers.stub import StubBroker
 
 import sentry_sdk
+from sentry_sdk.tracing import TransactionSource
+from sentry_sdk import start_transaction
+from sentry_sdk.consts import SPANSTATUS
 from sentry_sdk.integrations.dramatiq import DramatiqIntegration
 
+# from sentry_sdk.integrations.logging import LoggingIntegration
 
-@pytest.fixture
-def broker(sentry_init):
-    sentry_init(integrations=[DramatiqIntegration()])
+
+@pytest.fixture(scope="function")
+def broker(request, sentry_init):
+    sentry_init(
+        integrations=[DramatiqIntegration()],
+        traces_sample_rate=getattr(request, "param", None),
+        # disabled_integrations=[LoggingIntegration()],
+    )
     broker = StubBroker()
     broker.emit_after("process_boot")
     dramatiq.set_broker(broker)
@@ -44,22 +53,61 @@ def dummy_actor(x, y):
     assert exception["type"] == "ZeroDivisionError"
 
 
-def test_that_actor_name_is_set_as_transaction(broker, worker, capture_events):
+@pytest.mark.parametrize(
+    "broker,expected_span_status",
+    [
+        (1.0, SPANSTATUS.INTERNAL_ERROR),
+        (1.0, SPANSTATUS.OK),
+    ],
+    ids=["error", "success"],
+    indirect=["broker"],
+)
+def test_task_transaction(broker, worker, capture_events, expected_span_status):
     events = capture_events()
+    task_fails = expected_span_status == SPANSTATUS.INTERNAL_ERROR
 
     @dramatiq.actor(max_retries=0)
     def dummy_actor(x, y):
         return x / y
 
-    dummy_actor.send(1, 0)
+    dummy_actor.send(1, int(not task_fails))
     broker.join(dummy_actor.queue_name)
     worker.join()
 
+    if task_fails:
+        error_event = events.pop(0)
+        exception = error_event["exception"]["values"][0]
+        assert exception["type"] == "ZeroDivisionError"
+        # todo: failed assert. Logging instead of dramatiq
+        # assert exception["mechanism"]["type"] == DramatiqIntegration.identifier
+
     (event,) = events
+    assert event["type"] == "transaction"
     assert event["transaction"] == "dummy_actor"
+    assert event["transaction_info"] == {"source": TransactionSource.TASK}
+    assert event["contexts"]["trace"]["status"] == expected_span_status
 
 
-def test_that_dramatiq_message_id_is_set_as_extra(broker, worker, capture_events):
+@pytest.mark.parametrize("broker", [1.0], indirect=True)
+def test_dramatiq_propagate_trace(broker, worker, capture_events):
+    events = capture_events()
+
+    @dramatiq.actor(max_retries=0)
+    def propagated_trace_task():
+        pass
+
+    with start_transaction() as outer_transaction:
+        propagated_trace_task.send()
+        broker.join(propagated_trace_task.queue_name)
+        worker.join()
+
+    assert (
+        events[0]["transaction"] == "propagated_trace_task"
+    )  # the "inner" transaction
+    assert events[0]["contexts"]["trace"]["trace_id"] == outer_transaction.trace_id
+
+
+def test_that_dramatiq_message_id_is_set_as_tag(broker, worker, capture_events):
     events = capture_events()
 
     @dramatiq.actor(max_retries=0)
@@ -72,13 +120,14 @@ def dummy_actor(x, y):
     worker.join()
 
     event_message, event_error = events
-    assert "dramatiq_message_id" in event_message["extra"]
-    assert "dramatiq_message_id" in event_error["extra"]
+
+    assert "dramatiq_message_id" in event_message["tags"]
+    assert "dramatiq_message_id" in event_error["tags"]
     assert (
-        event_message["extra"]["dramatiq_message_id"]
-        == event_error["extra"]["dramatiq_message_id"]
+        event_message["tags"]["dramatiq_message_id"]
+        == event_error["tags"]["dramatiq_message_id"]
     )
-    msg_ids = [e["extra"]["dramatiq_message_id"] for e in events]
+    msg_ids = [e["tags"]["dramatiq_message_id"] for e in events]
     assert all(uuid.UUID(msg_id) and isinstance(msg_id, str) for msg_id in msg_ids)