ref: Patched functions decorator for integrations

sentry_sdk/_types.py

@@ -9,6 +9,7 @@
 
 
 if TYPE_CHECKING:
+    from collections.abc import Awaitable
     from types import TracebackType
     from typing import Any
     from typing import Callable
@@ -116,3 +117,6 @@
     FlushedMetricValue = Union[int, float]
 
     BucketKey = Tuple[MetricType, str, MeasurementUnit, MetricTagsInternal]
+
+    GenericCallable = Callable[..., Any]
+    GenericAsyncCallable = Callable[..., Awaitable[Any]]

sentry_sdk/integrations/gql.py

@@ -1,4 +1,4 @@
-from sentry_sdk.utils import event_from_exception, parse_version
+from sentry_sdk.utils import event_from_exception, integration_patched, parse_version
 from sentry_sdk.hub import Hub, _should_send_default_pii
 from sentry_sdk.integrations import DidNotEnable, Integration
 
@@ -85,13 +85,11 @@
     # type: () -> None
     real_execute = gql.Client.execute
 
+    @integration_patched(original_function=real_execute, integration=GQLIntegration)
     def sentry_patched_execute(self, document, *args, **kwargs):
         # type: (gql.Client, DocumentNode, Any, Any) -> Any
         hub = Hub.current
-        if hub.get_integration(GQLIntegration) is None:
-            return real_execute(self, document, *args, **kwargs)
-
-        with Hub.current.configure_scope() as scope:
+        with hub.configure_scope() as scope:
             scope.add_event_processor(_make_gql_event_processor(self, document))
 
         try:

sentry_sdk/integrations/starlette.py

@@ -27,6 +27,7 @@
     parse_version,
     transaction_from_function,
 )
+from sentry_sdk.utils_py3 import integration_patched_async
 
 if TYPE_CHECKING:
     from typing import Any, Awaitable, Callable, Dict, Optional, Tuple
@@ -103,61 +104,60 @@ def _enable_span_for_middleware(middleware_class):
     # type: (Any) -> type
     old_call = middleware_class.__call__
 
+    @integration_patched_async(
+        original_function=old_call, integration=StarletteIntegration
+    )
     async def _create_span_call(app, scope, receive, send, **kwargs):
         # type: (Any, Dict[str, Any], Callable[[], Awaitable[Dict[str, Any]]], Callable[[Dict[str, Any]], Awaitable[None]], Any) -> None
         hub = Hub.current
         integration = hub.get_integration(StarletteIntegration)
-        if integration is not None:
-            middleware_name = app.__class__.__name__
-
-            # Update transaction name with middleware name
-            with hub.configure_scope() as sentry_scope:
-                name, source = _get_transaction_from_middleware(app, scope, integration)
-                if name is not None:
-                    sentry_scope.set_transaction_name(
-                        name,
-                        source=source,
-                    )
+        middleware_name = app.__class__.__name__
+
+        # Update transaction name with middleware name
+        with hub.configure_scope() as sentry_scope:
+            name, source = _get_transaction_from_middleware(app, scope, integration)
+            if name is not None:
+                sentry_scope.set_transaction_name(
+                    name,
+                    source=source,
+                )
 
-            with hub.start_span(
-                op=OP.MIDDLEWARE_STARLETTE, description=middleware_name
-            ) as middleware_span:
-                middleware_span.set_tag("starlette.middleware_name", middleware_name)
-
-                # Creating spans for the "receive" callback
-                async def _sentry_receive(*args, **kwargs):
-                    # type: (*Any, **Any) -> Any
-                    hub = Hub.current
-                    with hub.start_span(
-                        op=OP.MIDDLEWARE_STARLETTE_RECEIVE,
-                        description=getattr(receive, "__qualname__", str(receive)),
-                    ) as span:
-                        span.set_tag("starlette.middleware_name", middleware_name)
-                        return await receive(*args, **kwargs)
-
-                receive_name = getattr(receive, "__name__", str(receive))
-                receive_patched = receive_name == "_sentry_receive"
-                new_receive = _sentry_receive if not receive_patched else receive
-
-                # Creating spans for the "send" callback
-                async def _sentry_send(*args, **kwargs):
-                    # type: (*Any, **Any) -> Any
-                    hub = Hub.current
-                    with hub.start_span(
-                        op=OP.MIDDLEWARE_STARLETTE_SEND,
-                        description=getattr(send, "__qualname__", str(send)),
-                    ) as span:
-                        span.set_tag("starlette.middleware_name", middleware_name)
-                        return await send(*args, **kwargs)
-
-                send_name = getattr(send, "__name__", str(send))
-                send_patched = send_name == "_sentry_send"
-                new_send = _sentry_send if not send_patched else send
-
-                return await old_call(app, scope, new_receive, new_send, **kwargs)
+        with hub.start_span(
+            op=OP.MIDDLEWARE_STARLETTE, description=middleware_name
+        ) as middleware_span:
+            middleware_span.set_tag("starlette.middleware_name", middleware_name)
 
