diff --git a/logfire/_internal/integrations/flask.py b/logfire/_internal/integrations/flask.py index 2363f18ae..8d911f643 100644 --- a/logfire/_internal/integrations/flask.py +++ b/logfire/_internal/integrations/flask.py @@ -1,10 +1,34 @@ -from typing import Any +from __future__ import annotations + +from typing import TYPE_CHECKING from flask.app import Flask from opentelemetry.instrumentation.flask import FlaskInstrumentor +if TYPE_CHECKING: + from wsgiref.types import WSGIEnvironment + + from opentelemetry.metrics import MeterProvider + from opentelemetry.trace import Span, TracerProvider + from typing_extensions import Protocol, TypedDict, Unpack + + class RequestHook(Protocol): + def __call__(self, span: Span, environment: WSGIEnvironment) -> None: ... + + class ResponseHook(Protocol): + def __call__(self, span: Span, status: str, response_headers: list[tuple[str, str]]) -> None: ... + + class FlaskInstrumentKwargs(TypedDict, total=False): + request_hook: RequestHook | None + response_hook: RequestHook | None + tracer_provider: TracerProvider | None + excluded_urls: str | None + enable_commenter: bool | None + commenter_options: dict[str, str] | None + meter_provider: MeterProvider | None + -def instrument_flask(app: Flask, **kwargs: Any): +def instrument_flask(app: Flask, **kwargs: Unpack[FlaskInstrumentKwargs]): """Instrument `app` so that spans are automatically created for each request. See the `Logfire.instrument_flask` method for details. diff --git a/logfire/_internal/integrations/psycopg.py b/logfire/_internal/integrations/psycopg.py index 67c36fc09..7f244b942 100644 --- a/logfire/_internal/integrations/psycopg.py +++ b/logfire/_internal/integrations/psycopg.py @@ -11,13 +11,25 @@ if TYPE_CHECKING: # pragma: no cover from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor + from typing_extensions import TypedDict, Unpack Instrumentor = PsycopgInstrumentor | Psycopg2Instrumentor + class CommenterOptions(TypedDict, total=False): + db_driver: bool + db_framework: bool + opentelemetry_values: bool + + class PsycopgInstrumentKwargs(TypedDict, total=False): + enable_commenter: bool + commenter_options: CommenterOptions + skip_dep_check: bool + + PACKAGE_NAMES = ('psycopg', 'psycopg2') -def instrument_psycopg(conn_or_module: Any = None, **kwargs: Any): +def instrument_psycopg(conn_or_module: Any = None, **kwargs: Unpack[PsycopgInstrumentKwargs]) -> None: """Instrument a `psycopg` connection or module so that spans are automatically created for each query. See the `Logfire.instrument_psycopg` method for details. @@ -50,7 +62,7 @@ def instrument_psycopg(conn_or_module: Any = None, **kwargs: Any): raise ValueError(f"Don't know how to instrument {conn_or_module!r}") -def _instrument_psycopg(name: str, conn: Any = None, **kwargs: Any): +def _instrument_psycopg(name: str, conn: Any = None, **kwargs: Unpack[PsycopgInstrumentKwargs]) -> None: try: instrumentor_module = importlib.import_module(f'opentelemetry.instrumentation.{name}') except ImportError: @@ -78,7 +90,7 @@ def _instrument_psycopg(name: str, conn: Any = None, **kwargs: Any): instrumentor.instrument_connection(conn) -def check_version(name: str, version: str, instrumentor: Instrumentor): +def check_version(name: str, version: str, instrumentor: Instrumentor) -> bool: with contextlib.suppress(Exception): # it's not worth raising an exception if this fails somehow. for dep in instrumentor.instrumentation_dependencies(): req = Requirement(dep) # dep is a string like 'psycopg >= 3.1.0' diff --git a/logfire/_internal/integrations/pymongo.py b/logfire/_internal/integrations/pymongo.py index 6480aeefb..ba61607b8 100644 --- a/logfire/_internal/integrations/pymongo.py +++ b/logfire/_internal/integrations/pymongo.py @@ -1,9 +1,35 @@ -from typing import Any +from __future__ import annotations + +from typing import TYPE_CHECKING, Any from opentelemetry.instrumentation.pymongo import PymongoInstrumentor +if TYPE_CHECKING: + from opentelemetry.metrics import MeterProvider + from opentelemetry.trace import TracerProvider + from pymongo.monitoring import CommandFailedEvent, CommandStartedEvent, CommandSucceededEvent + from typing_extensions import Protocol, TypedDict, Unpack + + class RequestHook(Protocol): + def __call__(self, span: Any, event: CommandStartedEvent) -> None: ... + + class ResponseHook(Protocol): + def __call__(self, span: Any, event: CommandSucceededEvent) -> None: ... + + class FailedHook(Protocol): + def __call__(self, span: Any, event: CommandFailedEvent) -> None: ... + + class PymongoInstrumentKwargs(TypedDict, total=False): + request_hook: RequestHook | None + response_hook: ResponseHook | None + failed_hook: FailedHook | None + capture_statement: bool | None + tracer_provider: TracerProvider | None + meter_provider: MeterProvider | None + skip_dep_check: bool + -def instrument_pymongo(**kwargs: Any): +def instrument_pymongo(**kwargs: Unpack[PymongoInstrumentKwargs]) -> None: """Instrument the `pymongo` module so that spans are automatically created for each operation. See the `Logfire.instrument_pymongo` method for details. diff --git a/logfire/_internal/integrations/redis.py b/logfire/_internal/integrations/redis.py index 479fa725f..4740e475c 100644 --- a/logfire/_internal/integrations/redis.py +++ b/logfire/_internal/integrations/redis.py @@ -1,9 +1,30 @@ -from typing import Any +from __future__ import annotations + +from typing import TYPE_CHECKING, Any from opentelemetry.instrumentation.redis import RedisInstrumentor +if TYPE_CHECKING: + from opentelemetry.metrics import MeterProvider + from opentelemetry.trace import Span, TracerProvider + from redis import Connection + from typing_extensions import Protocol, TypedDict, Unpack + + class RequestHook(Protocol): + def __call__(self, span: Span, instance: Connection, *args: Any, **kwargs: Any) -> None: ... + + class ResponseHook(Protocol): + def __call__(self, span: Span, instance: Connection, response: Any) -> None: ... + + class RedisInstrumentKwargs(TypedDict, total=False): + request_hook: RequestHook | None + response_hook: ResponseHook | None + tracer_provider: TracerProvider | None + meter_provider: MeterProvider | None + skip_dep_check: bool + -def instrument_redis(**kwargs: Any): +def instrument_redis(**kwargs: Unpack[RedisInstrumentKwargs]) -> None: """Instrument the `redis` module so that spans are automatically created for each operation. See the `Logfire.instrument_redis` method for details. diff --git a/logfire/_internal/integrations/sqlalchemy.py b/logfire/_internal/integrations/sqlalchemy.py index 5a3918188..81716bdd6 100644 --- a/logfire/_internal/integrations/sqlalchemy.py +++ b/logfire/_internal/integrations/sqlalchemy.py @@ -1,9 +1,26 @@ -from typing import Any +from __future__ import annotations + +from typing import TYPE_CHECKING from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor +if TYPE_CHECKING: + from sqlalchemy import Engine + from typing_extensions import TypedDict, Unpack + + class CommenterOptions(TypedDict, total=False): + db_driver: bool + db_framework: bool + 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: Any): +def instrument_sqlalchemy(**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. diff --git a/logfire/_internal/integrations/starlette.py b/logfire/_internal/integrations/starlette.py index 92e51700f..d9ed63e04 100644 --- a/logfire/_internal/integrations/starlette.py +++ b/logfire/_internal/integrations/starlette.py @@ -1,10 +1,33 @@ -from typing import Any +from __future__ import annotations + +from typing import TYPE_CHECKING, Any from opentelemetry.instrumentation.starlette import StarletteInstrumentor from starlette.applications import Starlette +if TYPE_CHECKING: + from opentelemetry.metrics import MeterProvider + from opentelemetry.trace import Span, TracerProvider + from typing_extensions import Protocol, TypedDict, Unpack + + class ServerRequestHook(Protocol): + def __call__(self, span: Span, scope: dict[str, Any]): ... + + class ClientRequestHook(Protocol): + def __call__(self, span: Span, scope: dict[str, Any]): ... + + class ClientResponseHook(Protocol): + def __call__(self, span: Span, message: dict[str, Any]): ... + + class StarletteInstrumentKwargs(TypedDict, total=False): + server_request_hook: ServerRequestHook | None + client_request_hook: ClientRequestHook | None + client_response_hook: ClientResponseHook | None + tracer_provider: TracerProvider | None + meter_provider: MeterProvider | None + -def instrument_starlette(app: Starlette, **kwargs: Any): +def instrument_starlette(app: Starlette, **kwargs: Unpack[StarletteInstrumentKwargs]): """Instrument `app` so that spans are automatically created for each request. See the `Logfire.instrument_starlette` method for details. diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index 0536798af..a5a1fbe5b 100644 --- a/logfire/_internal/main.py +++ b/logfire/_internal/main.py @@ -73,7 +73,14 @@ from starlette.applications import Starlette from starlette.requests import Request from starlette.websockets import WebSocket + from typing_extensions import Unpack + from .integrations.flask import FlaskInstrumentKwargs + from .integrations.psycopg import PsycopgInstrumentKwargs + from .integrations.pymongo import PymongoInstrumentKwargs + from .integrations.redis import RedisInstrumentKwargs + from .integrations.sqlalchemy import SQLAlchemyInstrumentKwargs + from .integrations.starlette import StarletteInstrumentKwargs try: from pydantic import ValidationError @@ -1103,7 +1110,7 @@ def instrument_requests(self, excluded_urls: str | None = None, **kwargs: Any): self._warn_if_not_initialized_for_instrumentation() return instrument_requests(excluded_urls=excluded_urls, **kwargs) - def instrument_psycopg(self, conn_or_module: Any = None, **kwargs: Any): + 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. Uses the OpenTelemetry instrumentation libraries for @@ -1127,7 +1134,7 @@ def instrument_psycopg(self, conn_or_module: Any = None, **kwargs: Any): self._warn_if_not_initialized_for_instrumentation() return instrument_psycopg(conn_or_module, **kwargs) - def instrument_flask(self, app: Flask, **kwargs: Any): + def instrument_flask(self, app: Flask, **kwargs: Unpack[FlaskInstrumentKwargs]) -> None: """Instrument `app` so that spans are automatically created for each request. Uses the @@ -1139,7 +1146,7 @@ def instrument_flask(self, app: Flask, **kwargs: Any): self._warn_if_not_initialized_for_instrumentation() return instrument_flask(app, **kwargs) - def instrument_starlette(self, app: Starlette, **kwargs: Any): + def instrument_starlette(self, app: Starlette, **kwargs: Unpack[StarletteInstrumentKwargs]) -> None: """Instrument `app` so that spans are automatically created for each request. Uses the @@ -1163,7 +1170,7 @@ def instrument_aiohttp_client(self, **kwargs: Any): self._warn_if_not_initialized_for_instrumentation() return instrument_aiohttp_client(**kwargs) - def instrument_sqlalchemy(self, **kwargs: Any): + def instrument_sqlalchemy(self, **kwargs: Unpack[SQLAlchemyInstrumentKwargs]) -> None: """Instrument the `sqlalchemy` module so that spans are automatically created for each query. Uses the @@ -1175,7 +1182,7 @@ def instrument_sqlalchemy(self, **kwargs: Any): self._warn_if_not_initialized_for_instrumentation() return instrument_sqlalchemy(**kwargs) - def instrument_pymongo(self, **kwargs: Any): + def instrument_pymongo(self, **kwargs: Unpack[PymongoInstrumentKwargs]) -> None: """Instrument the `pymongo` module so that spans are automatically created for each operation. Uses the @@ -1187,7 +1194,7 @@ def instrument_pymongo(self, **kwargs: Any): self._warn_if_not_initialized_for_instrumentation() return instrument_pymongo(**kwargs) - def instrument_redis(self, **kwargs: Any): + def instrument_redis(self, **kwargs: Unpack[RedisInstrumentKwargs]) -> None: """Instrument the `redis` module so that spans are automatically created for each operation. Uses the