Add ASGI & WSGI instrument methods (#324)

Co-authored-by: Alex Hall <[email protected]>

Makefile

@@ -30,7 +30,7 @@ typecheck:
 
 .PHONY: test  # Run the tests
 test:
-	uv run coverage run -m pytest
+	uv run --no-sync coverage run -m pytest
 
 .PHONY: generate-stubs  # Generate stubs for logfire-api
 generate-stubs:

docs/integrations/asgi.md

@@ -1,23 +1,20 @@
 # ASGI
 
-If the [ASGI][asgi] framework doesn't have a dedicated OpenTelemetry package, you can use the
-[OpenTelemetry ASGI middleware][opentelemetry-asgi].
+If the [ASGI][asgi] web framework you're using doesn't have a dedicated integration, you can use the
+[`logfire.instrument_asgi()`][logfire.Logfire.instrument_asgi] method to instrument it.
 
 ## Installation
 
-You need to install the `opentelemetry-instrumentation-asgi` package:
+Install `logfire` with the `asgi` extra:
 
-```bash
-pip install opentelemetry-instrumentation-asgi
-```
+{{ install_logfire(extras=['asgi']) }}
 
 ## Usage
 
 Below we have a minimal example using [Uvicorn][uvicorn]. You can run it with `python main.py`:
 
 ```py title="main.py"
 import logfire
-from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware
 
 
 logfire.configure()
@@ -34,23 +31,25 @@ async def app(scope, receive, send):
     )
     await send({"type": "http.response.body", "body": b"Hello, world!"})
 
-app = OpenTelemetryMiddleware(app)
+app = logfire.instrument_asgi(app)
 
 if __name__ == "__main__":
     import uvicorn
 
     uvicorn.run(app)
 ```
 
-You can read more about the OpenTelemetry ASGI middleware [here][opentelemetry-asgi].
+The keyword arguments of [`logfire.instrument_asgi()`][logfire.Logfire.instrument_asgi] are passed to the
+[`OpenTelemetryMiddleware`][opentelemetry.instrumentation.asgi.OpenTelemetryMiddleware] class
+of the OpenTelemetry ASGI Instrumentation package.
 
 ## Excluding URLs from instrumentation
 <!-- note that this section is duplicated for different frameworks but with slightly different links -->
 
 - [Quick guide](use-cases/web-frameworks.md#excluding-urls-from-instrumentation)
 
 !!! note
-    `OpenTelemetryMiddleware` does accept an `excluded_urls` parameter, but does not support specifying said URLs via an environment variable,
+    `instrument_asgi` does accept an `excluded_urls` parameter, but does not support specifying said URLs via an environment variable,
     unlike other instrumentations.
 
 ## Capturing request and response headers
@@ -60,5 +59,4 @@ You can read more about the OpenTelemetry ASGI middleware [here][opentelemetry-a
 - [OpenTelemetry Documentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/asgi/asgi.html#capture-http-request-and-response-headers)
 
 [asgi]: https://asgi.readthedocs.io/en/latest/
-[opentelemetry-asgi]: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/asgi/asgi.html
 [uvicorn]: https://www.uvicorn.org/

docs/integrations/wsgi.md

@@ -1,15 +1,13 @@
 # WSGI
 
-If the [WSGI][wsgi] framework doesn't have a dedicated OpenTelemetry package, you can use the
-[OpenTelemetry WSGI middleware][opentelemetry-wsgi].
+If the [WSGI][wsgi] web framework you're using doesn't have a dedicated integration, you can use the
+[`logfire.instrument_wsgi()`][logfire.Logfire.instrument_wsgi] method to instrument it.
 
 ## Installation
 
-You need to install the [`opentelemetry-instrumentation-wsgi`][pypi-otel-wsgi] package:
+Install `logfire` with the `wsgi` extra:
 
-```bash
-pip install opentelemetry-instrumentation-wsgi
-```
+{{ install_logfire(extras=['wsgi']) }}
 
 ## Usage
 
@@ -18,14 +16,16 @@ Below we have a minimal example using the standard library [`wsgiref`][wsgiref].
 ```py title="main.py"
 from wsgiref.simple_server import make_server
 
-from opentelemetry.instrumentation.wsgi import OpenTelemetryMiddleware
+import logfire
+
 
+logfire.configure()
 
 def app(env, start_response):
     start_response('200 OK', [('Content-Type','text/html')])
     return [b"Hello World"]
 
-app = OpenTelemetryMiddleware(app)
+app = logfire.instrument_wsgi(app)
 
 with make_server("", 8000, app) as httpd:
     print("Serving on port 8000...")
@@ -34,7 +34,10 @@ with make_server("", 8000, app) as httpd:
     httpd.serve_forever()
 ```
 
-You can read more about the OpenTelemetry WSGI middleware [here][opentelemetry-wsgi].
+The keyword arguments of [`logfire.instrument_wsgi()`][logfire.Logfire.instrument_wsgi] are passed to the
+[`OpenTelemetryMiddleware`][opentelemetry.instrumentation.wsgi.OpenTelemetryMiddleware] class of
+the OpenTelemetry WSGI Instrumentation package.
+
 
 ## Capturing request and response headers
 <!-- note that this section is duplicated for different frameworks but with slightly different links -->
@@ -43,6 +46,4 @@ You can read more about the OpenTelemetry WSGI middleware [here][opentelemetry-w
 - [OpenTelemetry Documentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/wsgi/wsgi.html#capture-http-request-and-response-headers)
 
 [wsgi]: https://wsgi.readthedocs.io/en/latest/
-[opentelemetry-wsgi]: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/wsgi/wsgi.html
-[pypi-otel-wsgi]: https://pypi.org/project/opentelemetry-instrumentation-wsgi/
 [wsgiref]: https://docs.python.org/3/library/wsgiref.html

logfire-api/logfire_api/__init__.py

@@ -94,6 +94,12 @@ def decorator(func):
 
                 return decorator
 
+            def instrument_asgi(self, app, *args, **kwargs):
+                return app
+
+            def instrument_wsgi(self, app, *args, **kwargs):
+                return app
+
             def instrument_fastapi(self, *args, **kwargs) -> ContextManager[None]:
                 return nullcontext()
 
@@ -148,6 +154,8 @@ def shutdown(self, *args, **kwargs) -> None: ...
         log_slow_async_callbacks = DEFAULT_LOGFIRE_INSTANCE.log_slow_async_callbacks
         install_auto_tracing = DEFAULT_LOGFIRE_INSTANCE.install_auto_tracing
         instrument = DEFAULT_LOGFIRE_INSTANCE.instrument
+        instrument_asgi = DEFAULT_LOGFIRE_INSTANCE.instrument_asgi
+        instrument_wsgi = DEFAULT_LOGFIRE_INSTANCE.instrument_wsgi
         instrument_pydantic = DEFAULT_LOGFIRE_INSTANCE.instrument_pydantic
         instrument_fastapi = DEFAULT_LOGFIRE_INSTANCE.instrument_fastapi
         instrument_openai = DEFAULT_LOGFIRE_INSTANCE.instrument_openai

logfire-api/logfire_api/__init__.pyi

@@ -12,7 +12,7 @@ from .version import VERSION as VERSION
 from logfire.sampling import SamplingOptions as SamplingOptions
 from typing import Any
 
-__all__ = ['Logfire', 'LogfireSpan', 'LevelName', 'AdvancedOptions', 'ConsoleOptions', 'CodeSource', 'PydanticPlugin', 'configure', 'span', 'instrument', 'log', 'trace', 'debug', 'notice', 'info', 'warn', 'error', 'exception', 'fatal', 'force_flush', 'log_slow_async_callbacks', 'install_auto_tracing', 'instrument_pydantic', 'instrument_fastapi', 'instrument_openai', 'instrument_anthropic', 'instrument_asyncpg', 'instrument_httpx', 'instrument_celery', 'instrument_requests', 'instrument_psycopg', 'instrument_django', 'instrument_flask', 'instrument_starlette', 'instrument_aiohttp_client', 'instrument_sqlalchemy', 'instrument_redis', 'instrument_pymongo', 'instrument_mysql', 'instrument_system_metrics', 'AutoTraceModule', 'with_tags', 'with_settings', 'shutdown', 'load_spans_from_file', 'no_auto_trace', 'ScrubMatch', 'ScrubbingOptions', 'VERSION', 'suppress_instrumentation', 'StructlogProcessor', 'LogfireLoggingHandler', 'loguru_handler', 'SamplingOptions', 'MetricsOptions']
+__all__ = ['Logfire', 'LogfireSpan', 'LevelName', 'AdvancedOptions', 'ConsoleOptions', 'CodeSource', 'PydanticPlugin', 'configure', 'span', 'instrument', 'log', 'trace', 'debug', 'notice', 'info', 'warn', 'error', 'exception', 'fatal', 'force_flush', 'log_slow_async_callbacks', 'install_auto_tracing', 'instrument_asgi', 'instrument_wsgi', 'instrument_pydantic', 'instrument_fastapi', 'instrument_openai', 'instrument_anthropic', 'instrument_asyncpg', 'instrument_httpx', 'instrument_celery', 'instrument_requests', 'instrument_psycopg', 'instrument_django', 'instrument_flask', 'instrument_starlette', 'instrument_aiohttp_client', 'instrument_sqlalchemy', 'instrument_redis', 'instrument_pymongo', 'instrument_mysql', 'instrument_system_metrics', 'AutoTraceModule', 'with_tags', 'with_settings', 'shutdown', 'load_spans_from_file', 'no_auto_trace', 'ScrubMatch', 'ScrubbingOptions', 'VERSION', 'suppress_instrumentation', 'StructlogProcessor', 'LogfireLoggingHandler', 'loguru_handler', 'SamplingOptions', 'MetricsOptions']
 
 DEFAULT_LOGFIRE_INSTANCE = Logfire()
 span = DEFAULT_LOGFIRE_INSTANCE.span
@@ -21,6 +21,8 @@ force_flush = DEFAULT_LOGFIRE_INSTANCE.force_flush
 log_slow_async_callbacks = DEFAULT_LOGFIRE_INSTANCE.log_slow_async_callbacks
 install_auto_tracing = DEFAULT_LOGFIRE_INSTANCE.install_auto_tracing
 instrument_pydantic = DEFAULT_LOGFIRE_INSTANCE.instrument_pydantic
+instrument_asgi = DEFAULT_LOGFIRE_INSTANCE.instrument_asgi
+instrument_wsgi = DEFAULT_LOGFIRE_INSTANCE.instrument_wsgi
 instrument_fastapi = DEFAULT_LOGFIRE_INSTANCE.instrument_fastapi
 instrument_openai = DEFAULT_LOGFIRE_INSTANCE.instrument_openai
 instrument_anthropic = DEFAULT_LOGFIRE_INSTANCE.instrument_anthropic

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

@@ -1,9 +1,28 @@
 from dataclasses import dataclass
 from logfire import Logfire as Logfire
-from logfire._internal.utils import is_asgi_send_receive_span_name as is_asgi_send_receive_span_name
+from logfire._internal.utils import is_asgi_send_receive_span_name as is_asgi_send_receive_span_name, maybe_capture_server_headers as maybe_capture_server_headers
 from opentelemetry.context import Context
 from opentelemetry.trace import Span, Tracer, TracerProvider
-from typing import Any
+from typing import Any, Awaitable, Callable, Protocol, TypedDict
+from typing_extensions import Unpack
+
+Scope = dict[str, Any]
+Receive = Callable[[], Awaitable[dict[str, Any]]]
+Send = Callable[[dict[str, Any]], Awaitable[None]]
+
+class ASGIApp(Protocol):
+    def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: ...
+Hook = Callable[[Span, dict[str, Any]], None]
+
+class ASGIInstrumentKwargs(TypedDict, total=False):
+    excluded_urls: str | None
+    default_span_details: Callable[[Scope], tuple[str, dict[str, Any]]]
+    server_request_hook: Hook | None
+    client_request_hook: Hook | None
+    client_response_hook: Hook | None
+    http_capture_headers_server_request: list[str] | None
+    http_capture_headers_server_response: list[str] | None
+    http_capture_headers_sanitize_fields: list[str] | None
 
 def tweak_asgi_spans_tracer_provider(logfire_instance: Logfire, record_send_receive: bool) -> TracerProvider:
     """If record_send_receive is False, return a TracerProvider that skips spans for ASGI send and receive events."""
@@ -18,3 +37,9 @@ class TweakAsgiSpansTracer(Tracer):
     tracer: Tracer
     def start_span(self, name: str, context: Context | None = None, *args: Any, **kwargs: Any) -> Span: ...
     start_as_current_span = ...
+
+def instrument_asgi(logfire_instance: Logfire, app: ASGIApp, *, record_send_receive: bool = False, capture_headers: bool = False, **kwargs: Unpack[ASGIInstrumentKwargs]) -> ASGIApp:
+    """Instrument `app` so that spans are automatically created for each request.
+
+    See the `Logfire.instrument_asgi` method for details.
+    """

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

@@ -0,0 +1,20 @@
+from logfire import Logfire as Logfire
+from logfire._internal.utils import maybe_capture_server_headers as maybe_capture_server_headers
+from opentelemetry.trace import Span
+from typing import Callable, Protocol, TypedDict
+from typing_extensions import Unpack
+from wsgiref.types import WSGIApplication, WSGIEnvironment
+
+class ResponseHook(Protocol):
+    def __call__(self, span: Span, environ: WSGIEnvironment, status_code: int, response_headers: list[tuple[str, str]]) -> None: ...
+RequestHook = Callable[[Span, WSGIEnvironment], None]
+
+class WSGIInstrumentKwargs(TypedDict, total=False):
+    request_hook: RequestHook | None
+    response_hook: ResponseHook | None
+
+def instrument_wsgi(logfire_instance: Logfire, app: WSGIApplication, *, capture_headers: bool = False, **kwargs: Unpack[WSGIInstrumentKwargs]) -> WSGIApplication:
+    """Instrument `app` so that spans are automatically created for each request.
+
+    See the `Logfire.instrument_wsgi` method for details.
+    """

logfire-api/logfire_api/_internal/main.pyi

@@ -9,6 +9,7 @@ from .config_params import PydanticPluginRecordValues as PydanticPluginRecordVal
 from .constants import ATTRIBUTES_JSON_SCHEMA_KEY as ATTRIBUTES_JSON_SCHEMA_KEY, ATTRIBUTES_LOG_LEVEL_NUM_KEY as ATTRIBUTES_LOG_LEVEL_NUM_KEY, ATTRIBUTES_MESSAGE_KEY as ATTRIBUTES_MESSAGE_KEY, ATTRIBUTES_MESSAGE_TEMPLATE_KEY as ATTRIBUTES_MESSAGE_TEMPLATE_KEY, ATTRIBUTES_SAMPLE_RATE_KEY as ATTRIBUTES_SAMPLE_RATE_KEY, ATTRIBUTES_SPAN_TYPE_KEY as ATTRIBUTES_SPAN_TYPE_KEY, ATTRIBUTES_TAGS_KEY as ATTRIBUTES_TAGS_KEY, ATTRIBUTES_VALIDATION_ERROR_KEY as ATTRIBUTES_VALIDATION_ERROR_KEY, DISABLE_CONSOLE_KEY as DISABLE_CONSOLE_KEY, LEVEL_NUMBERS as LEVEL_NUMBERS, LevelName as LevelName, NULL_ARGS_KEY as NULL_ARGS_KEY, OTLP_MAX_INT_SIZE as OTLP_MAX_INT_SIZE, log_level_attributes as log_level_attributes
 from .formatter import logfire_format as logfire_format, logfire_format_with_magic as logfire_format_with_magic
 from .instrument import instrument as instrument
+from .integrations.asgi import ASGIApp as ASGIApp, ASGIInstrumentKwargs as ASGIInstrumentKwargs
 from .integrations.asyncpg import AsyncPGInstrumentKwargs as AsyncPGInstrumentKwargs
 from .integrations.celery import CeleryInstrumentKwargs as CeleryInstrumentKwargs
 from .integrations.flask import FlaskInstrumentKwargs as FlaskInstrumentKwargs
@@ -20,6 +21,7 @@ from .integrations.redis import RedisInstrumentKwargs as RedisInstrumentKwargs
 from .integrations.sqlalchemy import SQLAlchemyInstrumentKwargs as SQLAlchemyInstrumentKwargs
 from .integrations.starlette import StarletteInstrumentKwargs as StarletteInstrumentKwargs
 from .integrations.system_metrics import Base as SystemMetricsBase, Config as SystemMetricsConfig
+from .integrations.wsgi import WSGIInstrumentKwargs as WSGIInstrumentKwargs
 from .json_encoder import logfire_json_dumps as logfire_json_dumps
 from .json_schema import JsonSchemaProperties as JsonSchemaProperties, attributes_json_schema as attributes_json_schema, attributes_json_schema_properties as attributes_json_schema_properties, create_json_schema as create_json_schema
 from .metrics import ProxyMeterProvider as ProxyMeterProvider
@@ -38,6 +40,7 @@ from starlette.requests import Request as Request
 from starlette.websockets import WebSocket as WebSocket
 from typing import Any, Callable, ContextManager, Iterable, Literal, Sequence, TypeVar
 from typing_extensions import LiteralString, ParamSpec, Unpack
+from wsgiref.types import WSGIApplication
 
 ExcInfo = SysExcInfo | BaseException | bool | None
 
@@ -600,25 +603,69 @@ class Logfire:
     def instrument_flask(self, app: Flask, *, capture_headers: bool = False, **kwargs: Unpack[FlaskInstrumentKwargs]) -> None:
         """Instrument `app` so that spans are automatically created for each request.
 
-        Set `capture_headers` to `True` to capture all request and response headers.
-
         Uses the
         [OpenTelemetry Flask Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/flask/flask.html)
         library, specifically `FlaskInstrumentor().instrument_app()`, to which it passes `**kwargs`.
+
+        Args:
+            app: The Flask app to instrument.
+            capture_headers: Set to `True` to capture all request and response headers.
+            **kwargs: Additional keyword arguments to pass to the OpenTelemetry Flask instrumentation.
         """
     def instrument_starlette(self, app: Starlette, *, capture_headers: bool = False, record_send_receive: bool = False, **kwargs: Unpack[StarletteInstrumentKwargs]) -> None:
         """Instrument `app` so that spans are automatically created for each request.
 
-        Set `capture_headers` to `True` to capture all request and response headers.
-
-        Set `record_send_receive` to `True` to allow the OpenTelemetry ASGI to create send/receive spans.
-        These are disabled by default to reduce overhead and the number of spans created,
-        since many can be created for a single request, and they are not often useful.
-        If enabled, they will be set to debug level, meaning they will usually still be hidden in the UI.
-
         Uses the
         [OpenTelemetry Starlette Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/starlette/starlette.html)
         library, specifically `StarletteInstrumentor.instrument_app()`, to which it passes `**kwargs`.
+
+        Args:
+            app: The Starlette app to instrument.
+            capture_headers: Set to `True` to capture all request and response headers.
+            record_send_receive: Set `record_send_receive` to `True` to allow the OpenTelemetry ASGI to create send/receive spans.
+
+                These are disabled by default to reduce overhead and the number of spans created,
+                since many can be created for a single request, and they are not often useful.
+                If enabled, they will be set to debug level, meaning they will usually still be hidden in the UI.
+            **kwargs: Additional keyword arguments to pass to the OpenTelemetry Starlette instrumentation.
+        """
+    def instrument_asgi(self, app: ASGIApp, capture_headers: bool = False, record_send_receive: bool = False, **kwargs: Unpack[ASGIInstrumentKwargs]) -> ASGIApp:
+        """Instrument `app` so that spans are automatically created for each request.
+
+        Uses the ASGI [`OpenTelemetryMiddleware`][opentelemetry.instrumentation.asgi.OpenTelemetryMiddleware] under
+        the hood, to which it passes `**kwargs`.
+
+        Warning:
+            Instead of modifying the app in place, this method returns the instrumented ASGI application.
+
+        Args:
+            app: The ASGI application to instrument.
+            capture_headers: Set to `True` to capture all request and response headers.
+            record_send_receive: Set to `True` to allow the OpenTelemetry ASGI to create send/receive spans.
+                These are disabled by default to reduce overhead and the number of spans created,
+                since many can be created for a single request, and they are not often useful.
+                If enabled, they will be set to debug level, meaning they will usually still be hidden in the UI.
+            **kwargs: Additional keyword arguments to pass to the OpenTelemetry ASGI middleware.
+
+        Returns:
+            The instrumented ASGI application.
+        """
+    def instrument_wsgi(self, app: WSGIApplication, capture_headers: bool = False, **kwargs: Unpack[WSGIInstrumentKwargs]) -> WSGIApplication:
+        """Instrument `app` so that spans are automatically created for each request.
+
+        Uses the WSGI [`OpenTelemetryMiddleware`][opentelemetry.instrumentation.wsgi.OpenTelemetryMiddleware] under
+        the hood, to which it passes `**kwargs`.
+
+        Warning:
+            Instead of modifying the app in place, this method returns the instrumented WSGI application.
+
+        Args:
+            app: The WSGI application to instrument.
+            capture_headers: Set to `True` to capture all request and response headers.
+            **kwargs: Additional keyword arguments to pass to the OpenTelemetry WSGI middleware.
+
+        Returns:
+            The instrumented WSGI application.
         """
     def instrument_aiohttp_client(self, **kwargs: Any) -> None:
         """Instrument the `aiohttp` module so that spans are automatically created for each client request.

logfire/__init__.py

@@ -32,6 +32,8 @@
 log_slow_async_callbacks = DEFAULT_LOGFIRE_INSTANCE.log_slow_async_callbacks
 install_auto_tracing = DEFAULT_LOGFIRE_INSTANCE.install_auto_tracing
 instrument_pydantic = DEFAULT_LOGFIRE_INSTANCE.instrument_pydantic
+instrument_asgi = DEFAULT_LOGFIRE_INSTANCE.instrument_asgi
+instrument_wsgi = DEFAULT_LOGFIRE_INSTANCE.instrument_wsgi
 instrument_fastapi = DEFAULT_LOGFIRE_INSTANCE.instrument_fastapi
 instrument_openai = DEFAULT_LOGFIRE_INSTANCE.instrument_openai
 instrument_anthropic = DEFAULT_LOGFIRE_INSTANCE.instrument_anthropic
@@ -111,6 +113,8 @@ def loguru_handler() -> dict[str, Any]:
     'force_flush',
     'log_slow_async_callbacks',
     'install_auto_tracing',
+    'instrument_asgi',
+    'instrument_wsgi',
     'instrument_pydantic',
     'instrument_fastapi',
     'instrument_openai',