Aiohttp response capture

docs/integrations/http-clients/aiohttp.md

@@ -42,6 +42,125 @@ if __name__ == "__main__":
 
 The keyword arguments of `logfire.instrument_aiohttp_client()` are passed to the `AioHttpClientInstrumentor().instrument()` method of the OpenTelemetry aiohttp client Instrumentation package, read more about it [here][opentelemetry-aiohttp].
 
+## Configuration
+
+The `logfire.instrument_aiohttp_client()` method accepts various parameters to configure the instrumentation.
+
+!!! note
+    The aiohttp client instrumentation captures request and response headers, and response bodies. Request bodies are not captured.
+
+
+### Capture HTTP Headers
+
+By default, **Logfire** doesn't capture HTTP headers. You can enable capturing both request and response headers by setting the `capture_headers` parameter to `True`.
+
+```py
+import aiohttp
+import logfire
+
+logfire.configure()
+logfire.instrument_aiohttp_client(capture_headers=True)
+
+async def main():
+    async with aiohttp.ClientSession() as session:
+        await session.get("https://httpbin.org/get")
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(main())
+```
+
+#### Capture Only Request Headers
+
+Instead of capturing both request and response headers, you can create a request hook to capture only the request headers:
+
+```py
+import aiohttp
+import logfire
+from aiohttp.tracing import TraceRequestStartParams
+from opentelemetry.trace import Span
+
+
+def capture_request_headers(span: Span, request: TraceRequestStartParams):
+    headers = request.headers
+    span.set_attributes(
+        {
+            f'http.request.header.{header_name}': headers.getall(header_name)
+            for header_name in headers.keys()
+        }
+    )
+
+
+logfire.configure()
+logfire.instrument_aiohttp_client(request_hook=capture_request_headers)
+
+async def main():
+    async with aiohttp.ClientSession() as session:
+        await session.get("https://httpbin.org/get")
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(main())
+```
+
+#### Capture Only Response Headers
+
+Similarly, you can create a response hook to capture only the response headers:
+
+```py
+import aiohttp
+import logfire
+from aiohttp.tracing import TraceRequestEndParams, TraceRequestExceptionParams
+from opentelemetry.trace import Span
+from typing import Union
+
+
+def capture_response_headers(span: Span, response: Union[TraceRequestEndParams, TraceRequestExceptionParams]):
+    if hasattr(response, 'response') and response.response:
+        headers = response.response.headers
+        span.set_attributes(
+            {f'http.response.header.{header_name}': headers.getall(header_name)
+             for header_name in headers.keys()}
+        )
+
+
+logfire.configure()
+logfire.instrument_aiohttp_client(response_hook=capture_response_headers)
+
+async def main():
+    async with aiohttp.ClientSession() as session:
+        await session.get('https://httpbin.org/get')
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(main())
+```
+
+You can also use the hooks to filter headers or modify them before capturing them.
+
+### Capture HTTP Response Bodies
+
+By default, **Logfire** doesn't capture HTTP response bodies.
+
+To capture response bodies, you can set the `capture_response_body` parameter to `True`.
+
+```py
+import aiohttp
+import logfire
+
+logfire.configure()
+logfire.instrument_aiohttp_client(capture_response_body=True)
+
+async def main():
+    async with aiohttp.ClientSession() as session:
+        response = await session.get("https://httpbin.org/get")
+        await response.text()
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(main())
+```
+
 ## Hiding sensitive URL parameters
 
 The `url_filter` keyword argument can be used to modify the URL that's recorded in spans. Here's an example of how to use this to redact query parameters:

logfire/_internal/integrations/aiohttp_client.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
+import functools
 from typing import TYPE_CHECKING, Any, Callable, Literal
 
 import attr
 from aiohttp.client_reqrep import ClientResponse
 from aiohttp.tracing import TraceRequestEndParams, TraceRequestExceptionParams, TraceRequestStartParams
-from opentelemetry.trace import Span
+from opentelemetry.trace import NonRecordingSpan, Span, use_span
 from yarl import URL
 
 try:
@@ -17,7 +18,7 @@
         "    pip install 'logfire[aiohttp-client]'"
     )
 
-from logfire import Logfire
+from logfire import Logfire, LogfireSpan
 from logfire._internal.utils import handle_internal_errors
 from logfire.integrations.aiohttp_client import AioHttpRequestHeaders, AioHttpResponseHeaders, RequestHook, ResponseHook
 
