celery

Author: sentrivana

sentry_sdk/integrations/celery/__init__.py

@@ -14,8 +14,9 @@
 )
 from sentry_sdk.integrations.celery.utils import _now_seconds_since_epoch
 from sentry_sdk.integrations.logging import ignore_logger
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, TransactionSource
-from sentry_sdk.tracing_utils import Baggage
+from sentry_sdk.tracing_utils import Baggage, has_span_streaming_enabled
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     ensure_integration_enabled,
@@ -247,7 +248,8 @@ def _wrap_task_run(f: "F") -> "F":
     def apply_async(*args: "Any", **kwargs: "Any") -> "Any":
         # Note: kwargs can contain headers=None, so no setdefault!
         # Unsure which backend though.
-        integration = sentry_sdk.get_client().get_integration(CeleryIntegration)
+        client = sentry_sdk.get_client()
+        integration = client.get_integration(CeleryIntegration)
         if integration is None:
             return f(*args, **kwargs)
 
@@ -266,17 +268,23 @@ def apply_async(*args: "Any", **kwargs: "Any") -> "Any":
         else:
             task_name = "<unknown Celery task>"
 
+        span_streaming = has_span_streaming_enabled(client.options)
+
         task_started_from_beat = sentry_sdk.get_isolation_scope()._name == "celery-beat"
 
