ref(integrations): Use ensure_integration_enabled decorator (#2906)

---------

Co-authored-by: Anton Pirker <[email protected]>

Author: sentrivana

sentry_sdk/integrations/aiohttp.py

@@ -190,12 +190,9 @@ def init(*args, **kwargs):
 def create_trace_config():
     # type: () -> TraceConfig
 
+    @ensure_integration_enabled_async(AioHttpIntegration)
     async def on_request_start(session, trace_config_ctx, params):
         # type: (ClientSession, SimpleNamespace, TraceRequestStartParams) -> None
-        client = sentry_sdk.get_client()
-        if client.get_integration(AioHttpIntegration) is None:
-            return
-
         method = params.method.upper()
 
         parsed_url = None
@@ -213,6 +210,8 @@ async def on_request_start(session, trace_config_ctx, params):
             span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query)
             span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment)
 
+        client = sentry_sdk.get_client()
+
         if should_propagate_trace(client, str(params.url)):
             for key, value in Scope.get_current_scope().iter_trace_propagation_headers(
                 span=span

sentry_sdk/integrations/asyncpg.py

@@ -58,14 +58,13 @@ def setup_once() -> None:
 
 
 def _wrap_execute(f: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
+    @ensure_integration_enabled_async(AsyncPGIntegration, f)
     async def _inner(*args: Any, **kwargs: Any) -> T:
-        integration = sentry_sdk.get_client().get_integration(AsyncPGIntegration)
-
         # Avoid recording calls to _execute twice.
         # Calls to Connection.execute with args also call
         # Connection._execute, which is recorded separately
         # args[0] = the connection object, args[1] is the query
-        if integration is None or len(args) > 2:
+        if len(args) > 2:
             return await f(*args, **kwargs)
 
         query = args[1]

sentry_sdk/integrations/atexit.py

@@ -6,7 +6,7 @@
 from sentry_sdk import Scope
 from sentry_sdk.utils import logger
 from sentry_sdk.integrations import Integration
-
+from sentry_sdk.utils import ensure_integration_enabled
 from sentry_sdk._types import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -44,13 +44,13 @@ def __init__(self, callback=None):
     def setup_once():
         # type: () -> None
         @atexit.register
+        @ensure_integration_enabled(AtexitIntegration)
         def _shutdown():
             # type: () -> None
             logger.debug("atexit: got shutdown signal")
             client = sentry_sdk.get_client()
             integration = client.get_integration(AtexitIntegration)
-            if integration is not None:
-                logger.debug("atexit: shutting down client")
 
-                Scope.get_isolation_scope().end_session()
-                client.close(callback=integration.callback)
+            logger.debug("atexit: shutting down client")
+            Scope.get_isolation_scope().end_session()
+            client.close(callback=integration.callback)

sentry_sdk/integrations/aws_lambda.py

@@ -11,6 +11,7 @@
 from sentry_sdk.utils import (
     AnnotatedValue,
     capture_internal_exceptions,
+    ensure_integration_enabled,
     event_from_exception,
     logger,
     TimeoutThread,
@@ -37,12 +38,10 @@
 
 def _wrap_init_error(init_error):
     # type: (F) -> F
+    @ensure_integration_enabled(AwsLambdaIntegration, init_error)
     def sentry_init_error(*args, **kwargs):
         # type: (*Any, **Any) -> Any
         client = sentry_sdk.get_client()
-        integration = client.get_integration(AwsLambdaIntegration)
-        if integration is None:
-            return init_error(*args, **kwargs)
 
         with capture_internal_exceptions():
             Scope.get_isolation_scope().clear_breadcrumbs()
@@ -63,6 +62,7 @@ def sentry_init_error(*args, **kwargs):
 
 def _wrap_handler(handler):
     # type: (F) -> F
+    @ensure_integration_enabled(AwsLambdaIntegration, handler)
     def sentry_handler(aws_event, aws_context, *args, **kwargs):
         # type: (Any, Any, *Any, **Any) -> Any
 
@@ -91,8 +91,6 @@ def sentry_handler(aws_event, aws_context, *args, **kwargs):
 
         client = sentry_sdk.get_client()
         integration = client.get_integration(AwsLambdaIntegration)
-        if integration is None:
-            return handler(aws_event, aws_context, *args, **kwargs)
 
         configured_time = aws_context.get_remaining_time_in_millis()
 

sentry_sdk/integrations/boto3.py

@@ -6,7 +6,12 @@
 from sentry_sdk.tracing import Span
 
 from sentry_sdk._types import TYPE_CHECKING
-from sentry_sdk.utils import capture_internal_exceptions, parse_url, parse_version
+from sentry_sdk.utils import (
+    capture_internal_exceptions,
+    ensure_integration_enabled,
+    parse_url,
+    parse_version,
+)
 
 if TYPE_CHECKING:
     from typing import Any
@@ -57,11 +62,9 @@ def sentry_patched_init(self, *args, **kwargs):
         BaseClient.__init__ = sentry_patched_init
 
 
+@ensure_integration_enabled(Boto3Integration)
 def _sentry_request_created(service_id, request, operation_name, **kwargs):
     # type: (str, AWSRequest, str, **Any) -> None
-    if sentry_sdk.get_client().get_integration(Boto3Integration) is None:
-        return
-
     description = "aws.%s.%s" % (service_id, operation_name)
     span = sentry_sdk.start_span(
         op=OP.HTTP_CLIENT,

sentry_sdk/integrations/bottle.py

@@ -77,11 +77,10 @@ def sentry_patched_wsgi_app(self, environ, start_response):
 
         old_handle = Bottle._handle
 
+        @ensure_integration_enabled(BottleIntegration, old_handle)
         def _patched_handle(self, environ):
             # type: (Bottle, Dict[str, Any]) -> Any
             integration = sentry_sdk.get_client().get_integration(BottleIntegration)
-            if integration is None:
-                return old_handle(self, environ)
 
             scope = Scope.get_isolation_scope()
             scope._name = "bottle"
@@ -96,13 +95,11 @@ def _patched_handle(self, environ):
 
         old_make_callback = Route._make_callback
 
+        @ensure_integration_enabled(BottleIntegration, old_make_callback)
         def patched_make_callback(self, *args, **kwargs):
             # type: (Route, *object, **object) -> Any
             client = sentry_sdk.get_client()
-            integration = client.get_integration(BottleIntegration)
             prepared_callback = old_make_callback(self, *args, **kwargs)
-            if integration is None:
-                return prepared_callback
 
             def wrapped_callback(*args, **kwargs):
                 # type: (*object, **object) -> Any

sentry_sdk/integrations/clickhouse_driver.py

@@ -4,7 +4,7 @@
 from sentry_sdk.tracing import Span
 from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.utils import capture_internal_exceptions
+from sentry_sdk.utils import capture_internal_exceptions, ensure_integration_enabled
 
 from typing import TypeVar
 
@@ -74,9 +74,8 @@ def setup_once() -> None:
 
 
 def _wrap_start(f: Callable[P, T]) -> Callable[P, T]:
+    @ensure_integration_enabled(ClickhouseDriverIntegration, f)
     def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
-        if sentry_sdk.get_client().get_integration(ClickhouseDriverIntegration) is None:
-            return f(*args, **kwargs)
         connection = args[0]
         query = args[1]
         query_id = args[2] if len(args) > 2 else kwargs.get("query_id")

sentry_sdk/integrations/django/__init__.py

@@ -392,13 +392,11 @@ def _set_transaction_name_and_source(scope, transaction_style, request):
         pass
 
 
+@ensure_integration_enabled(DjangoIntegration)
 def _before_get_response(request):
     # type: (WSGIRequest) -> None
     integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
 
-    if integration is None:
-        return
-
     _patch_drf()
 
     scope = Scope.get_current_scope()
@@ -423,10 +421,11 @@ def _attempt_resolve_again(request, scope, transaction_style):
     _set_transaction_name_and_source(scope, transaction_style, request)
 
 
+@ensure_integration_enabled(DjangoIntegration)
 def _after_get_response(request):
     # type: (WSGIRequest) -> None
     integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
-    if integration is None or integration.transaction_style != "url":
+    if integration.transaction_style != "url":
         return
 
     scope = Scope.get_current_scope()
@@ -492,21 +491,22 @@ def wsgi_request_event_processor(event, hint):
     return wsgi_request_event_processor
 
 
+@ensure_integration_enabled(DjangoIntegration)
 def _got_request_exception(request=None, **kwargs):
     # type: (WSGIRequest, **Any) -> None
     client = sentry_sdk.get_client()
     integration = client.get_integration(DjangoIntegration)
-    if integration is not None:
-        if request is not None and integration.transaction_style == "url":
-            scope = Scope.get_current_scope()
-            _attempt_resolve_again(request, scope, integration.transaction_style)
-
-        event, hint = event_from_exception(
-            sys.exc_info(),
-            client_options=client.options,
-            mechanism={"type": "django", "handled": False},
-        )
-        sentry_sdk.capture_event(event, hint=hint)
+
+    if request is not None and integration.transaction_style == "url":
+        scope = Scope.get_current_scope()
+        _attempt_resolve_again(request, scope, integration.transaction_style)
+
+    event, hint = event_from_exception(
+        sys.exc_info(),
+        client_options=client.options,
+        mechanism={"type": "django", "handled": False},
+    )
+    sentry_sdk.capture_event(event, hint=hint)
 
 
 class DjangoRequestExtractor(RequestExtractor):

sentry_sdk/integrations/django/templates.py

@@ -65,11 +65,9 @@ def patch_templates():
     real_rendered_content = SimpleTemplateResponse.rendered_content
 
     @property  # type: ignore
+    @ensure_integration_enabled(DjangoIntegration, real_rendered_content.fget)
     def rendered_content(self):
         # type: (SimpleTemplateResponse) -> str
-        if sentry_sdk.get_client().get_integration(DjangoIntegration) is None:
-            return real_rendered_content.fget(self)
-
         with sentry_sdk.start_span(
             op=OP.TEMPLATE_RENDER,
             description=_get_template_name_description(self.template_name),

sentry_sdk/integrations/excepthook.py

@@ -1,7 +1,11 @@
 import sys
 
 import sentry_sdk
-from sentry_sdk.utils import capture_internal_exceptions, event_from_exception
+from sentry_sdk.utils import (
+    capture_internal_exceptions,
+    ensure_integration_enabled,
+    event_from_exception,
+)
 from sentry_sdk.integrations import Integration
 
 from sentry_sdk._types import TYPE_CHECKING
@@ -43,11 +47,12 @@ def setup_once():
 
 def _make_excepthook(old_excepthook):
     # type: (Excepthook) -> Excepthook
+    @ensure_integration_enabled(ExcepthookIntegration, old_excepthook)
     def sentry_sdk_excepthook(type_, value, traceback):
         # type: (Type[BaseException], BaseException, Optional[TracebackType]) -> None
         integration = sentry_sdk.get_client().get_integration(ExcepthookIntegration)
 
-        if integration is not None and _should_send(integration.always_run):
+        if _should_send(integration.always_run):
             with capture_internal_exceptions():
                 event, hint = event_from_exception(
                     (type_, value, traceback),