feat(arq): Support span streaming

sentry_sdk/integrations/arq.py

@@ -1,13 +1,16 @@
 import sys
 
 import sentry_sdk
-from sentry_sdk.consts import OP, SPANSTATUS
+from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS
 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.traces import SegmentSource
 from sentry_sdk.tracing import Transaction, TransactionSource
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import (
     SENSITIVE_DATA_SUBSTITUTE,
+    _register_control_flow_exception,
     capture_internal_exceptions,
     ensure_integration_enabled,
     event_from_exception,
@@ -59,6 +62,8 @@
         patch_run_job()
         patch_create_worker()
 
+        _register_control_flow_exception(ARQ_CONTROL_FLOW_EXCEPTIONS)  # type: ignore
+
         ignore_logger("arq.worker")
 
 
@@ -69,10 +74,20 @@
     async def _sentry_enqueue_job(
         self: "ArqRedis", function: str, *args: "Any", **kwargs: "Any"
     ) -> "Optional[Job]":
-        integration = sentry_sdk.get_client().get_integration(ArqIntegration)
-        if integration is None:
+        client = sentry_sdk.get_client()
+        if client.get_integration(ArqIntegration) is None:
             return await old_enqueue_job(self, function, *args, **kwargs)
 
+        if has_span_streaming_enabled(client.options):
+            with sentry_sdk.traces.start_span(
+                name=function,
+                attributes={
+                    "sentry.op": OP.QUEUE_SUBMIT_ARQ,
+                    "sentry.origin": ArqIntegration.origin,
+                },
+            ):
+                return await old_enqueue_job(self, function, *args, **kwargs)
+
         with sentry_sdk.start_span(
             op=OP.QUEUE_SUBMIT_ARQ, name=function, origin=ArqIntegration.origin
         ):
@@ -86,14 +101,27 @@
     old_run_job = Worker.run_job
 
     async def _sentry_run_job(self: "Worker", job_id: str, score: int) -> None:
-        integration = sentry_sdk.get_client().get_integration(ArqIntegration)
-        if integration is None:
+        client = sentry_sdk.get_client()
+        if client.get_integration(ArqIntegration) is None:
             return await old_run_job(self, job_id, score)
 
         with sentry_sdk.isolation_scope() as scope:
             scope._name = "arq"
             scope.clear_breadcrumbs()
 
+            if has_span_streaming_enabled(client.options):
+                with sentry_sdk.traces.start_span(
+                    name="unknown arq task",
+                    attributes={
+                        "sentry.op": OP.QUEUE_TASK_ARQ,
+                        "sentry.origin": ArqIntegration.origin,
+                        "sentry.span.source": SegmentSource.TASK,
+                        SPANDATA.MESSAGING_MESSAGE_ID: job_id,
+                    },
+                    parent_span=None,
+                ):
+                    return await old_run_job(self, job_id, score)
+
             transaction = Transaction(
                 name="unknown arq task",
                 status="ok",
@@ -163,10 +191,19 @@
     async def _sentry_coroutine(
         ctx: "Dict[Any, Any]", *args: "Any", **kwargs: "Any"
     ) -> "Any":
-        integration = sentry_sdk.get_client().get_integration(ArqIntegration)
+        client = sentry_sdk.get_client()
+        integration = client.get_integration(ArqIntegration)
         if integration is None:
             return await coroutine(ctx, *args, **kwargs)
 
+        if has_span_streaming_enabled(client.options):
+            scope = sentry_sdk.get_current_scope()
+            span = scope.streamed_span
+            if span is not None:
+                span.name = name
+
+            scope.set_transaction_name(name)
+
         sentry_sdk.get_isolation_scope().add_event_processor(
             _make_event_processor({**ctx, "job_name": name}, *args, **kwargs)
         )