feat(aiohttp): add instrumentation of client requests

Author: md384

sentry_sdk/integrations/aiohttp.py

@@ -16,6 +16,7 @@
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
+    logger,
     parse_version,
     transaction_from_function,
     HAS_REAL_CONTEXTVARS,
@@ -27,6 +28,7 @@
     import asyncio
 
     from aiohttp import __version__ as AIOHTTP_VERSION
+    from aiohttp import ClientSession, TraceConfig
     from aiohttp.web import Application, HTTPException, UrlDispatcher
 except ImportError:
     raise DidNotEnable("AIOHTTP not installed")
@@ -36,6 +38,9 @@
 if TYPE_CHECKING:
     from aiohttp.web_request import Request
     from aiohttp.abc import AbstractMatchInfo
+    from aiohttp import TraceRequestStartParams
+    from aiohttp import TraceRequestEndParams
+    from types import SimpleNamespace
     from typing import Any
     from typing import Dict
     from typing import Optional
@@ -50,6 +55,57 @@
 TRANSACTION_STYLE_VALUES = ("handler_name", "method_and_path_pattern")
 
 
+def create_trace_config() -> TraceConfig:
+    async def on_request_start(
+        session: ClientSession,
+        trace_config_ctx: "SimpleNamespace",
+        params: "TraceRequestStartParams",
+    ):
+        hub = Hub.current
+        if hub.get_integration(AioHttpIntegration) is None:
+            return
+
+        method = params.method.upper()
+        request_url = str(params.url)
+
+        span = hub.start_span(
+            op=OP.HTTP_CLIENT, description="%s %s" % (method, request_url)
+        )
+        span.set_data("method", method)
+        span.set_data("url", request_url)
+
+        for key, value in hub.iter_trace_propagation_headers(span):
+            logger.debug(
+                "[Tracing] Adding `{key}` header {value} to outgoing request to {url}.".format(
+                    key=key, value=value, url=params.url
+                )
+            )
+            params.headers[key] = value
+
+        trace_config_ctx.span = span
+
+    async def on_request_end(
+        session: ClientSession,
+        trace_config_ctx: "SimpleNamespace",
+        params: "TraceRequestEndParams",
+    ):
+        if trace_config_ctx.span is None:
+            return
+
+        span = trace_config_ctx.span
+        span.set_data("status_code", int(params.response.status))
+        span.set_http_status(int(params.response.status))
+        span.set_data("reason", params.response.reason)
+        span.finish()
+
+    trace_config = TraceConfig()
+
+    trace_config.on_request_start.append(on_request_start)
+    trace_config.on_request_end.append(on_request_end)
+
+    return trace_config
+
+
 class AioHttpIntegration(Integration):
     identifier = "aiohttp"
 
@@ -164,6 +220,22 @@ async def sentry_urldispatcher_resolve(self, request):
 
         UrlDispatcher.resolve = sentry_urldispatcher_resolve
 
+        old_client_session_init = ClientSession.__init__
+
+        def init(*args, **kwargs):
+            hub = Hub.current
+            if hub.get_integration(AioHttpIntegration) is None:
+                return old_client_session_init(*args, **kwargs)
+
+            client_trace_configs = list(kwargs.get("trace_configs", ()))
+            trace_config = create_trace_config()
+            client_trace_configs.append(trace_config)
+
+            kwargs["trace_configs"] = client_trace_configs
+            return old_client_session_init(*args, **kwargs)
+
+        ClientSession.__init__ = init
+
 
 def _make_request_processor(weak_request):
     # type: (Callable[[], Request]) -> EventProcessor

tests/integrations/aiohttp/test_aiohttp.py

@@ -7,7 +7,7 @@
 from aiohttp.client import ServerDisconnectedError
 from aiohttp.web_request import Request
 
-from sentry_sdk import capture_message
+from sentry_sdk import capture_message, start_transaction
 from sentry_sdk.integrations.aiohttp import AioHttpIntegration
 
 try:
@@ -54,6 +54,7 @@ async def hello(request):
         "Accept-Encoding": "gzip, deflate",
         "Host": host,
         "User-Agent": request["headers"]["User-Agent"],
+        "sentry-trace": mock.ANY,
     }
 
 
@@ -427,3 +428,71 @@ async def hello(request):
 
     assert msg_event["contexts"]["trace"]["trace_id"] == trace_id
     assert error_event["contexts"]["trace"]["trace_id"] == trace_id
+
+
+async def test_crumb_capture(
+    sentry_init, aiohttp_raw_server, aiohttp_client, loop, capture_events
+):
+    def before_breadcrumb(crumb, hint):
+        crumb["data"]["extra"] = "foo"
+        return crumb
+
+    sentry_init(
+        integrations=[AioHttpIntegration()], before_breadcrumb=before_breadcrumb
+    )
+
+    async def handler(request):
+        return web.Response(text="OK")
+
+    raw_server = await aiohttp_raw_server(handler)
+
+    with start_transaction():
+        events = capture_events()
+
+        client = await aiohttp_client(raw_server)
+        resp = await client.get("/")
+        assert resp.status == 200
+        capture_message("Testing!")
+
+        (event,) = events
+
+        crumb = event["breadcrumbs"]["values"][0]
+        assert crumb["type"] == "http"
+        assert crumb["category"] == "httplib"
+        assert crumb["data"] == {
+            "url": "http://127.0.0.1:{}/".format(raw_server.port),
+            "method": "GET",
+            "status_code": 200,
+            "reason": "OK",
+            "extra": "foo",
+        }
+
+
+async def test_outgoing_trace_headers(sentry_init, aiohttp_raw_server, aiohttp_client):
+    sentry_init(
+        integrations=[AioHttpIntegration()],
+        traces_sample_rate=1.0,
+    )
+
+    async def handler(request):
+        return web.Response(text="OK")
+
+    raw_server = await aiohttp_raw_server(handler)
+
+    with start_transaction(
+        name="/interactions/other-dogs/new-dog",
+        op="greeting.sniff",
+        # make trace_id difference between transactions
+        trace_id="0123456789012345678901234567890",
+    ) as transaction:
+        client = await aiohttp_client(raw_server)
+        resp = await client.get("/")
+        request_span = transaction._span_recorder.spans[-1]
+
+        assert resp.request_info.headers[
+            "sentry-trace"
+        ] == "{trace_id}-{parent_span_id}-{sampled}".format(
+            trace_id=transaction.trace_id,
+            parent_span_id=request_span.span_id,
+            sampled=1,
+        )