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

Author: sentrivana

sentry_sdk/integrations/__init__.py

@@ -2,6 +2,7 @@
 from threading import Lock
 from typing import TYPE_CHECKING
 
+import sentry_sdk
 from sentry_sdk.utils import logger
 
 if TYPE_CHECKING:
@@ -279,6 +280,24 @@ def setup_integrations(
     return integrations
 
 
+def _enable_integration(integration: "Integration") -> "Optional[Integration]":
+    identifier = integration.identifier
+    client = sentry_sdk.get_client()
+
+    with _installer_lock:
+        logger.debug("Setting up integration %s", identifier)
+        try:
+            type(integration).setup_once()
+            integration.setup_once_with_options(client.options)
+        except DidNotEnable as e:
+            logger.debug("Did not enable integration %s: %s", identifier, e)
+        else:
+            _installed_integrations.add(identifier)
+            return integration
+
+        _processed_integrations.add(identifier)
+
+
 def _check_minimum_version(
     integration: "type[Integration]",
     version: "Optional[tuple[int, ...]]",

sentry_sdk/integrations/asyncio.py

@@ -3,7 +3,7 @@
 
 import sentry_sdk
 from sentry_sdk.consts import OP
-from sentry_sdk.integrations import Integration, DidNotEnable
+from sentry_sdk.integrations import Integration, DidNotEnable, _enable_integration
 from sentry_sdk.utils import event_from_exception, logger, reraise
 
 try:
@@ -138,3 +138,40 @@ 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.
+
+    The options need to correspond to the options currently accepted by the
+    AsyncioIntegration() constructor.
+
+    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()
+    ```
+
+    If AsyncioIntegration is already enabled (e.g. because it was provided in
+    sentry_sdk.init(integrations=[...])), this function will re-enable it.
+
+    If AsyncioIntegration was provided in
+    sentry_sdk.init(disabled_integrations=[...]), this function will ignore that
+    and enable it.
+    """
+    client = sentry_sdk.get_client()
+    if not client.is_active():
+        return
+
+    integration = _enable_integration(AsyncioIntegration(*args, **kwargs))
+    if integration is None:
+        return
+
+    client.integrations[integration.identifier] = 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
@@ -386,3 +390,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
+
+    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"]
+
+    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):
+    sentry_init(integrations=[AsyncioIntegration()], traces_sample_rate=1.0)
+
+    assert "asyncio" in sentry_sdk.get_client().integrations
+
+    original_integration = sentry_sdk.get_client().integrations["asyncio"]
+    enable_asyncio_integration()
+
+    assert "asyncio" in sentry_sdk.get_client().integrations
+
+    # The new asyncio integration should override the old one
+    assert sentry_sdk.get_client().integrations["asyncio"] is not original_integration
+
+
+@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"