ref: Patched functions decorator for integrations

sentry_sdk/utils.py

@@ -14,7 +14,7 @@
 from copy import copy
 from datetime import datetime
 from decimal import Decimal
-from functools import partial, partialmethod
+from functools import partial, partialmethod, wraps
 from numbers import Real
 from urllib.parse import parse_qs, unquote, urlencode, urlsplit, urlunsplit
 
@@ -26,11 +26,14 @@
     BaseExceptionGroup = None  # type: ignore
 
 import sentry_sdk
+import sentry_sdk.hub
 from sentry_sdk._compat import PY37
 from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.consts import DEFAULT_MAX_VALUE_LENGTH, EndpointType
 
 if TYPE_CHECKING:
+    from collections.abc import Awaitable
+
     from types import FrameType, TracebackType
     from typing import (
         Any,
@@ -41,14 +44,20 @@
         List,
         NoReturn,
         Optional,
+        ParamSpec,
         Set,
         Tuple,
         Type,
+        TypeVar,
         Union,
     )
 
+    import sentry_sdk.integrations
     from sentry_sdk._types import Event, ExcInfo
 
+    P = ParamSpec("P")
+    R = TypeVar("R")
+
 
 epoch = datetime(1970, 1, 1)
 
@@ -1622,6 +1631,74 @@ def reraise(tp, value, tb=None):
     raise value
 
 
+def ensure_integration_enabled(
+    integration,  # type: type[sentry_sdk.integrations.Integration]
+    original_function,  # type: Callable[P, R]
+):
+    # type: (...) -> Callable[[Callable[P, R]], Callable[P, R]]
+    """
+    Ensures a given integration is enabled prior to calling a Sentry-patched function.
+
+    The function takes as its parameters the integration that must be enabled and the original
+    function that the SDK is patching. The function returns a function that takes the
+    decorated (Sentry-patched) function as its parameter, and returns a function that, when
+    called, checks whether the given integration is enabled. If the integration is enabled, the
+    function calls the decorated, Sentry-patched function. If the integration is not enabled,
+    the original function is called.
+
+    The function also takes care of preserving the original function's signature and docstring.
+
+    Example usage:
+
+    ```python
+    @ensure_integration_enabled(MyIntegration, my_function)
+    def patch_my_function():
+        with sentry_sdk.start_transaction(...):
+            return my_function()
+    ```
+    """
+
+    def patcher(sentry_patched_function):
+        # type: (Callable[P, R]) -> Callable[P, R]
+        @wraps(original_function)
+        def runner(*args: "P.args", **kwargs: "P.kwargs"):
+            # type: (...) -> R
+            if sentry_sdk.get_client().get_integration(integration) is None:
+                return original_function(*args, **kwargs)
+
+            return sentry_patched_function(*args, **kwargs)
+
+        return runner
+
+    return patcher
+
+
+def ensure_integration_enabled_async(
+    integration,  # type: type[sentry_sdk.integrations.Integration]
+    original_function,  # type: Callable[P, Awaitable[R]]
+):
+    # type: (...) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]
+    """
+    Version of `ensure_integration_enabled` for decorating async functions.
+
+    Please refer to the `ensure_integration_enabled` documentation for more information.
+    """
+
+    def patcher(sentry_patched_function):
+        # type: (Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]
+        @wraps(original_function)
+        async def runner(*args: "P.args", **kwargs: "P.kwargs"):
+            # type: (...) -> R
+            if sentry_sdk.get_client().get_integration(integration) is None:
+                return await original_function(*args, **kwargs)
+
+            return await sentry_patched_function(*args, **kwargs)
+
+        return runner
+
+    return patcher
+
+
 if PY37:
 
     def nanosecond_time():

tests/integrations/cloud_resource_context/test_cloud_resource_context.py

@@ -395,6 +395,12 @@ def test_setup_once(
                 fake_set_context.assert_not_called()
 
             if warning_called:
-                assert fake_warning.call_count == 1
+                correct_warning_found = False
+                for call in fake_warning.call_args_list:
+                    if call[0][0].startswith("Invalid value for cloud_provider:"):
+                        correct_warning_found = True
+                        break
+
+                assert correct_warning_found
             else:
                 fake_warning.assert_not_called()

tests/test_utils.py

@@ -5,6 +5,7 @@
 import pytest
 
 import sentry_sdk
+from sentry_sdk.integrations import Integration
 from sentry_sdk.utils import (
     Components,
     Dsn,
@@ -21,9 +22,21 @@
     serialize_frame,
     is_sentry_url,
     _get_installed_modules,
+    ensure_integration_enabled,
+    ensure_integration_enabled_async,
 )
 
 
+class TestIntegration(Integration):
+    """
+    Test integration for testing ensure_integration_enabled and
+    ensure_integration_enabled_async decorators.
+    """
+
+    identifier = "test"
+    setup_once = mock.MagicMock()
+
+
 def _normalize_distribution_name(name):
     # type: (str) -> str
     """Normalize distribution name according to PEP-0503.
@@ -567,3 +580,75 @@ def test_default_release_empty_string():
         release = get_default_release()
 
     assert release is None
+
+
+def test_ensure_integration_enabled_integration_enabled(sentry_init):
+    def original_function():
+        return "original"
+
+    def function_to_patch():
+        return "patched"
+
+    sentry_init(integrations=[TestIntegration()])
+
+    # Test the decorator by applying to function_to_patch
+    patched_function = ensure_integration_enabled(TestIntegration, original_function)(
+        function_to_patch
+    )
+
+    assert patched_function() == "patched"
+
+
+def test_ensure_integration_enabled_integration_disabled(sentry_init):
+    def original_function():
+        return "original"
+
+    def function_to_patch():
+        return "patched"
+
+    sentry_init(integrations=[])  # TestIntegration is disabled
+
+    # Test the decorator by applying to function_to_patch
+    patched_function = ensure_integration_enabled(TestIntegration, original_function)(
+        function_to_patch
+    )
+
+    assert patched_function() == "original"
+
+
+@pytest.mark.asyncio
+async def test_ensure_integration_enabled_async_integration_enabled(sentry_init):
+    # Setup variables and functions for the test
+    async def original_function():
+        return "original"
+
+    async def function_to_patch():
+        return "patched"
+
+    sentry_init(integrations=[TestIntegration()])
+
+    # Test the decorator by applying to function_to_patch
+    patched_function = ensure_integration_enabled_async(
+        TestIntegration, original_function
+    )(function_to_patch)
+
+    assert await patched_function() == "patched"
+
+
+@pytest.mark.asyncio
+async def test_ensure_integration_enabled_async_integration_disabled(sentry_init):
+    # Setup variables and functions for the test
+    async def original_function():
+        return "original"
+
+    async def function_to_patch():
+        return "patched"
+
+    sentry_init(integrations=[])  # TestIntegration is disabled
+
+    # Test the decorator by applying to function_to_patch
+    patched_function = ensure_integration_enabled_async(
+        TestIntegration, original_function
+    )(function_to_patch)
+
+    assert await patched_function() == "original"