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/__init__.py b/sentry_sdk/integrations/__init__.py index 5ab181df25..8c85d5a193 100644 --- a/sentry_sdk/integrations/__init__.py +++ b/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,28 @@ def setup_integrations( return integrations +def _enable_integration(integration: "Integration") -> "Optional[Integration]": + identifier = integration.identifier + client = sentry_sdk.get_client() + + with _installer_lock: + if identifier in client.integrations: + logger.debug("Integration already enabled: %s", identifier) + return None + + logger.debug("Setting up integration %s", identifier) + _processed_integrations.add(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) + return None + else: + _installed_integrations.add(identifier) + return integration + + def _check_minimum_version( integration: "type[Integration]", version: "Optional[tuple[int, ...]]", diff --git a/sentry_sdk/integrations/asyncio.py b/sentry_sdk/integrations/asyncio.py index 39c7e3f879..43cc30cec8 100644 --- a/sentry_sdk/integrations/asyncio.py +++ b/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,39 @@ 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 is already enabled (e.g. because it was provided in + sentry_sdk.init(integrations=[...])), 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 + + integration = _enable_integration(AsyncioIntegration(*args, **kwargs)) + if integration is None: + return + + client.integrations[integration.identifier] = integration diff --git a/tests/integrations/asyncio/test_asyncio.py b/tests/integrations/asyncio/test_asyncio.py index 11b60fb0e1..5be675402f 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 @@ -386,3 +390,104 @@ 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): + integration = AsyncioIntegration() + sentry_init(integrations=[integration], traces_sample_rate=1.0) + + assert "asyncio" in sentry_sdk.get_client().integrations + + enable_asyncio_integration() + + assert "asyncio" in sentry_sdk.get_client().integrations + + # The new asyncio integration should not override the old one + assert sentry_sdk.get_client().integrations["asyncio"] == 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"