Skip to content

Commit 5d37a10

Browse files
Aiohttp response capture (#1409)
Co-authored-by: Alex Hall <[email protected]>
1 parent ed061d9 commit 5d37a10

File tree

4 files changed

+360
-4
lines changed

4 files changed

+360
-4
lines changed

docs/integrations/http-clients/aiohttp.md

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,125 @@ if __name__ == "__main__":
4242

4343
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].
4444

45+
## Configuration
46+
47+
The `logfire.instrument_aiohttp_client()` method accepts various parameters to configure the instrumentation.
48+
49+
!!! note
50+
The aiohttp client instrumentation captures request and response headers, and response bodies. Request bodies are not captured.
51+
52+
53+
### Capture HTTP Headers
54+
55+
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`.
56+
57+
```py
58+
import aiohttp
59+
import logfire
60+
61+
logfire.configure()
62+
logfire.instrument_aiohttp_client(capture_headers=True)
63+
64+
async def main():
65+
async with aiohttp.ClientSession() as session:
66+
await session.get("https://httpbin.org/get")
67+
68+
if __name__ == "__main__":
69+
import asyncio
70+
asyncio.run(main())
71+
```
72+
73+
#### Capture Only Request Headers
74+
75+
Instead of capturing both request and response headers, you can create a request hook to capture only the request headers:
76+
77+
```py
78+
import aiohttp
79+
import logfire
80+
from aiohttp.tracing import TraceRequestStartParams
81+
from opentelemetry.trace import Span
82+
83+
84+
def capture_request_headers(span: Span, request: TraceRequestStartParams):
85+
headers = request.headers
86+
span.set_attributes(
87+
{
88+
f'http.request.header.{header_name}': headers.getall(header_name)
89+
for header_name in headers.keys()
90+
}
91+
)
92+
93+
94+
logfire.configure()
95+
logfire.instrument_aiohttp_client(request_hook=capture_request_headers)
96+
97+
async def main():
98+
async with aiohttp.ClientSession() as session:
99+
await session.get("https://httpbin.org/get")
100+
101+
if __name__ == "__main__":
102+
import asyncio
103+
asyncio.run(main())
104+
```
105+
106+
#### Capture Only Response Headers
107+
108+
Similarly, you can create a response hook to capture only the response headers:
109+
110+
```py
111+
import aiohttp
112+
import logfire
113+
from aiohttp.tracing import TraceRequestEndParams, TraceRequestExceptionParams
114+
from opentelemetry.trace import Span
115+
from typing import Union
116+
117+
118+
def capture_response_headers(span: Span, response: Union[TraceRequestEndParams, TraceRequestExceptionParams]):
119+
if hasattr(response, 'response') and response.response:
120+
headers = response.response.headers
121+
span.set_attributes(
122+
{f'http.response.header.{header_name}': headers.getall(header_name)
123+
for header_name in headers.keys()}
124+
)
125+
126+
127+
logfire.configure()
128+
logfire.instrument_aiohttp_client(response_hook=capture_response_headers)
129+
130+
async def main():
131+
async with aiohttp.ClientSession() as session:
132+
await session.get('https://httpbin.org/get')
133+
134+
if __name__ == "__main__":
135+
import asyncio
136+
asyncio.run(main())
137+
```
138+
139+
You can also use the hooks to filter headers or modify them before capturing them.
140+
141+
### Capture HTTP Response Bodies
142+
143+
By default, **Logfire** doesn't capture HTTP response bodies.
144+
145+
To capture response bodies, you can set the `capture_response_body` parameter to `True`.
146+
147+
```py
148+
import aiohttp
149+
import logfire
150+
151+
logfire.configure()
152+
logfire.instrument_aiohttp_client(capture_response_body=True)
153+
154+
async def main():
155+
async with aiohttp.ClientSession() as session:
156+
response = await session.get("https://httpbin.org/get")
157+
await response.text()
158+
159+
if __name__ == "__main__":
160+
import asyncio
161+
asyncio.run(main())
162+
```
163+
45164
## Hiding sensitive URL parameters
46165

47166
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

Lines changed: 47 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
from __future__ import annotations
22

3+
import functools
34
from typing import TYPE_CHECKING, Any, Callable, Literal
45

56
import attr
67
from aiohttp.client_reqrep import ClientResponse
78
from aiohttp.tracing import TraceRequestEndParams, TraceRequestExceptionParams, TraceRequestStartParams
8-
from opentelemetry.trace import Span
9+
from opentelemetry.trace import NonRecordingSpan, Span, use_span
910
from yarl import URL
1011

1112
try:
@@ -17,7 +18,7 @@
1718
" pip install 'logfire[aiohttp-client]'"
1819
)
1920

20-
from logfire import Logfire
21+
from logfire import Logfire, LogfireSpan
2122
from logfire._internal.utils import handle_internal_errors
2223
from logfire.integrations.aiohttp_client import AioHttpRequestHeaders, AioHttpResponseHeaders, RequestHook, ResponseHook
2324

@@ -29,6 +30,7 @@
2930

3031
def instrument_aiohttp_client(
3132
logfire_instance: Logfire,
33+
capture_response_body: bool,
3234
capture_headers: bool,
3335
request_hook: RequestHook | None,
3436
response_hook: ResponseHook | None,
@@ -48,6 +50,7 @@ def instrument_aiohttp_client(
4850
response_hook,
4951
logfire_instance,
5052
capture_headers,
53+
capture_response_body,
5154
),
5255
'meter_provider': logfire_instance.config.get_meter_provider(),
5356
**kwargs,
@@ -67,7 +70,7 @@ def capture_headers(self):
6770
capture_request_or_response_headers(self.span, self.headers, 'request')
6871

6972

70-
@attr.s(auto_attribs=True, frozen=True, slots=True)
73+
@attr.s(auto_attribs=True, slots=True)
7174
class LogfireAioHttpResponseInfo(LogfireClientInfoMixin):
7275
span: Span
7376
method: str
@@ -76,11 +79,45 @@ class LogfireAioHttpResponseInfo(LogfireClientInfoMixin):
7679
response: ClientResponse | None
7780
exception: BaseException | None
7881
logfire_instance: Logfire
82+
body_captured: bool = False
7983

8084
def capture_headers(self):
8185
if self.response:
8286
capture_request_or_response_headers(self.span, self.response.headers, 'response')
8387

88+
def capture_body_if_text(self, attr_name: str = 'http.response.body.text') -> None:
89+
response = self.response
90+
if response is None:
91+
return
92+
93+
original_read = response.read
94+
95+
@functools.wraps(original_read)
96+
async def read() -> bytes:
97+
if self.body_captured:
98+
return await original_read()
99+
100+
with (
101+
use_span(NonRecordingSpan(self.span.get_span_context())),
102+
self.logfire_instance.span('Reading response body') as span,
103+
):
104+
body = await original_read()
105+
try:
106+
encoding = response.get_encoding()
107+
text = body.decode(encoding)
108+
except (UnicodeDecodeError, LookupError):
109+
self.body_captured = True
110+
return body
111+
self.capture_text_as_json(span, text=text, attr_name=attr_name)
112+
self.body_captured = True
113+
return body
114+
115+
response.read = read
116+
117+
def capture_text_as_json(self, span: LogfireSpan, *, text: str, attr_name: str) -> None:
118+
span.set_attribute(attr_name, {})
119+
span._span.set_attribute(attr_name, text) # type: ignore
120+
84121
@classmethod
85122
def create_from_trace_params(
86123
cls,
@@ -115,8 +152,9 @@ def make_response_hook(
115152
hook: ResponseHook | None,
116153
logfire_instance: Logfire,
117154
capture_headers: bool,
155+
capture_response_body: bool,
118156
) -> ResponseHook | None:
119-
if not (capture_headers or hook):
157+
if not (capture_headers or capture_response_body or hook):
120158
return None
121159

122160
def new_hook(span: Span, response: TraceRequestEndParams | TraceRequestExceptionParams) -> None:
@@ -126,6 +164,7 @@ def new_hook(span: Span, response: TraceRequestEndParams | TraceRequestException
126164
response,
127165
logfire_instance,
128166
capture_headers,
167+
capture_response_body,
129168
)
130169
run_hook(hook, span, response)
131170

@@ -150,6 +189,7 @@ def capture_response(
150189
response: TraceRequestEndParams | TraceRequestExceptionParams,
151190
logfire_instance: Logfire,
152191
capture_headers: bool,
192+
capture_response_body: bool,
153193
) -> LogfireAioHttpResponseInfo:
154194
response_info = LogfireAioHttpResponseInfo.create_from_trace_params(
155195
span=span, params=response, logfire_instance=logfire_instance
@@ -158,6 +198,9 @@ def capture_response(
158198
if capture_headers:
159199
response_info.capture_headers()
160200

201+
if capture_response_body:
202+
response_info.capture_body_if_text()
203+
161204
return response_info
162205

163206

logfire/_internal/main.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1768,6 +1768,7 @@ def instrument_aiohttp_client(
17681768
self,
17691769
*,
17701770
capture_headers: bool = False,
1771+
capture_response_body: bool = False,
17711772
request_hook: AiohttpClientRequestHook | None = None,
17721773
response_hook: AiohttpClientResponseHook | None = None,
17731774
**kwargs: Any,
@@ -1783,6 +1784,7 @@ def instrument_aiohttp_client(
17831784
self._warn_if_not_initialized_for_instrumentation()
17841785
return instrument_aiohttp_client(
17851786
self,
1787+
capture_response_body=capture_response_body,
17861788
capture_headers=capture_headers,
17871789
request_hook=request_hook,
17881790
response_hook=response_hook,

0 commit comments

Comments
 (0)