Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add capture_all to instrument_httpx #753

Closed
wants to merge 5 commits into from
Closed

Add capture_all to instrument_httpx #753

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1fd5d55
Drop httpx kwargs
Kludex Dec 30, 2024
0d89586
Merge branch 'main' into drop-httpx-kwargs
Kludex Dec 30, 2024
70d4554
Merge branch 'main' into drop-httpx-kwargs
Kludex Dec 30, 2024
5bb617f
Add `capture_all` to `instrument_httpx`
Kludex Dec 30, 2024
4db0b56
single line import
Kludex Dec 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
103 changes: 49 additions & 54 deletions logfire/_internal/integrations/httpx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,48 +37,43 @@
from logfire._internal.utils import handle_internal_errors

if TYPE_CHECKING:
from typing import ParamSpec, TypedDict, TypeVar

class AsyncClientKwargs(TypedDict, total=False):
request_hook: RequestHook | AsyncRequestHook
response_hook: ResponseHook | AsyncResponseHook
skip_dep_check: bool

class ClientKwargs(TypedDict, total=False):
request_hook: RequestHook
response_hook: ResponseHook
skip_dep_check: bool

class HTTPXInstrumentKwargs(TypedDict, total=False):
request_hook: RequestHook
response_hook: ResponseHook
async_request_hook: AsyncRequestHook
async_response_hook: AsyncResponseHook
skip_dep_check: bool

AnyRequestHook = TypeVar('AnyRequestHook', RequestHook, AsyncRequestHook)
AnyResponseHook = TypeVar('AnyResponseHook', ResponseHook, AsyncResponseHook)
Hook = TypeVar('Hook', RequestHook, ResponseHook)
AsyncHook = TypeVar('AsyncHook', AsyncRequestHook, AsyncResponseHook)
from typing import ParamSpec

P = ParamSpec('P')


def instrument_httpx(
logfire_instance: Logfire,
client: httpx.Client | httpx.AsyncClient | None,
capture_all: bool,
capture_headers: bool,
capture_request_json_body: bool,
capture_request_text_body: bool,
capture_response_json_body: bool,
capture_response_text_body: bool,
capture_request_form_data: bool,
request_hook: RequestHook | AsyncRequestHook | None = None,
response_hook: ResponseHook | AsyncResponseHook | None = None,
async_request_hook: AsyncRequestHook | None = None,
async_response_hook: AsyncResponseHook | None = None,
**kwargs: Any,
) -> None:
"""Instrument the `httpx` module so that spans are automatically created for each request.

See the `Logfire.instrument_httpx` method for details.
"""
if capture_all and (
capture_headers
or capture_request_json_body
or capture_request_text_body
or capture_response_json_body
or capture_response_text_body
or capture_request_form_data
):
warn_at_user_stacklevel(
'You should use either `capture_all` or the specific capture parameters, not both.', UserWarning
)
Comment on lines +73 to +75
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is better being a runtime error instead of UserWarning.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think a warning is fine.


capture_request_headers = kwargs.get('capture_request_headers')
capture_response_headers = kwargs.get('capture_response_headers')

Expand All @@ -91,8 +86,13 @@ 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_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_json_body = capture_request_json_body or capture_all
should_capture_request_text_body = capture_request_text_body or capture_all
should_capture_request_form_data = capture_request_form_data or capture_all
should_capture_response_json_body = capture_response_json_body or capture_all
should_capture_response_text_body = capture_response_text_body or capture_all

