feat(asyncio): Add on-demand way to enable AsyncioIntegration (#5288)

Initializing the SDK in an async environment can be a challenge, since
the requirement to initialize as early as possible might clash with the
requirement to initialize when there's an event loop running, if one
wants to use the AsyncioIntegration. These sort of scenarios would
benefit from a dedicated way to activate the integration after `init()`
has been called.


Supersedes the needlessly more general
#5285

Author: sentrivana

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": {}}

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

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"