-        span_mgr: "Union[Span, NoOpMgr]" = (
-            sentry_sdk.start_span(
-                op=OP.QUEUE_SUBMIT_CELERY,
-                name=task_name,
-                origin=CeleryIntegration.origin,
-            )
-            if not task_started_from_beat
-            else NoOpMgr()
-        )
+        span_mgr: "Union[StreamedSpan, Span, NoOpMgr]" = NoOpMgr()
+        if span_streaming:
+            if not task_started_from_beat:
+                span_mgr = sentry_sdk.traces.start_span(name=task_name)
+                span_mgr.set_op(OP.QUEUE_SUBMIT_CELERY)
+                span_mgr.set_origin(CeleryIntegration.origin)
+        else:
+            if not task_started_from_beat:
+                span_mgr = sentry_sdk.start_span(
+                    op=OP.QUEUE_SUBMIT_CELERY,
+                    name=task_name,
+                    origin=CeleryIntegration.origin,
+                )
 
         with span_mgr as span:
             kwargs["headers"] = _update_celery_task_headers(
@@ -295,8 +303,13 @@ def _wrap_tracer(task: "Any", f: "F") -> "F":
     # Also because in Celery 3, signal dispatch returns early if one handler
     # crashes.
     @wraps(f)
-    @ensure_integration_enabled(CeleryIntegration, f)
     def _inner(*args: "Any", **kwargs: "Any") -> "Any":
+        client = sentry_sdk.get_client()
+        if client.get_integration(CeleryIntegration) is None:
+            return f(*args, **kwargs)
+
+        span_streaming = has_span_streaming_enabled(client.options)
+
         with isolation_scope() as scope:
             scope._name = "celery"
             scope.clear_breadcrumbs()
@@ -308,37 +321,53 @@ def _inner(*args: "Any", **kwargs: "Any") -> "Any":
             # something such as attribute access can fail.
             with capture_internal_exceptions():
                 headers = args[3].get("headers") or {}
-                transaction = continue_trace(
-                    headers,
-                    op=OP.QUEUE_TASK_CELERY,
-                    name="unknown celery task",
-                    source=TransactionSource.TASK,
-                    origin=CeleryIntegration.origin,
-                )
-                transaction.name = task.name
-                transaction.set_status(SPANSTATUS.OK)
+                if span_streaming:
+                    sentry_sdk.traces.continue_trace(headers)
+                    transaction = sentry_sdk.traces.start_span(
+                        name="unknown celery task"
+                    )
+                    transaction.set_origin(CeleryIntegration.origin)
+                    transaction.set_source(TransactionSource.TASK)
+                    transaction.set_op(OP.QUEUE_TASK_CELERY)
+
+                    span_ctx = transaction
+
+                else:
+                    transaction = continue_trace(
+                        headers,
+                        op=OP.QUEUE_TASK_CELERY,
+                        name="unknown celery task",
+                        source=TransactionSource.TASK,
+                        origin=CeleryIntegration.origin,
+                    )
+                    transaction.name = task.name
+                    transaction.set_status(SPANSTATUS.OK)
+
+                    span_ctx = sentry_sdk.start_transaction(
+                        transaction,
+                        custom_sampling_context={
+                            "celery_job": {
+                                "task": task.name,
+                                # for some reason, args[1] is a list if non-empty but a
+                                # tuple if empty
+                                "args": list(args[1]),
+                                "kwargs": args[2],
+                            }
+                        },
+                    )
 
             if transaction is None:
                 return f(*args, **kwargs)
 
-            with sentry_sdk.start_transaction(
-                transaction,
-                custom_sampling_context={
-                    "celery_job": {
-                        "task": task.name,
-                        # for some reason, args[1] is a list if non-empty but a
-                        # tuple if empty
-                        "args": list(args[1]),
-                        "kwargs": args[2],
-                    }
-                },
-            ):
+            with span_ctx:
                 return f(*args, **kwargs)
 
     return _inner  # type: ignore
 
 
-def _set_messaging_destination_name(task: "Any", span: "Span") -> None:
+def _set_messaging_destination_name(
+    task: "Any", span: "Union[StreamedSpan, Span]"
+) -> None:
     """Set "messaging.destination.name" tag for span"""
     with capture_internal_exceptions():
         delivery_info = task.request.delivery_info
@@ -347,7 +376,10 @@ def _set_messaging_destination_name(task: "Any", span: "Span") -> None:
             if delivery_info.get("exchange") == "" and routing_key is not None:
                 # Empty exchange indicates the default exchange, meaning the tasks
                 # are sent to the queue with the same name as the routing key.
-                span.set_data(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key)
+                if isinstance(span, StreamedSpan):
+                    span.set_attribute(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key)
+                else:
+                    span.set_data(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key)
 
 
 def _wrap_task_call(task: "Any", f: "F") -> "F":
@@ -359,14 +391,32 @@ def _wrap_task_call(task: "Any", f: "F") -> "F":
     # but if we ever remove the @ensure_integration_enabled decorator, we need
     # to add @functools.wraps(f) here.
     # https://github.com/getsentry/sentry-python/issues/421
-    @ensure_integration_enabled(CeleryIntegration, f)
     def _inner(*args: "Any", **kwargs: "Any") -> "Any":
+        client = sentry_sdk.get_client()
+        if client.get_integration(CeleryIntegration) is None:
+            return f(*args, **kwargs)
+
+        span_streaming = has_span_streaming_enabled(client.options)
+
         try:
-            with sentry_sdk.start_span(
-                op=OP.QUEUE_PROCESS,
-                name=task.name,
-                origin=CeleryIntegration.origin,
-            ) as span:
+            span: "Union[Span, StreamedSpan]"
+            if span_streaming:
+                span = sentry_sdk.traces.start_span(name=task.name)
+                span.set_op(OP.QUEUE_PROCESS)
+                span.set_origin(CeleryIntegration.origin)
+            else:
+                span = sentry_sdk.start_span(
+                    op=OP.QUEUE_PROCESS,
+                    name=task.name,
+                    origin=CeleryIntegration.origin,
+                )
+
+            with span:
+                if isinstance(span, StreamedSpan):
+                    set_on_span = span.set_attribute
+                else:
+                    set_on_span = span.set_data
+
                 _set_messaging_destination_name(task, span)
 
                 latency = None
@@ -381,18 +431,18 @@ def _inner(*args: "Any", **kwargs: "Any") -> "Any":
 
                 if latency is not None:
                     latency *= 1000  # milliseconds
-                    span.set_data(SPANDATA.MESSAGING_MESSAGE_RECEIVE_LATENCY, latency)
+                    set_on_span(SPANDATA.MESSAGING_MESSAGE_RECEIVE_LATENCY, latency)
 
                 with capture_internal_exceptions():
-                    span.set_data(SPANDATA.MESSAGING_MESSAGE_ID, task.request.id)
+                    set_on_span(SPANDATA.MESSAGING_MESSAGE_ID, task.request.id)
 
                 with capture_internal_exceptions():
-                    span.set_data(
+                    set_on_span(
                         SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT, task.request.retries
                     )
 
                 with capture_internal_exceptions():
-                    span.set_data(
+                    set_on_span(
                         SPANDATA.MESSAGING_SYSTEM,
                         task.app.connection().transport.driver_type,
                     )
@@ -467,8 +517,13 @@ def sentry_workloop(*args: "Any", **kwargs: "Any") -> "Any":
 def _patch_producer_publish() -> None:
     original_publish = Producer.publish
 
-    @ensure_integration_enabled(CeleryIntegration, original_publish)
     def sentry_publish(self: "Producer", *args: "Any", **kwargs: "Any") -> "Any":
+        client = sentry_sdk.get_client()
+        if client.get_integration(CeleryIntegration) is None:
+            return original_publish(self, *args, **kwargs)
+
+        span_streaming = has_span_streaming_enabled(client.options)
+
         kwargs_headers = kwargs.get("headers", {})
         if not isinstance(kwargs_headers, Mapping):
             # Ensure kwargs_headers is a Mapping, so we can safely call get().
@@ -485,24 +540,37 @@ def sentry_publish(self: "Producer", *args: "Any", **kwargs: "Any") -> "Any":
         routing_key = kwargs.get("routing_key")
         exchange = kwargs.get("exchange")
 
-        with sentry_sdk.start_span(
-            op=OP.QUEUE_PUBLISH,
-            name=task_name,
-            origin=CeleryIntegration.origin,
-        ) as span:
+        span: "Union[StreamedSpan, Span]"
+        if span_streaming:
+            span = sentry_sdk.traces.start_span(name=task_name)
+            span.set_op(OP.QUEUE_PUBLISH)
+            span.set_origin(CeleryIntegration.origin)
+        else:
+            span = sentry_sdk.start_span(
+                op=OP.QUEUE_PUBLISH,
+                name=task_name,
+                origin=CeleryIntegration.origin,
+            )
+
+        with span:
+            if isinstance(span, StreamedSpan):
+                set_on_span = span.set_attribute
+            else:
+                set_on_span = span.set_data
+
             if task_id is not None:
-                span.set_data(SPANDATA.MESSAGING_MESSAGE_ID, task_id)
+                set_on_span(SPANDATA.MESSAGING_MESSAGE_ID, task_id)
 
             if exchange == "" and routing_key is not None:
                 # Empty exchange indicates the default exchange, meaning messages are
                 # routed to the queue with the same name as the routing key.
-                span.set_data(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key)
+                set_on_span(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key)
 
             if retries is not None:
-                span.set_data(SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT, retries)
+                set_on_span(SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT, retries)
 
             with capture_internal_exceptions():
-                span.set_data(
+                set_on_span(
                     SPANDATA.MESSAGING_SYSTEM, self.connection.transport.driver_type
                 )