ref(arq): Use new scopes API (#2878)

Author: sentrivana

sentry_sdk/integrations/arq.py

@@ -1,15 +1,16 @@
 import sys
 
+import sentry_sdk
 from sentry_sdk._types import TYPE_CHECKING
-from sentry_sdk import Hub
 from sentry_sdk.consts import OP
-from sentry_sdk.hub import _should_send_default_pii
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.integrations.logging import ignore_logger
-from sentry_sdk.scope import Scope
+from sentry_sdk.scope import Scope, should_send_default_pii
 from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_TASK
 from sentry_sdk.utils import (
     capture_internal_exceptions,
+    ensure_integration_enabled,
+    ensure_integration_enabled_async,
     event_from_exception,
     SENSITIVE_DATA_SUBSTITUTE,
     parse_version,
@@ -70,14 +71,10 @@ def patch_enqueue_job():
     # type: () -> None
     old_enqueue_job = ArqRedis.enqueue_job
 
+    @ensure_integration_enabled_async(ArqIntegration, old_enqueue_job)
     async def _sentry_enqueue_job(self, function, *args, **kwargs):
         # type: (ArqRedis, str, *Any, **Any) -> Optional[Job]
-        hub = Hub.current
-
-        if hub.get_integration(ArqIntegration) is None:
-            return await old_enqueue_job(self, function, *args, **kwargs)
-
-        with hub.start_span(op=OP.QUEUE_SUBMIT_ARQ, description=function):
+        with sentry_sdk.start_span(op=OP.QUEUE_SUBMIT_ARQ, description=function):
             return await old_enqueue_job(self, function, *args, **kwargs)
 
     ArqRedis.enqueue_job = _sentry_enqueue_job
@@ -87,14 +84,10 @@ def patch_run_job():
     # type: () -> None
     old_run_job = Worker.run_job
 
+    @ensure_integration_enabled_async(ArqIntegration, old_run_job)
     async def _sentry_run_job(self, job_id, score):
         # type: (Worker, str, int) -> None
-        hub = Hub(Hub.current)
-
-        if hub.get_integration(ArqIntegration) is None:
-            return await old_run_job(self, job_id, score)
-
-        with hub.push_scope() as scope:
+        with sentry_sdk.isolation_scope() as scope:
             scope._name = "arq"
             scope.clear_breadcrumbs()
 
@@ -105,7 +98,7 @@ async def _sentry_run_job(self, job_id, score):
                 source=TRANSACTION_SOURCE_TASK,
             )
 
-            with hub.start_transaction(transaction):
+            with sentry_sdk.start_transaction(transaction):
                 return await old_run_job(self, job_id, score)
 
     Worker.run_job = _sentry_run_job
@@ -127,7 +120,7 @@ def _capture_exception(exc_info):
         client_options=Scope.get_client().options,
         mechanism={"type": ArqIntegration.identifier, "handled": False},
     )
-    scope.capture_event(event, hint=hint)
+    sentry_sdk.capture_event(event, hint=hint)
 
 
 def _make_event_processor(ctx, *args, **kwargs):
@@ -148,10 +141,10 @@ def event_processor(event, hint):
             extra["arq-job"] = {
                 "task": ctx["job_name"],
                 "args": (
-                    args if _should_send_default_pii() else SENSITIVE_DATA_SUBSTITUTE
+                    args if should_send_default_pii() else SENSITIVE_DATA_SUBSTITUTE
                 ),
                 "kwargs": (
-                    kwargs if _should_send_default_pii() else SENSITIVE_DATA_SUBSTITUTE
+                    kwargs if should_send_default_pii() else SENSITIVE_DATA_SUBSTITUTE
                 ),
                 "retry": ctx["job_try"],
             }
@@ -163,13 +156,11 @@ def event_processor(event, hint):
 
 def _wrap_coroutine(name, coroutine):
     # type: (str, WorkerCoroutine) -> WorkerCoroutine
+
+    @ensure_integration_enabled_async(ArqIntegration, coroutine)
     async def _sentry_coroutine(ctx, *args, **kwargs):
         # type: (Dict[Any, Any], *Any, **Any) -> Any
-        hub = Hub.current
-        if hub.get_integration(ArqIntegration) is None:
-            return await coroutine(ctx, *args, **kwargs)
-
-        hub.scope.add_event_processor(
+        Scope.get_isolation_scope().add_event_processor(
             _make_event_processor({**ctx, "job_name": name}, *args, **kwargs)
         )
 
@@ -189,13 +180,9 @@ def patch_create_worker():
     # type: () -> None
     old_create_worker = arq.worker.create_worker
 
+    @ensure_integration_enabled(ArqIntegration, old_create_worker)
     def _sentry_create_worker(*args, **kwargs):
         # type: (*Any, **Any) -> Worker
-        hub = Hub.current
-
-        if hub.get_integration(ArqIntegration) is None:
-            return old_create_worker(*args, **kwargs)
-
         settings_cls = args[0]
 
         if hasattr(settings_cls, "functions"):

tests/integrations/arq/test_arq.py

@@ -1,7 +1,7 @@
 import asyncio
 import pytest
 
-from sentry_sdk import start_transaction, Hub
+from sentry_sdk import get_client, start_transaction
 from sentry_sdk.integrations.arq import ArqIntegration
 
 import arq.worker
@@ -60,7 +60,6 @@ def inner(
             integrations=[ArqIntegration()],
             traces_sample_rate=1.0,
             send_default_pii=True,
-            debug=True,
         )
 
         server = FakeRedis()
@@ -245,7 +244,7 @@ async def dummy_job(_ctx):
 
     pool, worker = init_arq([dummy_job])
     # remove the integration to trigger the edge case
-    Hub.current.client.integrations.pop("arq")
+    get_client().integrations.pop("arq")
 
     job = await pool.enqueue_job("dummy_job")