Add capture_request_form_data param to instrument_httpx (#711)

Author: alexmojaki

logfire/_internal/integrations/httpx.py

@@ -3,9 +3,10 @@
 import inspect
 from contextlib import suppress
 from email.message import Message
-from typing import TYPE_CHECKING, Any, Callable, Literal, cast, overload
+from typing import TYPE_CHECKING, Any, Callable, Literal, Mapping, cast, overload
 
 import httpx
+import opentelemetry.sdk.trace
 
 from logfire.propagate import attach_context, get_context
 
@@ -67,6 +68,7 @@ def instrument_httpx(
         capture_response_headers: bool,
         capture_request_json_body: bool,
         capture_response_json_body: bool,
+        capture_request_form_data: bool,
         **kwargs: Unpack[ClientKwargs],
     ) -> None: ...
 
@@ -78,6 +80,7 @@ def instrument_httpx(
         capture_response_headers: bool,
         capture_request_json_body: bool,
         capture_response_json_body: bool,
+        capture_request_form_data: bool,
         **kwargs: Unpack[AsyncClientKwargs],
     ) -> None: ...
 
@@ -89,6 +92,7 @@ def instrument_httpx(
         capture_response_headers: bool,
         capture_request_json_body: bool,
         capture_response_json_body: bool,
+        capture_request_form_data: bool,
         **kwargs: Unpack[HTTPXInstrumentKwargs],
     ) -> None: ...
 
@@ -100,6 +104,7 @@ def instrument_httpx(
     capture_response_headers: bool,
     capture_request_json_body: bool,
     capture_response_json_body: bool,
+    capture_request_form_data: bool,
     **kwargs: Any,
 ) -> None:
     """Instrument the `httpx` module so that spans are automatically created for each request.
@@ -122,13 +127,13 @@ def instrument_httpx(
         async_request_hook = cast('AsyncRequestHook | None', final_kwargs.get('async_request_hook'))
         async_response_hook = cast('AsyncResponseHook | None', final_kwargs.get('async_response_hook'))
         final_kwargs['request_hook'] = make_request_hook(
-            request_hook, capture_request_headers, capture_request_json_body
+            request_hook, capture_request_headers, capture_request_json_body, capture_request_form_data
         )
         final_kwargs['response_hook'] = make_response_hook(
             response_hook, capture_response_headers, capture_response_json_body, logfire_instance
         )
         final_kwargs['async_request_hook'] = make_async_request_hook(
-            async_request_hook, capture_request_headers, capture_request_json_body
+            async_request_hook, capture_request_headers, capture_request_json_body, capture_request_form_data
         )
         final_kwargs['async_response_hook'] = make_async_response_hook(
             async_response_hook, capture_response_headers, capture_response_json_body, logfire_instance
@@ -140,15 +145,19 @@ def instrument_httpx(
             request_hook = cast('RequestHook | AsyncRequestHook | None', final_kwargs.get('request_hook'))
             response_hook = cast('ResponseHook | AsyncResponseHook | None', final_kwargs.get('response_hook'))
 
-            request_hook = make_async_request_hook(request_hook, capture_request_headers, capture_request_json_body)
+            request_hook = make_async_request_hook(
+                request_hook, capture_request_headers, capture_request_json_body, capture_request_form_data
+            )
             response_hook = make_async_response_hook(
                 response_hook, capture_response_headers, capture_response_json_body, logfire_instance
             )
         else:
             request_hook = cast('RequestHook | None', final_kwargs.get('request_hook'))
             response_hook = cast('ResponseHook | None', final_kwargs.get('response_hook'))
 
-            request_hook = make_request_hook(request_hook, capture_request_headers, capture_request_json_body)
+            request_hook = make_request_hook(
+                request_hook, capture_request_headers, capture_request_json_body, capture_request_form_data
+            )
             response_hook = make_response_hook(
                 response_hook, capture_response_headers, capture_response_json_body, logfire_instance
             )
@@ -158,39 +167,51 @@ def instrument_httpx(
 
 
 def make_request_hook(
-    hook: RequestHook | None, should_capture_headers: bool, should_capture_json: bool
+    hook: RequestHook | None, should_capture_headers: bool, should_capture_json: bool, should_capture_form_data: bool
 ) -> RequestHook | None:
-    if not should_capture_headers and not should_capture_json and not hook:
+    if not should_capture_headers and not should_capture_json and not should_capture_form_data and not hook:
         return None
 
     def new_hook(span: Span, request: RequestInfo) -> None:
         with handle_internal_errors():
-            if should_capture_headers:
-                capture_request_headers(span, request)
-            if should_capture_json:
-                capture_request_body(span, request)
+            capture_request(request, span, should_capture_headers, should_capture_json, should_capture_form_data)
             run_hook(hook, span, request)
 
     return new_hook
 
 
 def make_async_request_hook(
-    hook: AsyncRequestHook | RequestHook | None, should_capture_headers: bool, should_capture_json: bool
+    hook: AsyncRequestHook | RequestHook | None,
+    should_capture_headers: bool,
+    should_capture_json: bool,
+    should_capture_form_data: bool,
 ) -> AsyncRequestHook | None:
-    if not should_capture_headers and not should_capture_json and not hook:
+    if not should_capture_headers and not should_capture_json and not should_capture_form_data and not hook:
         return None
 
     async def new_hook(span: Span, request: RequestInfo) -> None:
         with handle_internal_errors():
-            if should_capture_headers:
-                capture_request_headers(span, request)
-            if should_capture_json:
-                capture_request_body(span, request)
+            capture_request(request, span, should_capture_headers, should_capture_json, should_capture_form_data)
             await run_async_hook(hook, span, request)
 
     return new_hook
 
 
+def capture_request(
+    request: RequestInfo,
+    span: Span,
+    should_capture_headers: bool,
+    should_capture_json: bool,
+    should_capture_form_data: bool,
+) -> None:
+    if should_capture_headers:
+        capture_request_headers(span, request)
+    if should_capture_json:
+        capture_request_body(span, request)
+    if should_capture_form_data:
+        capture_request_form_data(span, request)
+
+
 def make_response_hook(
     hook: ResponseHook | None, should_capture_headers: bool, should_capture_json: bool, logfire_instance: Logfire
 ) -> ResponseHook | None:
@@ -338,3 +359,34 @@ def capture_request_body(span: Span, request: RequestInfo) -> None:
     attr_name = 'http.request.body.json'
     set_user_attributes_on_raw_span(span, {attr_name: {}})  # type: ignore
     span.set_attribute(attr_name, body)
+
+
+CODES_FOR_METHODS_WITH_DATA_PARAM = [
+    inspect.unwrap(method).__code__
+    for method in [
+        httpx.Client.request,
+        httpx.Client.stream,
+        httpx.AsyncClient.request,
+        httpx.AsyncClient.stream,
+    ]
+]
+
+
+def capture_request_form_data(span: Span, request: RequestInfo) -> None:
+    content_type = cast('httpx.Headers', request.headers).get('content-type', '')
+    if content_type != 'application/x-www-form-urlencoded':
+        return
+
+    frame = inspect.currentframe().f_back.f_back.f_back  # type: ignore
+    while frame:
+        if frame.f_code in CODES_FOR_METHODS_WITH_DATA_PARAM:
+            break
+        frame = frame.f_back
+    else:  # pragma: no cover
+        return
+
+    data = frame.f_locals.get('data')
+    if not (data and isinstance(data, Mapping)):  # pragma: no cover
+        return
+    span = cast(opentelemetry.sdk.trace.Span, span)
+    set_user_attributes_on_raw_span(span, {'http.request.body.form': data})

logfire/_internal/main.py

@@ -1153,42 +1153,50 @@ def instrument_asyncpg(self, **kwargs: Unpack[AsyncPGInstrumentKwargs]) -> None:
     def instrument_httpx(
         self,
         client: httpx.Client,
+        *,
         capture_request_headers: bool = False,
         capture_response_headers: bool = False,
         capture_request_json_body: bool = False,
         capture_response_json_body: bool = False,
+        capture_request_form_data: bool = False,
         **kwargs: Unpack[ClientKwargs],
     ) -> None: ...
 
     @overload
     def instrument_httpx(
         self,
         client: httpx.AsyncClient,
+        *,
         capture_request_headers: bool = False,
         capture_response_headers: bool = False,
         capture_request_json_body: bool = False,
         capture_response_json_body: bool = False,
+        capture_request_form_data: bool = False,
         **kwargs: Unpack[AsyncClientKwargs],
     ) -> None: ...
 
     @overload
     def instrument_httpx(
         self,
         client: None = None,
+        *,
         capture_request_headers: bool = False,
         capture_response_headers: bool = False,
         capture_request_json_body: bool = False,
         capture_response_json_body: bool = False,
+        capture_request_form_data: bool = False,
         **kwargs: Unpack[HTTPXInstrumentKwargs],
     ) -> None: ...
 
     def instrument_httpx(
         self,
         client: httpx.Client | httpx.AsyncClient | None = None,
+        *,
         capture_request_headers: bool = False,
         capture_response_headers: bool = False,
         capture_request_json_body: bool = False,
         capture_response_json_body: bool = False,
+        capture_request_form_data: bool = False,
         **kwargs: Any,
     ) -> None:
         """Instrument the `httpx` module so that spans are automatically created for each request.
@@ -1205,7 +1213,17 @@ def instrument_httpx(
             capture_request_headers: Set to `True` to capture all request headers.
             capture_response_headers: Set to `True` to capture all response headers.
             capture_request_json_body: Set to `True` to capture the request JSON body.
+                Specifically captures the raw request body whenever the content type is `application/json`.
+                Doesn't check if the body is actually JSON.
             capture_response_json_body: Set to `True` to capture the response JSON body.
+                Specifically captures the raw response body whenever the content type is `application/json`
+                when the `response.read()` or `.aread()` method is first called,
+                which happens automatically for non-streaming requests.
+                For streaming requests, the body is not captured if it's merely iterated over.
+                Doesn't check if the body is actually JSON.
+            capture_request_form_data: Set to `True` to capture the request form data.
+                Specifically captures the `data` argument of `httpx` methods like `post` and `put`.
+                Doesn't inspect or parse the raw request body.
             **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` method, for future compatibility.
         """
         from .integrations.httpx import instrument_httpx
@@ -1214,10 +1232,11 @@ def instrument_httpx(
         return instrument_httpx(
             self,
             client,
-            capture_request_headers,
-            capture_response_headers,
+            capture_request_headers=capture_request_headers,
+            capture_response_headers=capture_response_headers,
             capture_request_json_body=capture_request_json_body,
             capture_response_json_body=capture_response_json_body,
+            capture_request_form_data=capture_request_form_data,
             **kwargs,
         )
 

tests/otel_integrations/test_httpx.py

@@ -15,6 +15,7 @@
 
 import logfire
 import logfire._internal.integrations.httpx
+from logfire._internal.integrations.httpx import CODES_FOR_METHODS_WITH_DATA_PARAM
 from logfire.testing import TestExporter
 
 pytestmark = pytest.mark.anyio
@@ -452,6 +453,7 @@ async def test_async_httpx_client_capture_full(exporter: TestExporter):
                 capture_request_json_body=True,
                 capture_response_headers=True,
                 capture_response_json_body=True,
+                capture_request_form_data=True,
             )
             response = await client.post('https://example.org/', json={'hello': 'world'})
             checker(response)
@@ -558,3 +560,42 @@ async def test_httpx_async_client_capture_json_response_checks_header(exporter:
     assert len(spans) == 1
     assert spans[0]['name'] == 'POST'
     assert 'http.response.body.json' not in str(spans)
+
+
+def test_httpx_client_capture_request_form_data(exporter: TestExporter):
+    assert len({code.co_filename for code in CODES_FOR_METHODS_WITH_DATA_PARAM}) == 1
+    assert [code.co_name for code in CODES_FOR_METHODS_WITH_DATA_PARAM] == ['request', 'stream', 'request', 'stream']
+
+    with httpx.Client(transport=create_transport()) as client:
+        logfire.instrument_httpx(client, capture_request_form_data=True)
+        client.post('https://example.org/', data={'form': 'values'})
+
+    assert exporter.exported_spans_as_dict() == snapshot(
+        [
+            {
+                'name': 'POST',
+                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
+                'parent': None,
+                'start_time': 1000000000,
+                'end_time': 2000000000,
+                'attributes': {
+                    'http.method': 'POST',
+                    'http.request.method': 'POST',
+                    'http.url': 'https://example.org/',
+                    'url.full': 'https://example.org/',
+                    'http.host': 'example.org',
+                    'server.address': 'example.org',
+                    'network.peer.address': 'example.org',
+                    'logfire.span_type': 'span',
+                    'logfire.msg': 'POST /',
+                    'http.request.body.form': '{"form":"values"}',
+                    'logfire.json_schema': '{"type":"object","properties":{"http.request.body.form":{"type":"object"}}}',
+                    'http.status_code': 200,
+                    'http.response.status_code': 200,
+                    'http.flavor': '1.1',
+                    'network.protocol.version': '1.1',
+                    'http.target': '/',
+                },
+            }
+        ]
+    )