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)
  )