ref: Patched functions decorator for integrations (getsentry#2454)

szokeasaurusrex · web-flow · commit 7a2c153ca8f7 · 2024-03-19T09:54:57.000+01:00
This commit introduces two new decorators in sentry_sdk.utils that we can use in our integrations to automate the checks for whether the integration is still enabled. Since these decorators use the new scopes API, adopting these decorators may simplify the change to the new Scopes API.
diff --git a/sentry_sdk/utils.py b/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():
diff --git a/tests/integrations/cloud_resource_context/test_cloud_resource_context.py b/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()
diff --git a/tests/test_utils.py b/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"
