feat(rq): Support span streaming

Author: alexander-alderman-webb

sentry_sdk/integrations/rq.py

@@ -2,15 +2,15 @@
 
 import sentry_sdk
 from sentry_sdk.api import continue_trace
-from sentry_sdk.consts import OP
+from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version
 from sentry_sdk.integrations.logging import ignore_logger
-from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.scope import Scope, should_send_default_pii
 from sentry_sdk.tracing import TransactionSource
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import (
     SENSITIVE_DATA_SUBSTITUTE,
     capture_internal_exceptions,
-    ensure_integration_enabled,
     event_from_exception,
     format_timestamp,
     parse_version,
@@ -61,30 +61,57 @@ def setup_once() -> None:
 
         old_perform_job = worker_cls.perform_job
 
-        @ensure_integration_enabled(RqIntegration, old_perform_job)
         def sentry_patched_perform_job(
             self: "Any", job: "Job", *args: "Queue", **kwargs: "Any"
         ) -> bool:
+            client = sentry_sdk.get_client()
+            if client.get_integration(RqIntegration) is None:
+                return old_perform_job(self, job, *args, **kwargs)
+
             with sentry_sdk.new_scope() as scope:
                 scope.clear_breadcrumbs()
                 scope.add_event_processor(_make_event_processor(weakref.ref(job)))
 
-                transaction = continue_trace(
-                    job.meta.get("_sentry_trace_headers") or {},
-                    op=OP.QUEUE_TASK_RQ,
-                    name="unknown RQ task",
-                    source=TransactionSource.TASK,
-                    origin=RqIntegration.origin,
-                )
-
-                with capture_internal_exceptions():
-                    transaction.name = job.func_name
-
-                with sentry_sdk.start_transaction(
-                    transaction,
-                    custom_sampling_context={"rq_job": job},
-                ):
-                    rv = old_perform_job(self, job, *args, **kwargs)
+                if has_span_streaming_enabled(client.options):
+                    sentry_sdk.traces.continue_trace(
+                        job.meta.get("_sentry_trace_headers") or {}
+                    )
+
+                    Scope.set_custom_sampling_context({"rq_job": job})
+
+                    span_name = "unknown RQ task"
+                    with capture_internal_exceptions():
+                        span_name = job.func_name
+
+                    with sentry_sdk.traces.start_span(
+                        name=span_name,
+                        attributes={
+                            "sentry.op": OP.QUEUE_TASK_RQ,
+                            "sentry.origin": RqIntegration.origin,
+                            "sentry.span.source": TransactionSource.TASK,
+                            SPANDATA.MESSAGING_MESSAGE_ID: job.id,
+                            SPANDATA.CODE_FUNCTION_NAME: job.func_name,
+                        },
+                        parent_span=None,
+                    ):
+                        rv = old_perform_job(self, job, *args, **kwargs)
+                else:
+                    transaction = continue_trace(
+                        job.meta.get("_sentry_trace_headers") or {},
+                        op=OP.QUEUE_TASK_RQ,
+                        name="unknown RQ task",
+                        source=TransactionSource.TASK,
+                        origin=RqIntegration.origin,
+                    )
+
+                    with capture_internal_exceptions():
+                        transaction.name = job.func_name
+
+                    with sentry_sdk.start_transaction(
+                        transaction,
+                        custom_sampling_context={"rq_job": job},
+                    ):
+                        rv = old_perform_job(self, job, *args, **kwargs)
 
             if self.is_horse:
                 # We're inside of a forked process and RQ is
@@ -116,12 +143,20 @@ def sentry_patched_handle_exception(
 
         old_enqueue_job = Queue.enqueue_job
 
-        @ensure_integration_enabled(RqIntegration, old_enqueue_job)
         def sentry_patched_enqueue_job(
             self: "Queue", job: "Any", **kwargs: "Any"
         ) -> "Any":
+            client = sentry_sdk.get_client()
+            if client.get_integration(RqIntegration) is None:
+                return old_enqueue_job(self, job, **kwargs)
+
             scope = sentry_sdk.get_current_scope()
-            if scope.span is not None:
+            span_streaming = has_span_streaming_enabled(client.options)
+            if span_streaming and scope.streamed_span is not None:
+                job.meta["_sentry_trace_headers"] = dict(
+                    scope.iter_trace_propagation_headers()
+                )
+            elif not span_streaming and scope.span is not None:
                 job.meta["_sentry_trace_headers"] = dict(
                     scope.iter_trace_propagation_headers()
                 )