feat(aiohttp): add instrumentation of client requests (#1761)

Author: md384

sentry_sdk/integrations/aiohttp.py

@@ -3,7 +3,7 @@
 
 from sentry_sdk.api import continue_trace
 from sentry_sdk._compat import reraise
-from sentry_sdk.consts import OP
+from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.hub import Hub
 from sentry_sdk.integrations import Integration, DidNotEnable
 from sentry_sdk.integrations.logging import ignore_logger
@@ -13,20 +13,25 @@
     request_body_within_bounds,
 )
 from sentry_sdk.tracing import SOURCE_FOR_STYLE, TRANSACTION_SOURCE_ROUTE
+from sentry_sdk.tracing_utils import should_propagate_trace
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
+    logger,
+    parse_url,
     parse_version,
     transaction_from_function,
     HAS_REAL_CONTEXTVARS,
     CONTEXTVARS_ERROR_MESSAGE,
+    SENSITIVE_DATA_SUBSTITUTE,
     AnnotatedValue,
 )
 
 try:
     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 +41,8 @@
 if TYPE_CHECKING:
     from aiohttp.web_request import Request
     from aiohttp.abc import AbstractMatchInfo
+    from aiohttp import TraceRequestStartParams, TraceRequestEndParams
+    from types import SimpleNamespace
     from typing import Any
     from typing import Dict
     from typing import Optional
@@ -164,6 +171,76 @@ async def sentry_urldispatcher_resolve(self, request):
 
         UrlDispatcher.resolve = sentry_urldispatcher_resolve
 
+        old_client_session_init = ClientSession.__init__
+
+        def init(*args, **kwargs):
+            # type: (Any, Any) -> ClientSession
+            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 create_trace_config():
+    # type: () -> TraceConfig
+    async def on_request_start(session, trace_config_ctx, params):
+        # type: (ClientSession, SimpleNamespace, TraceRequestStartParams) -> None
+        hub = Hub.current
+        if hub.get_integration(AioHttpIntegration) is None:
+            return
+
+        method = params.method.upper()
+
+        parsed_url = None
+        with capture_internal_exceptions():
+            parsed_url = parse_url(str(params.url), sanitize=False)
+
+        span = hub.start_span(
+            op=OP.HTTP_CLIENT,
+            description="%s %s"
+            % (method, parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE),
+        )
+        span.set_data(SPANDATA.HTTP_METHOD, method)
+        span.set_data("url", parsed_url.url)
+        span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query)
+        span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment)
+
+        if should_propagate_trace(hub, str(params.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, trace_config_ctx, params):
+        # type: (ClientSession, SimpleNamespace, TraceRequestEndParams) -> None
+        if trace_config_ctx.span is None:
+            return
+
+        span = trace_config_ctx.span
+        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
+
 
 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,8 @@ async def hello(request):
         "Accept-Encoding": "gzip, deflate",
         "Host": host,
         "User-Agent": request["headers"]["User-Agent"],
+        "baggage": mock.ANY,
+        "sentry-trace": mock.ANY,
     }
 
 
@@ -372,11 +374,13 @@ async def hello(request):
 
     events = capture_events()
 
-    trace_id = "582b43a4192642f0b136d5159a501701"
-    sentry_trace_header = "{}-{}-{}".format(trace_id, "6e8f22c393e68f19", 1)
-
+    # The aiohttp_client is instrumented so will generate the sentry-trace header and add request.
+    # Get the sentry-trace header from the request so we can later compare with transaction events.
     client = await aiohttp_client(app)
-    resp = await client.get("/", headers={"sentry-trace": sentry_trace_header})
+    resp = await client.get("/")
+    sentry_trace_header = resp.request_info.headers.get("sentry-trace")
+    trace_id = sentry_trace_header.split("-")[0]
+
     assert resp.status == 500
 
     msg_event, error_event, transaction_event = events
@@ -410,11 +414,13 @@ async def hello(request):
 
     events = capture_events()
 
-    trace_id = "582b43a4192642f0b136d5159a501701"
-    sentry_trace_header = "{}-{}-{}".format(trace_id, "6e8f22c393e68f19", 1)
-
+    # The aiohttp_client is instrumented so will generate the sentry-trace header and add request.
+    # Get the sentry-trace header from the request so we can later compare with transaction events.
     client = await aiohttp_client(app)
-    resp = await client.get("/", headers={"sentry-trace": sentry_trace_header})
+    resp = await client.get("/")
+    sentry_trace_header = resp.request_info.headers.get("sentry-trace")
+    trace_id = sentry_trace_header.split("-")[0]
+
     assert resp.status == 500
 
     msg_event, error_event = events
@@ -427,3 +433,75 @@ async def hello(request):
 
     assert msg_event["contexts"]["trace"]["trace_id"] == trace_id
     assert error_event["contexts"]["trace"]["trace_id"] == trace_id
+
+
+@pytest.mark.asyncio
+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),
+            "http.fragment": "",
+            "http.method": "GET",
+            "http.query": "",
+            "http.response.status_code": 200,
+            "reason": "OK",
+            "extra": "foo",
+        }
+
+
+@pytest.mark.asyncio
+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,
+        )