-        else:
-            return await old_call(app, scope, receive, send, **kwargs)
+            # Creating spans for the "receive" callback
+            async def _sentry_receive(*args, **kwargs):
+                # type: (*Any, **Any) -> Any
+                hub = Hub.current
+                with hub.start_span(
+                    op=OP.MIDDLEWARE_STARLETTE_RECEIVE,
+                    description=getattr(receive, "__qualname__", str(receive)),
+                ) as span:
+                    span.set_tag("starlette.middleware_name", middleware_name)
+                    return await receive(*args, **kwargs)
+
+            receive_name = getattr(receive, "__name__", str(receive))
+            receive_patched = receive_name == "_sentry_receive"
+            new_receive = _sentry_receive if not receive_patched else receive
+
+            # Creating spans for the "send" callback
+            async def _sentry_send(*args, **kwargs):
+                # type: (*Any, **Any) -> Any
+                hub = Hub.current
+                with hub.start_span(
+                    op=OP.MIDDLEWARE_STARLETTE_SEND,
+                    description=getattr(send, "__qualname__", str(send)),
+                ) as span:
+                    span.set_tag("starlette.middleware_name", middleware_name)
+                    return await send(*args, **kwargs)
+
+            send_name = getattr(send, "__name__", str(send))
+            send_patched = send_name == "_sentry_send"
+            new_send = _sentry_send if not send_patched else send
+
+            return await old_call(app, scope, new_receive, new_send, **kwargs)
 
     not_yet_patched = old_call.__name__ not in [
         "_create_span_call",

sentry_sdk/utils.py

@@ -12,6 +12,7 @@
 from collections import namedtuple
 from copy import copy
 from decimal import Decimal
+from functools import wraps
 from numbers import Real
 
 try:
@@ -48,6 +49,7 @@
     _PARTIALMETHOD_AVAILABLE = False
 
 import sentry_sdk
+import sentry_sdk.hub
 from sentry_sdk._compat import PY2, PY33, PY37, implements_str, text_type, urlparse
 from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.consts import DEFAULT_MAX_VALUE_LENGTH
@@ -67,8 +69,9 @@
         Type,
         Union,
     )
+    from sentry_sdk.integrations import Integration
 
-    from sentry_sdk._types import EndpointType, ExcInfo
+    from sentry_sdk._types import EndpointType, ExcInfo, GenericCallable
 
 
 epoch = datetime(1970, 1, 1)
@@ -1563,6 +1566,26 @@
     return release_tuple
 
 
+def integration_patched(original_function, integration=None):
+    # type: (GenericCallable, Optional[type[Integration]]) -> Callable[[GenericCallable], GenericCallable]
+    def patcher(sentry_patched_function):
+        # type: (GenericCallable) -> GenericCallable
+        @wraps(original_function)
+        def runner(*args, **kwargs):
+            # type: (*object, **object) -> object
+            if (
+                integration is not None
+                and sentry_sdk.hub.Hub.current.get_integration(integration) is None
+            ):
+                return original_function(*args, **kwargs)
+
+            return sentry_patched_function(*args, **kwargs)
+
+        return runner
+
+    return patcher
+
+
 if PY37:
 
     def nanosecond_time():

sentry_sdk/utils_py3.py

@@ -0,0 +1,31 @@
+from sentry_sdk.hub import Hub
+from functools import wraps
+
+from sentry_sdk._types import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Optional
+    from sentry_sdk.integrations import Integration
+
+    from sentry_sdk._types import GenericAsyncCallable
+
+
+def integration_patched_async(original_function, integration=None):
+    # type: (GenericAsyncCallable, Optional[type[Integration]]) -> Callable[[GenericAsyncCallable], GenericAsyncCallable]
+    def patcher(sentry_patched_function):
+        # type: (GenericAsyncCallable) -> GenericAsyncCallable
+        @wraps(original_function)
+        async def runner(*args, **kwargs):
+            # type: (*object, **object) -> object
+            if (
+                integration is not None
+                and Hub.current.get_integration(integration) is None
+            ):
+                return await original_function(*args, **kwargs)
+
+            return await sentry_patched_function(*args, **kwargs)
+
+        return runner
+
+    return patcher