@@ -29,6 +30,7 @@
 
 def instrument_aiohttp_client(
     logfire_instance: Logfire,
+    capture_response_body: bool,
     capture_headers: bool,
     request_hook: RequestHook | None,
     response_hook: ResponseHook | None,
@@ -48,6 +50,7 @@ def instrument_aiohttp_client(
                 response_hook,
                 logfire_instance,
                 capture_headers,
+                capture_response_body,
             ),
             'meter_provider': logfire_instance.config.get_meter_provider(),
             **kwargs,
@@ -67,7 +70,7 @@ def capture_headers(self):
         capture_request_or_response_headers(self.span, self.headers, 'request')
 
 
-@attr.s(auto_attribs=True, frozen=True, slots=True)
+@attr.s(auto_attribs=True, slots=True)
 class LogfireAioHttpResponseInfo(LogfireClientInfoMixin):
     span: Span
     method: str
@@ -76,11 +79,45 @@ class LogfireAioHttpResponseInfo(LogfireClientInfoMixin):
     response: ClientResponse | None
     exception: BaseException | None
     logfire_instance: Logfire
+    body_captured: bool = False
 
     def capture_headers(self):
         if self.response:
             capture_request_or_response_headers(self.span, self.response.headers, 'response')
 
+    def capture_body_if_text(self, attr_name: str = 'http.response.body.text') -> None:
+        response = self.response
+        if response is None:
+            return
+
+        original_read = response.read
+
+        @functools.wraps(original_read)
+        async def read() -> bytes:
+            if self.body_captured:
+                return await original_read()
+
+            with (
+                use_span(NonRecordingSpan(self.span.get_span_context())),
+                self.logfire_instance.span('Reading response body') as span,
+            ):
+                body = await original_read()
+                try:
+                    encoding = response.get_encoding()
+                    text = body.decode(encoding)
+                except (UnicodeDecodeError, LookupError):
+                    self.body_captured = True
+                    return body
+                self.capture_text_as_json(span, text=text, attr_name=attr_name)
+                self.body_captured = True
+                return body
+
+        response.read = read
+
+    def capture_text_as_json(self, span: LogfireSpan, *, text: str, attr_name: str) -> None:
+        span.set_attribute(attr_name, {})
+        span._span.set_attribute(attr_name, text)  # type: ignore
+
     @classmethod
     def create_from_trace_params(
         cls,
@@ -115,8 +152,9 @@ def make_response_hook(
     hook: ResponseHook | None,
     logfire_instance: Logfire,
     capture_headers: bool,
+    capture_response_body: bool,
 ) -> ResponseHook | None:
-    if not (capture_headers or hook):
+    if not (capture_headers or capture_response_body or hook):
         return None
 
     def new_hook(span: Span, response: TraceRequestEndParams | TraceRequestExceptionParams) -> None:
@@ -126,6 +164,7 @@ def new_hook(span: Span, response: TraceRequestEndParams | TraceRequestException
                 response,
                 logfire_instance,
                 capture_headers,
+                capture_response_body,
             )
             run_hook(hook, span, response)
 
@@ -150,6 +189,7 @@ def capture_response(
     response: TraceRequestEndParams | TraceRequestExceptionParams,
     logfire_instance: Logfire,
     capture_headers: bool,
+    capture_response_body: bool,
 ) -> LogfireAioHttpResponseInfo:
     response_info = LogfireAioHttpResponseInfo.create_from_trace_params(
         span=span, params=response, logfire_instance=logfire_instance
@@ -158,6 +198,9 @@ def capture_response(
     if capture_headers:
         response_info.capture_headers()
 
+    if capture_response_body:
+        response_info.capture_body_if_text()
+
     return response_info
 
 

logfire/_internal/main.py

@@ -1768,6 +1768,7 @@ def instrument_aiohttp_client(
         self,
         *,
         capture_headers: bool = False,
+        capture_response_body: bool = False,
         request_hook: AiohttpClientRequestHook | None = None,
         response_hook: AiohttpClientResponseHook | None = None,
         **kwargs: Any,
@@ -1783,6 +1784,7 @@ def instrument_aiohttp_client(
         self._warn_if_not_initialized_for_instrumentation()
         return instrument_aiohttp_client(
             self,
+            capture_response_body=capture_response_body,
             capture_headers=capture_headers,
             request_hook=request_hook,
             response_hook=response_hook,