diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py index 3be5fe6779..a0d0184a72 100644 --- a/sentry_sdk/consts.py +++ b/sentry_sdk/consts.py @@ -111,6 +111,7 @@ class OP: DB = "db" DB_REDIS = "db.redis" EVENT_DJANGO = "event.django" + FUNCTION = "function" FUNCTION_AWS = "function.aws" FUNCTION_GCP = "function.gcp" HTTP_CLIENT = "http.client" diff --git a/sentry_sdk/integrations/asyncio.py b/sentry_sdk/integrations/asyncio.py new file mode 100644 index 0000000000..ab07ffc3cb --- /dev/null +++ b/sentry_sdk/integrations/asyncio.py @@ -0,0 +1,64 @@ +from __future__ import absolute_import + +from sentry_sdk.consts import OP +from sentry_sdk.hub import Hub +from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk._types import MYPY + +try: + import asyncio + from asyncio.tasks import Task +except ImportError: + raise DidNotEnable("asyncio not available") + + +if MYPY: + from typing import Any + + +def _sentry_task_factory(loop, coro): + # type: (Any, Any) -> Task[None] + + async def _coro_creating_hub_and_span(): + # type: () -> None + hub = Hub(Hub.current) + with hub: + with hub.start_span(op=OP.FUNCTION, description=coro.__qualname__): + await coro + + # Trying to use user set task factory (if there is one) + orig_factory = loop.get_task_factory() + if orig_factory: + return orig_factory(loop, _coro_creating_hub_and_span) + + # The default task factory in `asyncio` does not have its own function + # but is just a couple of lines in `asyncio.base_events.create_task()` + # Those lines are copied here. + + # WARNING: + # If the default behavior of the task creation in asyncio changes, + # this will break! + task = Task(_coro_creating_hub_and_span, loop=loop) # type: ignore + if task._source_traceback: # type: ignore + del task._source_traceback[-1] # type: ignore + + return task + + +def patch_asyncio(): + # type: () -> None + try: + loop = asyncio.get_running_loop() + loop.set_task_factory(_sentry_task_factory) + except RuntimeError: + # When there is no running loop, we have nothing to patch. + pass + + +class AsyncioIntegration(Integration): + identifier = "asyncio" + + @staticmethod + def setup_once(): + # type: () -> None + patch_asyncio() diff --git a/tests/integrations/asyncio/__init__.py b/tests/integrations/asyncio/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/integrations/asyncio/test_asyncio.py b/tests/integrations/asyncio/test_asyncio.py new file mode 100644 index 0000000000..2e0643c4d2 --- /dev/null +++ b/tests/integrations/asyncio/test_asyncio.py @@ -0,0 +1,118 @@ +import asyncio +import sys + +import pytest +import pytest_asyncio + +import sentry_sdk +from sentry_sdk.consts import OP +from sentry_sdk.integrations.asyncio import AsyncioIntegration + + +minimum_python_36 = pytest.mark.skipif( + sys.version_info < (3, 6), reason="ASGI is only supported in Python >= 3.6" +) + + +async def foo(): + await asyncio.sleep(0.01) + + +async def bar(): + await asyncio.sleep(0.01) + + +@pytest_asyncio.fixture(scope="session") +def event_loop(request): + """Create an instance of the default event loop for each test case.""" + loop = asyncio.get_event_loop_policy().new_event_loop() + yield loop + loop.close() + + +@minimum_python_36 +@pytest.mark.asyncio +async def test_create_task( + sentry_init, + capture_events, + event_loop, +): + sentry_init( + traces_sample_rate=1.0, + send_default_pii=True, + debug=True, + integrations=[ + AsyncioIntegration(), + ], + ) + + events = capture_events() + + with sentry_sdk.start_transaction(name="test_transaction_for_create_task"): + with sentry_sdk.start_span(op="root", description="not so important"): + tasks = [event_loop.create_task(foo()), event_loop.create_task(bar())] + await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION) + + sentry_sdk.flush() + + (transaction_event,) = events + + assert transaction_event["spans"][0]["op"] == "root" + assert transaction_event["spans"][0]["description"] == "not so important" + + assert transaction_event["spans"][1]["op"] == OP.FUNCTION + assert transaction_event["spans"][1]["description"] == "foo" + assert ( + transaction_event["spans"][1]["parent_span_id"] + == transaction_event["spans"][0]["span_id"] + ) + + assert transaction_event["spans"][2]["op"] == OP.FUNCTION + assert transaction_event["spans"][2]["description"] == "bar" + assert ( + transaction_event["spans"][2]["parent_span_id"] + == transaction_event["spans"][0]["span_id"] + ) + + +@minimum_python_36 +@pytest.mark.asyncio +async def test_gather( + sentry_init, + capture_events, +): + sentry_init( + traces_sample_rate=1.0, + send_default_pii=True, + debug=True, + integrations=[ + AsyncioIntegration(), + ], + ) + + events = capture_events() + + with sentry_sdk.start_transaction(name="test_transaction_for_gather"): + with sentry_sdk.start_span(op="root", description="not so important"): + await asyncio.gather(foo(), bar(), return_exceptions=True) + + sentry_sdk.flush() + + (transaction_event,) = events + + assert transaction_event["spans"][0]["op"] == "root" + assert transaction_event["spans"][0]["description"] == "not so important" + + assert transaction_event["spans"][1]["op"] == OP.FUNCTION + assert transaction_event["spans"][1]["description"] == "foo" + assert ( + transaction_event["spans"][1]["parent_span_id"] + == transaction_event["spans"][0]["span_id"] + ) + + assert transaction_event["spans"][2]["op"] == OP.FUNCTION + assert transaction_event["spans"][2]["description"] == "bar" + assert ( + transaction_event["spans"][2]["parent_span_id"] + == transaction_event["spans"][0]["span_id"] + )