Skip to content

Commit

Permalink
capture_request_text_body
Browse files Browse the repository at this point in the history
  • Loading branch information
@alexmojaki
alexmojaki committed Dec 23, 2024
1 parent bd76022 commit acdd595
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 9 deletions.
55 changes: 46 additions & 9 deletions logfire/_internal/integrations/httpx.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ def instrument_httpx(
client: httpx.Client | httpx.AsyncClient | None,
capture_headers: bool,
capture_request_json_body: bool,
capture_request_text_body: bool,
capture_response_json_body: bool,
capture_request_form_data: bool,
**kwargs: Any,
Expand Down Expand Up @@ -108,13 +109,21 @@ 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, should_capture_request_headers, capture_request_json_body, capture_request_form_data
request_hook,
should_capture_request_headers,
capture_request_json_body,
capture_request_text_body,
capture_request_form_data,
)
final_kwargs['response_hook'] = make_response_hook(
response_hook, should_capture_response_headers, capture_response_json_body, logfire_instance
)
final_kwargs['async_request_hook'] = make_async_request_hook(
async_request_hook, should_capture_request_headers, capture_request_json_body, capture_request_form_data
async_request_hook,
should_capture_request_headers,
capture_request_json_body,
capture_request_text_body,
capture_request_form_data,
)
final_kwargs['async_response_hook'] = make_async_response_hook(
async_response_hook, should_capture_response_headers, capture_response_json_body, logfire_instance
Expand All @@ -127,7 +136,11 @@ def instrument_httpx(
response_hook = cast('ResponseHook | AsyncResponseHook | None', final_kwargs.get('response_hook'))

request_hook = make_async_request_hook(
request_hook, should_capture_request_headers, capture_request_json_body, capture_request_form_data
request_hook,
should_capture_request_headers,
capture_request_json_body,
capture_request_text_body,
capture_request_form_data,
)
response_hook = make_async_response_hook(
response_hook, should_capture_response_headers, capture_response_json_body, logfire_instance
Expand All @@ -137,7 +150,11 @@ def instrument_httpx(
response_hook = cast('ResponseHook | None', final_kwargs.get('response_hook'))

request_hook = make_request_hook(
request_hook, should_capture_request_headers, capture_request_json_body, capture_request_form_data
request_hook,
should_capture_request_headers,
capture_request_json_body,
capture_request_text_body,
capture_request_form_data,
)
response_hook = make_response_hook(
response_hook, should_capture_response_headers, capture_response_json_body, logfire_instance
Expand All @@ -157,6 +174,10 @@ def capture_body_if_json(self, attr_name: str = 'http.request.body.json'):
if not self.body_is_streaming and self.content_type_is_json:
self.capture_text_as_json(attr_name=attr_name)

def capture_body_if_text(self, attr_name: str = 'http.request.body.text'):
if not self.body_is_streaming and self.content_type_is_text:
self.span.set_attribute(attr_name, self.text)

Check warning on line 179 in logfire/_internal/integrations/httpx.py

View check run for this annotation

Codecov / codecov/patch

logfire/_internal/integrations/httpx.py#L179 

Added line #L179 was not covered by tests

def capture_body_if_form(self, attr_name: str = 'http.request.body.form'):
if not self.content_type_is_form:
return
Expand Down Expand Up @@ -188,6 +209,10 @@ def content_type_header_string(self) -> str:
def content_type_is_json(self):
return is_json_type(self.content_type_header_string)

@property
def content_type_is_text(self):
return is_text_type(self.content_type_header_string)

Check warning on line 214 in logfire/_internal/integrations/httpx.py

View check run for this annotation

Codecov / codecov/patch

logfire/_internal/integrations/httpx.py#L214 

Added line #L214 was not covered by tests

@property
def content_type_is_form(self):
content_type = self.content_type_header_string
Expand Down Expand Up @@ -224,16 +249,22 @@ def set_complex_span_attributes(self, attributes: dict[str, Any]):


def make_request_hook(
hook: RequestHook | None, should_capture_headers: bool, should_capture_json: bool, should_capture_form_data: bool
hook: RequestHook | None,
should_capture_headers: bool,
should_capture_json: bool,
should_capture_text: bool,
should_capture_form_data: bool,
) -> RequestHook | None:
if not should_capture_headers and not should_capture_json and not should_capture_form_data and not hook:
if not (should_capture_headers or should_capture_json or should_capture_text or should_capture_form_data or hook):
return None

def new_hook(span: Span, request: RequestInfo) -> None:
with handle_internal_errors():
request = LogfireHttpxRequestInfo(*request)
request.span = span
capture_request(request, should_capture_headers, should_capture_json, should_capture_form_data)
capture_request(
request, should_capture_headers, should_capture_json, should_capture_text, should_capture_form_data
)
run_hook(hook, span, request)

return new_hook
Expand All @@ -243,16 +274,19 @@ def make_async_request_hook(
hook: AsyncRequestHook | RequestHook | None,
should_capture_headers: bool,
should_capture_json: bool,
should_capture_text: bool,
should_capture_form_data: bool,
) -> AsyncRequestHook | None:
if not should_capture_headers and not should_capture_json and not should_capture_form_data and not hook:
if not (should_capture_headers or should_capture_json or should_capture_text or should_capture_form_data or hook):
return None

async def new_hook(span: Span, request: RequestInfo) -> None:
with handle_internal_errors():
request = LogfireHttpxRequestInfo(*request)
request.span = span
capture_request(request, should_capture_headers, should_capture_json, should_capture_form_data)
capture_request(
request, should_capture_headers, should_capture_json, should_capture_text, should_capture_form_data
)
await run_async_hook(hook, span, request)

return new_hook
Expand All @@ -262,12 +296,15 @@ def capture_request(
request: LogfireHttpxRequestInfo,
should_capture_headers: bool,
should_capture_json: bool,
should_capture_text: bool,
should_capture_form_data: bool,
) -> None:
if should_capture_headers:
request.capture_headers()
if should_capture_json:
request.capture_body_if_json()
if should_capture_text and not (should_capture_json and request.content_type_is_json):
request.capture_body_if_text()

Check warning on line 307 in logfire/_internal/integrations/httpx.py

View check run for this annotation

Codecov / codecov/patch

logfire/_internal/integrations/httpx.py#L307 

Added line #L307 was not covered by tests
if should_capture_form_data:
request.capture_body_if_form()

Expand Down
8 changes: 8 additions & 0 deletions logfire/_internal/main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1157,6 +1157,7 @@ def instrument_httpx(
client: httpx.Client,
*,
capture_headers: bool = False,
capture_request_text_body: bool = False,
capture_request_json_body: bool = False,
capture_response_json_body: bool = False,
capture_request_form_data: bool = False,
Expand All @@ -1170,6 +1171,7 @@ def instrument_httpx(
*,
capture_headers: bool = False,
capture_request_json_body: bool = False,
capture_request_text_body: bool = False,
capture_response_json_body: bool = False,
capture_request_form_data: bool = False,
**kwargs: Unpack[AsyncClientKwargs],
Expand All @@ -1182,6 +1184,7 @@ def instrument_httpx(
*,
capture_headers: bool = False,
capture_request_json_body: bool = False,
capture_request_text_body: bool = False,
capture_response_json_body: bool = False,
capture_request_form_data: bool = False,
**kwargs: Unpack[HTTPXInstrumentKwargs],
Expand All @@ -1193,6 +1196,7 @@ def instrument_httpx(
*,
capture_headers: bool = False,
capture_request_json_body: bool = False,
capture_request_text_body: bool = False,
capture_response_json_body: bool = False,
capture_request_form_data: bool = False,
**kwargs: Any,
Expand All @@ -1212,6 +1216,9 @@ def instrument_httpx(
If you don't want to capture all headers, you can customize the headers captured. See the
[Capture Headers](https://logfire.pydantic.dev/docs/guides/advanced/capture_headers/) section for more info.
capture_request_text_body: Set to `True` to capture the request text body
if the content type is either `text/*`
or `application/` followed by a known human-readable text format, e.g. XML.
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.
Expand All @@ -1234,6 +1241,7 @@ def instrument_httpx(
client,
capture_headers=capture_headers,
capture_request_json_body=capture_request_json_body,
capture_request_text_body=capture_request_text_body,
capture_response_json_body=capture_response_json_body,
capture_request_form_data=capture_request_form_data,
**kwargs,
Expand Down

0 comments on commit acdd595

Please sign in to comment.