diff --git a/sentry_sdk/integrations/huey.py b/sentry_sdk/integrations/huey.py index b765491877..9b457c08d6 100644 --- a/sentry_sdk/integrations/huey.py +++ b/sentry_sdk/integrations/huey.py @@ -1,20 +1,20 @@ import sys from datetime import datetime +import sentry_sdk from sentry_sdk._types import TYPE_CHECKING -from sentry_sdk import Hub from sentry_sdk.api import continue_trace, get_baggage, get_traceparent 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.scope import Scope, should_send_default_pii from sentry_sdk.tracing import ( BAGGAGE_HEADER_NAME, SENTRY_TRACE_HEADER_NAME, TRANSACTION_SOURCE_TASK, ) -from sentry_sdk.scope import Scope from sentry_sdk.utils import ( capture_internal_exceptions, + ensure_integration_enabled, event_from_exception, SENSITIVE_DATA_SUBSTITUTE, reraise, @@ -52,14 +52,10 @@ def patch_enqueue(): # type: () -> None old_enqueue = Huey.enqueue + @ensure_integration_enabled(HueyIntegration, old_enqueue) def _sentry_enqueue(self, task): # type: (Huey, Task) -> Optional[Union[Result, ResultGroup]] - hub = Hub.current - - if hub.get_integration(HueyIntegration) is None: - return old_enqueue(self, task) - - with hub.start_span(op=OP.QUEUE_SUBMIT_HUEY, description=task.name): + with sentry_sdk.start_span(op=OP.QUEUE_SUBMIT_HUEY, description=task.name): if not isinstance(task, PeriodicTask): # Attach trace propagation data to task kwargs. We do # not do this for periodic tasks, as these don't @@ -87,12 +83,12 @@ def event_processor(event, hint): "task": task.name, "args": ( task.args - if _should_send_default_pii() + if should_send_default_pii() else SENSITIVE_DATA_SUBSTITUTE ), "kwargs": ( task.kwargs - if _should_send_default_pii() + if should_send_default_pii() else SENSITIVE_DATA_SUBSTITUTE ), "retry": (task.default_retries or 0) - task.retries, @@ -122,12 +118,10 @@ def _capture_exception(exc_info): def _wrap_task_execute(func): # type: (F) -> F + + @ensure_integration_enabled(HueyIntegration, func) def _sentry_execute(*args, **kwargs): # type: (*Any, **Any) -> Any - hub = Hub.current - if hub.get_integration(HueyIntegration) is None: - return func(*args, **kwargs) - try: result = func(*args, **kwargs) except Exception: @@ -144,14 +138,10 @@ def patch_execute(): # type: () -> None old_execute = Huey._execute + @ensure_integration_enabled(HueyIntegration, old_execute) def _sentry_execute(self, task, timestamp=None): # type: (Huey, Task, Optional[datetime]) -> Any - hub = Hub.current - - if hub.get_integration(HueyIntegration) is None: - return old_execute(self, task, timestamp) - - with hub.push_scope() as scope: + with sentry_sdk.isolation_scope() as scope: with capture_internal_exceptions(): scope._name = "huey" scope.clear_breadcrumbs() @@ -171,7 +161,7 @@ def _sentry_execute(self, task, timestamp=None): task.execute = _wrap_task_execute(task.execute) task._sentry_is_patched = True - with hub.start_transaction(transaction): + with sentry_sdk.start_transaction(transaction): return old_execute(self, task, timestamp) Huey._execute = _sentry_execute