httpx: support capture_headers (#719)

docs/integrations/http-clients/httpx.md

@@ -71,5 +71,80 @@ Let's see a minimal example below. You can run it with `python main.py`:
 **OpenTelemetry HTTPX Instrumentation** package,
 which you can find more information about [here][opentelemetry-httpx].
 
+## Fine Tuning
+
+The `logfire.instrument_httpx()` method accepts different parameters to fine-tune the instrumentation.
+
+### Capture HTTP Headers
+
+By default, **Logfire** doesn't capture HTTP headers. You can enable it by setting the `capture_headers` parameter to `True`.
+
+```py
+import httpx
+import logfire
+
+logfire.configure()
+logfire.instrument_httpx(capture_headers=True)
+
+client = httpx.Client()
+client.get("https://httpbin.org/get")
+```
+
+#### Capture Only Request Headers
+
+Instead of capturing both request and response headers, you can create a request hook to capture only the request headers:
+
+```py
+import httpx
+import logfire
+from logfire.integrations.httpx import RequestInfo
+from opentelemetry.trace import Span
+
+
+def capture_request_headers(span: Span, request: RequestInfo):
+    headers = request.headers
+    span.set_attributes(
+        {
+            f'http.request.header.{header_name}': headers.get_list(header_name)
+            for header_name in headers.keys()
+        }
+    )
+
+
+logfire.configure()
+logfire.instrument_httpx(request_hook=capture_request_headers)
+
+client = httpx.Client()
+client.get("https://httpbin.org/get")
+```
+
+#### Capture Only Response Headers
+
+Similarly, you can create a response hook to capture only the response headers:
+
+```py
+import httpx
+import logfire
+from opentelemetry.trace import Span
+from logfire.integrations.httpx import RequestInfo, ResponseInfo
+
+
+def capture_response_headers(span: Span, request: RequestInfo, response: ResponseInfo):
+    headers = response.headers
+    span.set_attributes(
+        {f'http.response.header.{header_name}': headers.get_list(header_name)
+        for header_name in headers.keys()}
+    )
+
+
+logfire.configure()
+logfire.instrument_httpx(response_hook=capture_response_headers)
+
+client = httpx.Client()
+client.get('https://httpbin.org/get')
+```
+
+You can also use the hooks to filter headers or modify them before capturing them.
+
 [httpx]: https://www.python-httpx.org/
 [opentelemetry-httpx]: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/httpx/httpx.html

logfire-api/logfire_api/_internal/integrations/httpx.pyi

@@ -1,11 +1,12 @@
 import httpx
+from _typeshed import Incomplete
 from logfire import Logfire as Logfire
 from logfire._internal.main import set_user_attributes_on_raw_span as set_user_attributes_on_raw_span
 from logfire._internal.utils import handle_internal_errors as handle_internal_errors
+from logfire.integrations.httpx import AsyncRequestHook as AsyncRequestHook, AsyncResponseHook as AsyncResponseHook, RequestHook as RequestHook, RequestInfo as RequestInfo, ResponseHook as ResponseHook, ResponseInfo as ResponseInfo
 from logfire.propagate import attach_context as attach_context, get_context as get_context
-from opentelemetry.instrumentation.httpx import AsyncRequestHook, AsyncResponseHook, RequestHook, RequestInfo, ResponseHook, ResponseInfo
 from opentelemetry.trace import Span
-from typing import Any, Callable, Literal, ParamSpec, TypeVar, TypedDict, Unpack, overload
+from typing import Any, Callable, Literal, ParamSpec, TypeVar, TypedDict
 
 class AsyncClientKwargs(TypedDict, total=False):
     request_hook: RequestHook | AsyncRequestHook
@@ -29,14 +30,14 @@ Hook = TypeVar('Hook', RequestHook, ResponseHook)
 AsyncHook = TypeVar('AsyncHook', AsyncRequestHook, AsyncResponseHook)
 P = ParamSpec('P')
 
-@overload
-def instrument_httpx(logfire_instance: Logfire, client: httpx.Client, capture_request_headers: bool, capture_response_headers: bool, capture_request_json_body: bool, capture_response_json_body: bool, **kwargs: Unpack[ClientKwargs]) -> None: ...
-@overload
-def instrument_httpx(logfire_instance: Logfire, client: httpx.AsyncClient, capture_request_headers: bool, capture_response_headers: bool, capture_request_json_body: bool, capture_response_json_body: bool, **kwargs: Unpack[AsyncClientKwargs]) -> None: ...
-@overload
-def instrument_httpx(logfire_instance: Logfire, client: None, capture_request_headers: bool, capture_response_headers: bool, capture_request_json_body: bool, capture_response_json_body: bool, **kwargs: Unpack[HTTPXInstrumentKwargs]) -> None: ...
-def make_request_hook(hook: RequestHook | None, should_capture_headers: bool, should_capture_json: bool) -> RequestHook | None: ...
-def make_async_request_hook(hook: AsyncRequestHook | RequestHook | None, should_capture_headers: bool, should_capture_json: bool) -> AsyncRequestHook | None: ...
+def instrument_httpx(logfire_instance: Logfire, client: httpx.Client | httpx.AsyncClient | None, capture_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.
+
+    See the `Logfire.instrument_httpx` method for details.
+    """
+def make_request_hook(hook: RequestHook | None, should_capture_headers: bool, should_capture_json: bool, should_capture_form_data: bool) -> RequestHook | None: ...
+def make_async_request_hook(hook: AsyncRequestHook | RequestHook | None, should_capture_headers: bool, should_capture_json: bool, should_capture_form_data: bool) -> AsyncRequestHook | None: ...
+def capture_request(request: RequestInfo, span: Span, should_capture_headers: bool, should_capture_json: bool, should_capture_form_data: bool) -> None: ...
 def make_response_hook(hook: ResponseHook | None, should_capture_headers: bool, should_capture_json: bool, logfire_instance: Logfire) -> ResponseHook | None: ...
 def make_async_response_hook(hook: ResponseHook | AsyncResponseHook | None, should_capture_headers: bool, should_capture_json: bool, logfire_instance: Logfire) -> AsyncResponseHook | None: ...
 def capture_response_json(logfire_instance: Logfire, response_info: ResponseInfo, is_async: bool) -> None: ...
@@ -48,3 +49,7 @@ def capture_headers(span: Span, headers: httpx.Headers, request_or_response: Lit
 def get_charset(content_type: str) -> str: ...
 def decode_body(body: bytes, content_type: str): ...
 def capture_request_body(span: Span, request: RequestInfo) -> None: ...
+
+CODES_FOR_METHODS_WITH_DATA_PARAM: Incomplete
+
+def capture_request_form_data(span: Span, request: RequestInfo) -> None: ...

logfire-api/logfire_api/_internal/integrations/sqlalchemy.pyi

@@ -1,4 +1,5 @@
 from sqlalchemy import Engine
+from sqlalchemy.ext.asyncio import AsyncEngine
 from typing_extensions import TypedDict, Unpack
 
 class CommenterOptions(TypedDict, total=False):
@@ -7,12 +8,11 @@ class CommenterOptions(TypedDict, total=False):
     opentelemetry_values: bool
 
 class SQLAlchemyInstrumentKwargs(TypedDict, total=False):
-    engine: Engine | None
     enable_commenter: bool | None
     commenter_options: CommenterOptions | None
     skip_dep_check: bool
 
-def instrument_sqlalchemy(**kwargs: Unpack[SQLAlchemyInstrumentKwargs]) -> None:
+def instrument_sqlalchemy(engine: AsyncEngine | Engine | None, **kwargs: Unpack[SQLAlchemyInstrumentKwargs]) -> None:
     """Instrument the `sqlalchemy` module so that spans are automatically created for each query.
 
     See the `Logfire.instrument_sqlalchemy` method for details.

logfire-api/logfire_api/_internal/main.pyi

@@ -38,6 +38,8 @@ from opentelemetry.metrics import CallbackT as CallbackT, Counter, Histogram, Up
 from opentelemetry.sdk.trace import ReadableSpan, Span
 from opentelemetry.trace import SpanContext, Tracer
 from opentelemetry.util import types as otel_types
+from sqlalchemy import Engine
+from sqlalchemy.ext.asyncio import AsyncEngine
 from starlette.applications import Starlette
 from starlette.requests import Request as Request
 from starlette.websockets import WebSocket as WebSocket
@@ -551,11 +553,11 @@ class Logfire:
     def instrument_asyncpg(self, **kwargs: Unpack[AsyncPGInstrumentKwargs]) -> None:
         """Instrument the `asyncpg` module so that spans are automatically created for each query."""
     @overload
-    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, **kwargs: Unpack[ClientKwargs]) -> None: ...
+    def instrument_httpx(self, client: httpx.Client, *, capture_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, **kwargs: Unpack[AsyncClientKwargs]) -> None: ...
+    def instrument_httpx(self, client: httpx.AsyncClient, *, capture_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, **kwargs: Unpack[HTTPXInstrumentKwargs]) -> None: ...
+    def instrument_httpx(self, client: None = None, *, capture_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_celery(self, **kwargs: Any) -> None:
         """Instrument `celery` so that spans are automatically created for each task.
 
@@ -700,12 +702,16 @@ class Logfire:
         [OpenTelemetry aiohttp client Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/aiohttp_client/aiohttp_client.html)
         library, specifically `AioHttpClientInstrumentor().instrument()`, to which it passes `**kwargs`.
         """
-    def instrument_sqlalchemy(self, **kwargs: Unpack[SQLAlchemyInstrumentKwargs]) -> None:
+    def instrument_sqlalchemy(self, engine: AsyncEngine | Engine | None = None, **kwargs: Unpack[SQLAlchemyInstrumentKwargs]) -> None:
         """Instrument the `sqlalchemy` module so that spans are automatically created for each query.
 
         Uses the
         [OpenTelemetry SQLAlchemy Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/sqlalchemy/sqlalchemy.html)
         library, specifically `SQLAlchemyInstrumentor().instrument()`, to which it passes `**kwargs`.
+
+        Args:
+            engine: The `sqlalchemy` engine to instrument, or `None` to instrument all engines.
+            **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods.
         """
     def instrument_sqlite3(self, conn: SQLite3Connection = None, **kwargs: Unpack[SQLite3InstrumentKwargs]) -> SQLite3Connection:
         """Instrument the `sqlite3` module or a specific connection so that spans are automatically created for each operation.

logfire-api/logfire_api/integrations/httpx.pyi

@@ -0,0 +1,28 @@
+import httpx
+from opentelemetry.trace import Span
+from typing import Any, Awaitable, Callable, NamedTuple
+
+class RequestInfo(NamedTuple):
+    """Information about an HTTP request.
+
+    This is the second parameter passed to the `RequestHook` function.
+    """
+    method: bytes
+    url: httpx.URL
+    headers: httpx.Headers
+    stream: httpx.SyncByteStream | httpx.AsyncByteStream | None
+    extensions: dict[str, Any] | None
+
+class ResponseInfo(NamedTuple):
+    """Information about an HTTP response.
+
+    This is the second parameter passed to the `ResponseHook` function.
+    """
+    status_code: int
+    headers: httpx.Headers
+    stream: httpx.SyncByteStream | httpx.AsyncByteStream | None
+    extensions: dict[str, Any] | None
+RequestHook = Callable[[Span, RequestInfo], None]
+ResponseHook = Callable[[Span, RequestInfo, ResponseInfo], None]
+AsyncRequestHook = Callable[[Span, RequestInfo], Awaitable[None]]
+AsyncResponseHook = Callable[[Span, RequestInfo, ResponseInfo], Awaitable[None]]