final_kwargs: dict[str, Any] = {
'tracer_provider': logfire_instance.config.get_tracer_provider(),
Expand All @@ -105,75 +105,70 @@ def instrument_httpx(
logfire_instance = logfire_instance.with_settings(custom_scope_suffix='httpx')

if client is None:
request_hook = cast('RequestHook | None', final_kwargs.get('request_hook'))
response_hook = cast('ResponseHook | None', final_kwargs.get('response_hook'))
async_request_hook = cast('AsyncRequestHook | None', final_kwargs.get('async_request_hook'))
async_response_hook = cast('AsyncResponseHook | None', final_kwargs.get('async_response_hook'))
request_hook = cast('RequestHook | None', request_hook)
response_hook = cast('ResponseHook | None', response_hook)
final_kwargs['request_hook'] = make_request_hook(
request_hook,
should_capture_request_headers,
capture_request_json_body,
capture_request_text_body,
capture_request_form_data,
should_capture_request_json_body,
should_capture_request_text_body,
should_capture_request_form_data,
)
final_kwargs['response_hook'] = make_response_hook(
response_hook,
should_capture_response_headers,
capture_response_json_body,
capture_response_text_body,
should_capture_response_json_body,
should_capture_response_text_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_text_body,
capture_request_form_data,
should_capture_request_json_body,
should_capture_request_text_body,
should_capture_request_form_data,
)
final_kwargs['async_response_hook'] = make_async_response_hook(
async_response_hook,
should_capture_response_headers,
capture_response_json_body,
capture_response_text_body,
should_capture_response_json_body,
should_capture_response_text_body,
logfire_instance,
)

instrumentor.instrument(**final_kwargs)
else:
if isinstance(client, httpx.AsyncClient):
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,
should_capture_request_headers,
capture_request_json_body,
capture_request_text_body,
capture_request_form_data,
should_capture_request_json_body,
should_capture_request_text_body,
should_capture_request_form_data,
)
response_hook = make_async_response_hook(
response_hook,
should_capture_response_headers,
capture_response_json_body,
capture_response_text_body,
should_capture_response_json_body,
should_capture_response_text_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 = cast('RequestHook | None', request_hook)
response_hook = cast('ResponseHook | None', response_hook)

request_hook = make_request_hook(
request_hook,
should_capture_request_headers,
capture_request_json_body,
capture_request_text_body,
capture_request_form_data,
should_capture_request_json_body,
should_capture_request_text_body,
should_capture_request_form_data,
)
response_hook = make_response_hook(
response_hook,
should_capture_response_headers,
capture_response_json_body,
capture_response_text_body,
should_capture_response_json_body,
should_capture_response_text_body,
logfire_instance,
)

Expand Down
38 changes: 34 additions & 4 deletions logfire/_internal/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,13 +88,18 @@
RequestHook as FlaskRequestHook,
ResponseHook as FlaskResponseHook,
)
from ..integrations.httpx import (
AsyncRequestHook as HttpxAsyncRequestHook,
AsyncResponseHook as HttpxAsyncResponseHook,
RequestHook as HttpxRequestHook,
ResponseHook as HttpxResponseHook,
)
from ..integrations.wsgi import (
RequestHook as WSGIRequestHook,
ResponseHook as WSGIResponseHook,
)
from .integrations.asgi import ASGIApp, ASGIInstrumentKwargs
from .integrations.aws_lambda import LambdaEvent, LambdaHandler
from .integrations.httpx import AsyncClientKwargs, ClientKwargs, HTTPXInstrumentKwargs
from .integrations.mysql import MySQLConnection, MySQLInstrumentKwargs
from .integrations.psycopg import PsycopgInstrumentKwargs
from .integrations.pymongo import PymongoInstrumentKwargs
Expand Down Expand Up @@ -1185,47 +1190,62 @@ def instrument_httpx(
capture_response_json_body: bool = False,
capture_response_text_body: bool = False,
capture_request_form_data: bool = False,
**kwargs: Unpack[ClientKwargs],
request_hook: HttpxRequestHook | None = None,
response_hook: HttpxResponseHook | None = None,
**kwargs: Any,
) -> None: ...

@overload
def instrument_httpx(
self,
client: httpx.AsyncClient,
*,
capture_all: bool = False,
capture_headers: bool = False,
capture_request_json_body: bool = False,
capture_request_text_body: bool = False,
capture_response_json_body: bool = False,
capture_response_text_body: bool = False,
capture_request_form_data: bool = False,
**kwargs: Unpack[AsyncClientKwargs],
request_hook: HttpxRequestHook | HttpxAsyncRequestHook | None = None,
response_hook: HttpxResponseHook | HttpxAsyncResponseHook | None = None,
**kwargs: Any,
) -> None: ...

@overload
def instrument_httpx(
self,
client: None = None,
*,
capture_all: bool = False,
capture_headers: bool = False,
capture_request_json_body: bool = False,
capture_request_text_body: bool = False,
capture_response_json_body: bool = False,
capture_response_text_body: bool = False,
capture_request_form_data: bool = False,
**kwargs: Unpack[HTTPXInstrumentKwargs],
request_hook: HttpxRequestHook | None = None,
response_hook: HttpxResponseHook | None = None,
async_request_hook: HttpxAsyncRequestHook | None = None,
async_response_hook: HttpxAsyncResponseHook | None = None,
**kwargs: Any,
) -> None: ...

def instrument_httpx(
self,
client: httpx.Client | httpx.AsyncClient | None = None,
*,
capture_all: bool = False,
capture_headers: bool = False,
capture_request_json_body: bool = False,
capture_request_text_body: bool = False,
capture_response_json_body: bool = False,
capture_response_text_body: bool = False,
capture_request_form_data: bool = False,
request_hook: HttpxRequestHook | HttpxAsyncRequestHook | None = None,
response_hook: HttpxResponseHook | HttpxAsyncResponseHook | None = None,
async_request_hook: HttpxAsyncRequestHook | None = None,
async_response_hook: HttpxAsyncResponseHook | None = None,
**kwargs: Any,
) -> None:
"""Instrument the `httpx` module so that spans are automatically created for each request.
Expand All @@ -1239,6 +1259,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 request and response headers and 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
Expand All @@ -1261,6 +1282,10 @@ def instrument_httpx(
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.
request_hook: A function called right after a span is created for a request.
response_hook: A function called right before a span is finished for the response.
async_request_hook: A function called right after a span is created for an async request.
async_response_hook: A function called right before a span is finished for an async response.
**kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` method, for future compatibility.
"""
from .integrations.httpx import instrument_httpx
Expand All @@ -1269,12 +1294,17 @@ def instrument_httpx(
return instrument_httpx(
self,
client,
capture_all=capture_all,
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_response_text_body=capture_response_text_body,
capture_request_form_data=capture_request_form_data,
request_hook=request_hook,
response_hook=response_hook,
async_request_hook=async_request_hook,
async_response_hook=async_response_hook,
**kwargs,
)

Expand Down
Loading
Loading