From 7e6b140e5c45d81892fb5976292395fe8b4e056d Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Wed, 18 Dec 2024 11:30:37 +0100 Subject: [PATCH] Document Requests kwargs (#679) --- .../_internal/integrations/celery.pyi | 8 +--- .../_internal/integrations/requests.pyi | 7 ++-- logfire-api/logfire_api/_internal/main.pyi | 14 ++++--- logfire/_internal/integrations/celery.py | 20 ++-------- logfire/_internal/integrations/requests.py | 24 +++++++----- logfire/_internal/main.py | 39 +++++++++++++++---- 6 files changed, 65 insertions(+), 47 deletions(-) diff --git a/logfire-api/logfire_api/_internal/integrations/celery.pyi b/logfire-api/logfire_api/_internal/integrations/celery.pyi index cce9a3094..360dddf5d 100644 --- a/logfire-api/logfire_api/_internal/integrations/celery.pyi +++ b/logfire-api/logfire_api/_internal/integrations/celery.pyi @@ -1,10 +1,6 @@ -from logfire import Logfire as Logfire -from typing_extensions import TypedDict, Unpack +from typing import Any -class CeleryInstrumentKwargs(TypedDict, total=False): - skip_dep_check: bool - -def instrument_celery(logfire_instance: Logfire, **kwargs: Unpack[CeleryInstrumentKwargs]) -> None: +def instrument_celery(**kwargs: Any) -> None: """Instrument the `celery` module so that spans are automatically created for each task. See the `Logfire.instrument_celery` method for details. diff --git a/logfire-api/logfire_api/_internal/integrations/requests.pyi b/logfire-api/logfire_api/_internal/integrations/requests.pyi index a1bbdbaa4..5582e5add 100644 --- a/logfire-api/logfire_api/_internal/integrations/requests.pyi +++ b/logfire-api/logfire_api/_internal/integrations/requests.pyi @@ -1,7 +1,8 @@ -from logfire import Logfire as Logfire -from typing import Any +import requests +from opentelemetry.sdk.trace import Span as Span +from typing import Any, Callable -def instrument_requests(logfire_instance: Logfire, excluded_urls: str | None = None, **kwargs: Any): +def instrument_requests(excluded_urls: str | None = None, request_hook: Callable[[Span, requests.PreparedRequest], None] | None = None, response_hook: Callable[[Span, requests.PreparedRequest, requests.Response], None] | None = None, **kwargs: Any) -> None: """Instrument the `requests` module so that spans are automatically created for each request. See the `Logfire.instrument_requests` method for details. diff --git a/logfire-api/logfire_api/_internal/main.pyi b/logfire-api/logfire_api/_internal/main.pyi index 8539212bc..ab3414953 100644 --- a/logfire-api/logfire_api/_internal/main.pyi +++ b/logfire-api/logfire_api/_internal/main.pyi @@ -2,6 +2,7 @@ import anthropic import httpx import openai import opentelemetry.trace as trace_api +import requests from . import async_ as async_ from ..version import VERSION as VERSION from .auto_trace import AutoTraceModule as AutoTraceModule, install_auto_tracing as install_auto_tracing @@ -13,7 +14,6 @@ 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.aws_lambda import AwsLambdaInstrumentKwargs as AwsLambdaInstrumentKwargs, LambdaHandler as LambdaHandler -from .integrations.celery import CeleryInstrumentKwargs as CeleryInstrumentKwargs from .integrations.flask import FlaskInstrumentKwargs as FlaskInstrumentKwargs from .integrations.httpx import AsyncClientKwargs as AsyncClientKwargs, ClientKwargs as ClientKwargs, HTTPXInstrumentKwargs as HTTPXInstrumentKwargs from .integrations.mysql import MySQLConnection as MySQLConnection, MySQLInstrumentKwargs as MySQLInstrumentKwargs @@ -556,12 +556,15 @@ class Logfire: def instrument_httpx(self, client: httpx.AsyncClient, capture_request_headers: bool = False, capture_response_headers: bool = False, **kwargs: Unpack[AsyncClientKwargs]) -> None: ... @overload def instrument_httpx(self, client: None = None, capture_request_headers: bool = False, capture_response_headers: bool = False, **kwargs: Unpack[HTTPXInstrumentKwargs]) -> None: ... - def instrument_celery(self, **kwargs: Unpack[CeleryInstrumentKwargs]) -> None: + def instrument_celery(self, **kwargs: Any) -> None: """Instrument `celery` so that spans are automatically created for each task. Uses the [OpenTelemetry Celery Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/celery/celery.html) library. + + Args: + **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` method, for future compatibility. """ def instrument_django(self, capture_headers: bool = False, is_sql_commentor_enabled: bool | None = None, request_hook: Callable[[Span, HttpRequest], None] | None = None, response_hook: Callable[[Span, HttpRequest, HttpResponse], None] | None = None, excluded_urls: str | None = None, **kwargs: Any) -> None: """Instrument `django` so that spans are automatically created for each web request. @@ -594,13 +597,14 @@ class Logfire: for future compatibility. """ - def instrument_requests(self, excluded_urls: str | None = None, **kwargs: Any) -> None: + def instrument_requests(self, excluded_urls: str | None = None, request_hook: Callable[[Span, requests.PreparedRequest], None] | None = None, response_hook: Callable[[Span, requests.PreparedRequest, requests.Response], None] | None = None, **kwargs: Any) -> None: """Instrument the `requests` module so that spans are automatically created for each request. Args: excluded_urls: A string containing a comma-delimited list of regexes used to exclude URLs from tracking - **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods, - particularly `request_hook` and `response_hook`. + 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. + **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods, for future compatibility. """ def instrument_psycopg(self, conn_or_module: Any = None, **kwargs: Unpack[PsycopgInstrumentKwargs]) -> None: """Instrument a `psycopg` connection or module so that spans are automatically created for each query. diff --git a/logfire/_internal/integrations/celery.py b/logfire/_internal/integrations/celery.py index 093dc6b71..253a6f756 100644 --- a/logfire/_internal/integrations/celery.py +++ b/logfire/_internal/integrations/celery.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import Any try: from opentelemetry.instrumentation.celery import CeleryInstrumentor @@ -11,24 +11,10 @@ " pip install 'logfire[celery]'" ) -from logfire import Logfire -if TYPE_CHECKING: - from typing_extensions import TypedDict, Unpack - - class CeleryInstrumentKwargs(TypedDict, total=False): - skip_dep_check: bool - - -def instrument_celery(logfire_instance: Logfire, **kwargs: Unpack[CeleryInstrumentKwargs]) -> None: +def instrument_celery(**kwargs: Any) -> None: """Instrument the `celery` module so that spans are automatically created for each task. See the `Logfire.instrument_celery` method for details. """ - return CeleryInstrumentor().instrument( - **{ - 'tracer_provider': logfire_instance.config.get_tracer_provider(), - 'meter_provider': logfire_instance.config.get_meter_provider(), - **kwargs, - } - ) + return CeleryInstrumentor().instrument(**kwargs) diff --git a/logfire/_internal/integrations/requests.py b/logfire/_internal/integrations/requests.py index af9d0857a..3481f99d4 100644 --- a/logfire/_internal/integrations/requests.py +++ b/logfire/_internal/integrations/requests.py @@ -1,4 +1,9 @@ -from typing import Any, Optional +from __future__ import annotations + +from typing import Any, Callable + +import requests +from opentelemetry.sdk.trace import Span try: from opentelemetry.instrumentation.requests import RequestsInstrumentor @@ -9,19 +14,20 @@ " pip install 'logfire[requests]'" ) -from logfire import Logfire - -def instrument_requests(logfire_instance: Logfire, excluded_urls: Optional[str] = None, **kwargs: Any): +def instrument_requests( + excluded_urls: str | None = None, + request_hook: Callable[[Span, requests.PreparedRequest], None] | None = None, + response_hook: Callable[[Span, requests.PreparedRequest, requests.Response], None] | None = None, + **kwargs: Any, +) -> None: """Instrument the `requests` module so that spans are automatically created for each request. See the `Logfire.instrument_requests` method for details. """ RequestsInstrumentor().instrument( excluded_urls=excluded_urls, - **{ - 'tracer_provider': logfire_instance.config.get_tracer_provider(), - 'meter_provider': logfire_instance.config.get_meter_provider(), - **kwargs, - }, + request_hook=request_hook, + response_hook=response_hook, + **kwargs, ) diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index 9449ff041..ed5d9e678 100644 --- a/logfire/_internal/main.py +++ b/logfire/_internal/main.py @@ -72,6 +72,7 @@ import anthropic import httpx import openai + import requests from django.http import HttpRequest, HttpResponse from fastapi import FastAPI from flask.app import Flask @@ -84,7 +85,6 @@ from .integrations.asgi import ASGIApp, ASGIInstrumentKwargs from .integrations.asyncpg import AsyncPGInstrumentKwargs from .integrations.aws_lambda import AwsLambdaInstrumentKwargs, LambdaHandler - from .integrations.celery import CeleryInstrumentKwargs from .integrations.flask import FlaskInstrumentKwargs from .integrations.httpx import AsyncClientKwargs, ClientKwargs, HTTPXInstrumentKwargs from .integrations.mysql import MySQLConnection, MySQLInstrumentKwargs @@ -1224,17 +1224,26 @@ def instrument_httpx( **kwargs, ) - def instrument_celery(self, **kwargs: Unpack[CeleryInstrumentKwargs]) -> None: + def instrument_celery(self, **kwargs: Any) -> None: """Instrument `celery` so that spans are automatically created for each task. Uses the [OpenTelemetry Celery Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/celery/celery.html) library. + + Args: + **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` method, for future compatibility. """ from .integrations.celery import instrument_celery self._warn_if_not_initialized_for_instrumentation() - return instrument_celery(self, **kwargs) + return instrument_celery( + **{ + 'tracer_provider': self._config.get_tracer_provider(), + 'meter_provider': self._config.get_meter_provider(), + **kwargs, + }, + ) def instrument_django( self, @@ -1288,18 +1297,34 @@ def instrument_django( **kwargs, ) - def instrument_requests(self, excluded_urls: str | None = None, **kwargs: Any) -> None: + def instrument_requests( + self, + excluded_urls: str | None = None, + request_hook: Callable[[Span, requests.PreparedRequest], None] | None = None, + response_hook: Callable[[Span, requests.PreparedRequest, requests.Response], None] | None = None, + **kwargs: Any, + ) -> None: """Instrument the `requests` module so that spans are automatically created for each request. Args: excluded_urls: A string containing a comma-delimited list of regexes used to exclude URLs from tracking - **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods, - particularly `request_hook` and `response_hook`. + 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. + **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods, for future compatibility. """ from .integrations.requests import instrument_requests self._warn_if_not_initialized_for_instrumentation() - return instrument_requests(self, excluded_urls=excluded_urls, **kwargs) + return instrument_requests( + excluded_urls=excluded_urls, + request_hook=request_hook, + response_hook=response_hook, + **{ + 'tracer_provider': self._config.get_tracer_provider(), + 'meter_provider': self._config.get_meter_provider(), + **kwargs, + }, + ) def instrument_psycopg(self, conn_or_module: Any = None, **kwargs: Unpack[PsycopgInstrumentKwargs]) -> None: """Instrument a `psycopg` connection or module so that spans are automatically created for each query.