diff --git a/sentry_sdk/client.py b/sentry_sdk/client.py index 259196d1c6..e3821f48ca 100644 --- a/sentry_sdk/client.py +++ b/sentry_sdk/client.py @@ -188,6 +188,7 @@ def __init__(self, options: "Optional[Dict[str, Any]]" = None) -> None: self.monitor: "Optional[Monitor]" = None self.log_batcher: "Optional[LogBatcher]" = None self.metrics_batcher: "Optional[MetricsBatcher]" = None + self.integrations: "dict[str, Integration]" = {} def __getstate__(self, *args: "Any", **kwargs: "Any") -> "Any": return {"options": {}} diff --git a/sentry_sdk/integrations/asyncio.py b/sentry_sdk/integrations/asyncio.py index 39c7e3f879..afaca73d33 100644 --- a/sentry_sdk/integrations/asyncio.py +++ b/sentry_sdk/integrations/asyncio.py @@ -47,6 +47,10 @@ def patch_asyncio() -> None: loop = asyncio.get_running_loop() orig_task_factory = loop.get_task_factory() + # Check if already patched + if getattr(orig_task_factory, "_is_sentry_task_factory", False): + return + def _sentry_task_factory( loop: "asyncio.AbstractEventLoop", coro: "Coroutine[Any, Any, Any]", @@ -102,6 +106,7 @@ async def _task_with_sentry_span_creation() -> "Any": return task + _sentry_task_factory._is_sentry_task_factory = True # type: ignore loop.set_task_factory(_sentry_task_factory) # type: ignore except RuntimeError: @@ -138,3 +143,48 @@ class AsyncioIntegration(Integration): @staticmethod def setup_once() -> None: patch_asyncio() + + +def enable_asyncio_integration(*args: "Any", **kwargs: "Any") -> None: + """ + Enable AsyncioIntegration with the provided options. + + This is useful in scenarios where Sentry needs to be initialized before + an event loop is set up, but you still want to instrument asyncio once there + is an event loop. In that case, you can sentry_sdk.init() early on without + the AsyncioIntegration and then, once the event loop has been set up, + execute: + + ```python + from sentry_sdk.integrations.asyncio import enable_asyncio_integration + + async def async_entrypoint(): + enable_asyncio_integration() + ``` + + Any arguments provided will be passed to AsyncioIntegration() as is. + + If AsyncioIntegration has already patched the current event loop, this + function won't have any effect. + + If AsyncioIntegration was provided in + sentry_sdk.init(disabled_integrations=[...]), this function will ignore that + and the integration will be enabled. + """ + client = sentry_sdk.get_client() + if not client.is_active(): + return + + # This function purposefully bypasses the integration machinery in + # integrations/__init__.py. _installed_integrations/_processed_integrations + # is used to prevent double patching the same module, but in the case of + # the AsyncioIntegration, we don't monkeypatch the standard library directly, + # we patch the currently running event loop, and we keep the record of doing + # that on the loop itself. + logger.debug("Setting up integration asyncio") + + integration = AsyncioIntegration(*args, **kwargs) + integration.setup_once() + + if "asyncio" not in client.integrations: + client.integrations["asyncio"] = integration diff --git a/tests/integrations/asyncio/test_asyncio.py b/tests/integrations/asyncio/test_asyncio.py index 11b60fb0e1..58072b0071 100644 --- a/tests/integrations/asyncio/test_asyncio.py +++ b/tests/integrations/asyncio/test_asyncio.py @@ -7,7 +7,11 @@ import sentry_sdk from sentry_sdk.consts import OP -from sentry_sdk.integrations.asyncio import AsyncioIntegration, patch_asyncio +from sentry_sdk.integrations.asyncio import ( + AsyncioIntegration, + patch_asyncio, + enable_asyncio_integration, +) try: from contextvars import Context, ContextVar @@ -229,6 +233,7 @@ def test_patch_asyncio(mock_get_running_loop): Test that the patch_asyncio function will patch the task factory. """ mock_loop = mock_get_running_loop.return_value + mock_loop.get_task_factory.return_value._is_sentry_task_factory = False patch_asyncio() @@ -278,6 +283,7 @@ def test_sentry_task_factory_with_factory(mock_get_running_loop): # The original task factory will be mocked out here, let's retrieve the value for later orig_task_factory = mock_loop.get_task_factory.return_value + orig_task_factory._is_sentry_task_factory = False # Retieve sentry task factory (since it is an inner function within patch_asyncio) sentry_task_factory = get_sentry_task_factory(mock_get_running_loop) @@ -340,6 +346,7 @@ def test_sentry_task_factory_context_with_factory(mock_get_running_loop): # The original task factory will be mocked out here, let's retrieve the value for later orig_task_factory = mock_loop.get_task_factory.return_value + orig_task_factory._is_sentry_task_factory = False # Retieve sentry task factory (since it is an inner function within patch_asyncio) sentry_task_factory = get_sentry_task_factory(mock_get_running_loop) @@ -386,3 +393,113 @@ async def test_span_origin( assert event["contexts"]["trace"]["origin"] == "manual" assert event["spans"][0]["origin"] == "auto.function.asyncio" + + +@minimum_python_38 +@pytest.mark.asyncio +async def test_delayed_enable_integration(sentry_init, capture_events): + sentry_init(traces_sample_rate=1.0) + + assert "asyncio" not in sentry_sdk.get_client().integrations + + events = capture_events() + + with sentry_sdk.start_transaction(name="test"): + await asyncio.create_task(foo()) + + assert len(events) == 1 + (transaction,) = events + assert not transaction["spans"] + + enable_asyncio_integration() + + events = capture_events() + + assert "asyncio" in sentry_sdk.get_client().integrations + + with sentry_sdk.start_transaction(name="test"): + await asyncio.create_task(foo()) + + assert len(events) == 1 + (transaction,) = events + assert transaction["spans"] + assert transaction["spans"][0]["origin"] == "auto.function.asyncio" + + +@minimum_python_38 +@pytest.mark.asyncio +async def test_delayed_enable_integration_with_options(sentry_init, capture_events): + sentry_init(traces_sample_rate=1.0) + + assert "asyncio" not in sentry_sdk.get_client().integrations + + mock_init = MagicMock(return_value=None) + mock_setup_once = MagicMock() + with patch( + "sentry_sdk.integrations.asyncio.AsyncioIntegration.__init__", mock_init + ): + with patch( + "sentry_sdk.integrations.asyncio.AsyncioIntegration.setup_once", + mock_setup_once, + ): + enable_asyncio_integration("arg", kwarg="kwarg") + + assert "asyncio" in sentry_sdk.get_client().integrations + mock_init.assert_called_once_with("arg", kwarg="kwarg") + mock_setup_once.assert_called_once() + + +@minimum_python_38 +@pytest.mark.asyncio +async def test_delayed_enable_enabled_integration(sentry_init, uninstall_integration): + # Ensure asyncio integration is not already installed from previous tests + uninstall_integration("asyncio") + + integration = AsyncioIntegration() + sentry_init(integrations=[integration], traces_sample_rate=1.0) + + assert "asyncio" in sentry_sdk.get_client().integrations + + # Get the task factory after initial setup - it should be Sentry's + loop = asyncio.get_running_loop() + task_factory_before = loop.get_task_factory() + assert getattr(task_factory_before, "_is_sentry_task_factory", False) is True + + enable_asyncio_integration() + + assert "asyncio" in sentry_sdk.get_client().integrations + + # The task factory should be the same (loop not re-patched) + task_factory_after = loop.get_task_factory() + assert task_factory_before is task_factory_after + + +@minimum_python_38 +@pytest.mark.asyncio +async def test_delayed_enable_integration_after_disabling(sentry_init, capture_events): + sentry_init(disabled_integrations=[AsyncioIntegration()], traces_sample_rate=1.0) + + assert "asyncio" not in sentry_sdk.get_client().integrations + + events = capture_events() + + with sentry_sdk.start_transaction(name="test"): + await asyncio.create_task(foo()) + + assert len(events) == 1 + (transaction,) = events + assert not transaction["spans"] + + enable_asyncio_integration() + + events = capture_events() + + assert "asyncio" in sentry_sdk.get_client().integrations + + with sentry_sdk.start_transaction(name="test"): + await asyncio.create_task(foo()) + + assert len(events) == 1 + (transaction,) = events + assert transaction["spans"] + assert transaction["spans"][0]["origin"] == "auto.function.asyncio"