From 425b4412df931ca989b195b59f0eb0ae50c6b84a Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Tue, 7 Jan 2025 12:35:54 +0100 Subject: [PATCH] Add `capture_all` to `instrument_httpx` (#780) --- logfire/_internal/integrations/httpx.py | 14 ++-- logfire/_internal/main.py | 6 ++ tests/otel_integrations/test_httpx.py | 95 +++++++++++++++++++++++++ 3 files changed, 111 insertions(+), 4 deletions(-) diff --git a/logfire/_internal/integrations/httpx.py b/logfire/_internal/integrations/httpx.py index b0d33c28..5b81bddd 100644 --- a/logfire/_internal/integrations/httpx.py +++ b/logfire/_internal/integrations/httpx.py @@ -44,6 +44,7 @@ def instrument_httpx( logfire_instance: Logfire, client: httpx.Client | httpx.AsyncClient | None, + capture_all: bool, capture_headers: bool, capture_request_body: bool, capture_response_body: bool, @@ -57,6 +58,11 @@ def instrument_httpx( See the `Logfire.instrument_httpx` method for details. """ + if capture_all and (capture_headers or capture_request_body or capture_response_body): + warn_at_user_stacklevel( + 'You should use either `capture_all` or the specific capture parameters, not both.', UserWarning + ) + capture_request_headers = kwargs.get('capture_request_headers') capture_response_headers = kwargs.get('capture_response_headers') @@ -69,10 +75,10 @@ def instrument_httpx( 'The `capture_response_headers` parameter is deprecated. Use `capture_headers` instead.', DeprecationWarning ) - should_capture_request_headers = capture_request_headers or capture_headers - should_capture_response_headers = capture_response_headers or capture_headers - should_capture_request_body = capture_request_body - should_capture_response_body = capture_response_body + should_capture_request_headers = capture_request_headers or capture_headers or capture_all + should_capture_response_headers = capture_response_headers or capture_headers or capture_all + should_capture_request_body = capture_request_body or capture_all + should_capture_response_body = capture_response_body or capture_all final_kwargs: dict[str, Any] = { 'tracer_provider': logfire_instance.config.get_tracer_provider(), diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index 21b15c6a..7069a24c 100644 --- a/logfire/_internal/main.py +++ b/logfire/_internal/main.py @@ -1183,6 +1183,7 @@ def instrument_httpx( self, client: httpx.Client, *, + capture_all: bool = False, capture_headers: bool = False, capture_request_body: bool = False, capture_response_body: bool = False, @@ -1196,6 +1197,7 @@ def instrument_httpx( self, client: httpx.AsyncClient, *, + capture_all: bool = False, capture_headers: bool = False, capture_request_body: bool = False, capture_response_body: bool = False, @@ -1209,6 +1211,7 @@ def instrument_httpx( self, client: None = None, *, + capture_all: bool = False, capture_headers: bool = False, capture_request_body: bool = False, capture_response_body: bool = False, @@ -1223,6 +1226,7 @@ def instrument_httpx( self, client: httpx.Client | httpx.AsyncClient | None = None, *, + capture_all: bool = False, capture_headers: bool = False, capture_request_body: bool = False, capture_response_body: bool = False, @@ -1243,6 +1247,7 @@ def instrument_httpx( Args: client: The `httpx.Client` or `httpx.AsyncClient` instance to instrument. If `None`, the default, all clients will be instrumented. + capture_all: Set to `True` to capture all HTTP headers, request and response bodies. capture_headers: Set to `True` to capture all HTTP headers. If you don't want to capture all headers, you can customize the headers captured. See the @@ -1261,6 +1266,7 @@ def instrument_httpx( return instrument_httpx( self, client, + capture_all=capture_all, capture_headers=capture_headers, capture_request_body=capture_request_body, capture_response_body=capture_response_body, diff --git a/tests/otel_integrations/test_httpx.py b/tests/otel_integrations/test_httpx.py index 6fbadd1a..92091264 100644 --- a/tests/otel_integrations/test_httpx.py +++ b/tests/otel_integrations/test_httpx.py @@ -673,6 +673,101 @@ def test_is_json_type(): assert not is_json_type('application//json') +async def test_httpx_client_capture_all(exporter: TestExporter): + with check_traceparent_header() as checker: + async with httpx.AsyncClient(transport=create_transport()) as client: + logfire.instrument_httpx(client, capture_all=True) + response = await client.post('https://example.org/', json={'hello': 'world'}) + checker(response) + assert response.json() == {'good': 'response'} + assert await response.aread() == b'{"good": "response"}' + + assert exporter.exported_spans_as_dict() == snapshot( + [ + { + 'name': 'POST', + 'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False}, + 'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False}, + 'start_time': 2000000000, + 'end_time': 3000000000, + '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.header.host': ('example.org',), + 'http.request.header.accept': ('*/*',), + 'http.request.header.accept-encoding': ('gzip, deflate',), + 'http.request.header.connection': ('keep-alive',), + 'http.request.header.user-agent': ('python-httpx/0.28.1',), + 'http.request.header.content-length': ('17',), + 'http.request.header.content-type': ('application/json',), + 'logfire.json_schema': '{"type":"object","properties":{"http.request.body.text":{"type":"object"}}}', + 'http.request.body.text': '{"hello":"world"}', + 'http.status_code': 200, + 'http.response.status_code': 200, + 'http.flavor': '1.1', + 'network.protocol.version': '1.1', + 'http.response.header.host': ('example.org',), + 'http.response.header.accept': ('*/*',), + 'http.response.header.accept-encoding': ('gzip, deflate',), + 'http.response.header.connection': ('keep-alive',), + 'http.response.header.user-agent': ('python-httpx/0.28.1',), + 'http.response.header.content-length': ('17',), + 'http.response.header.content-type': ('application/json',), + 'http.response.header.traceparent': ('00-00000000000000000000000000000001-0000000000000003-01',), + 'http.target': '/', + }, + }, + { + 'name': 'Reading response body', + 'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False}, + 'parent': {'trace_id': 1, 'span_id': 3, 'is_remote': False}, + 'start_time': 4000000000, + 'end_time': 5000000000, + 'attributes': { + 'code.filepath': 'test_httpx.py', + 'code.function': 'test_httpx_client_capture_all', + 'code.lineno': 123, + 'logfire.msg_template': 'Reading response body', + 'logfire.msg': 'Reading response body', + 'logfire.span_type': 'span', + 'http.response.body.text': '{"good": "response"}', + 'logfire.json_schema': '{"type":"object","properties":{"http.response.body.text":{}}}', + }, + }, + { + 'name': 'test span', + 'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False}, + 'parent': None, + 'start_time': 1000000000, + 'end_time': 6000000000, + 'attributes': { + 'code.filepath': 'test_httpx.py', + 'code.function': 'check_traceparent_header', + 'code.lineno': 123, + 'logfire.msg_template': 'test span', + 'logfire.msg': 'test span', + 'logfire.span_type': 'span', + }, + }, + ] + ) + + +def test_httpx_capture_all_and_other_flags_should_warn(exporter: TestExporter): + with httpx.Client(transport=create_transport()) as client: + with pytest.warns( + UserWarning, match='You should use either `capture_all` or the specific capture parameters, not both.' + ): + logfire.instrument_httpx(client, capture_all=True, capture_request_body=True) + + def test_missing_opentelemetry_dependency() -> None: with mock.patch.dict('sys.modules', {'opentelemetry.instrumentation.httpx': None}): with pytest.raises(RuntimeError) as exc_info: