From 27bfae27bcb48f901d4ce291002313219297a13b Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Tue, 31 Dec 2024 13:17:39 +0100 Subject: [PATCH] Drop PyMongo kwargs (#758) --- docs/integrations/databases/pymongo.md | 2 - logfire/_internal/integrations/pymongo.py | 46 +++++++++++------------ logfire/_internal/main.py | 28 ++++++++++++-- logfire/integrations/pymongo.py | 13 +++++++ 4 files changed, 60 insertions(+), 29 deletions(-) create mode 100644 logfire/integrations/pymongo.py diff --git a/docs/integrations/databases/pymongo.md b/docs/integrations/databases/pymongo.md index 35ee0458b..cb557308b 100644 --- a/docs/integrations/databases/pymongo.md +++ b/docs/integrations/databases/pymongo.md @@ -2,8 +2,6 @@ integration: otel --- -# PyMongo - The [`logfire.instrument_pymongo()`][logfire.Logfire.instrument_pymongo] method will create a span for every operation performed using your [PyMongo][pymongo] clients. !!! success "Also works with Motor... 🚗" diff --git a/logfire/_internal/integrations/pymongo.py b/logfire/_internal/integrations/pymongo.py index bbbc81c8c..b726e9c01 100644 --- a/logfire/_internal/integrations/pymongo.py +++ b/logfire/_internal/integrations/pymongo.py @@ -1,9 +1,14 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import Any try: - from opentelemetry.instrumentation.pymongo import PymongoInstrumentor + from opentelemetry.instrumentation.pymongo import ( + PymongoInstrumentor, + dummy_callback, # type: ignore[reportUnknownVariableType] + ) + + from logfire.integrations.pymongo import FailedHook, RequestHook, ResponseHook except ImportError: raise RuntimeError( '`logfire.instrument_pymongo()` requires the `opentelemetry-instrumentation-pymongo` package.\n' @@ -11,30 +16,23 @@ " pip install 'logfire[pymongo]'" ) -if TYPE_CHECKING: - 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 - skip_dep_check: bool - - -def instrument_pymongo(**kwargs: Unpack[PymongoInstrumentKwargs]) -> None: +def instrument_pymongo( + *, + capture_statement: bool, + request_hook: RequestHook | None, + response_hook: ResponseHook | None, + failed_hook: FailedHook | None, + **kwargs: Any, +) -> None: """Instrument the `pymongo` module so that spans are automatically created for each operation. See the `Logfire.instrument_pymongo` method for details. """ - PymongoInstrumentor().instrument(**kwargs) + PymongoInstrumentor().instrument( + request_hook=request_hook or dummy_callback, + response_hook=response_hook or dummy_callback, + failed_hook=failed_hook or dummy_callback, + capture_statement=capture_statement, + **kwargs, + ) diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index f3ba8d0a5..b5aecd6d4 100644 --- a/logfire/_internal/main.py +++ b/logfire/_internal/main.py @@ -94,6 +94,11 @@ RequestHook as HttpxRequestHook, ResponseHook as HttpxResponseHook, ) + from ..integrations.pymongo import ( + FailedHook as PymongoFailedHook, + RequestHook as PymongoRequestHook, + ResponseHook as PymongoResponseHook, + ) from ..integrations.sqlalchemy import CommenterOptions as SQLAlchemyCommenterOptions from ..integrations.wsgi import ( RequestHook as WSGIRequestHook, @@ -103,7 +108,6 @@ from .integrations.aws_lambda import LambdaEvent, LambdaHandler from .integrations.mysql import MySQLConnection from .integrations.psycopg import PsycopgInstrumentKwargs - from .integrations.pymongo import PymongoInstrumentKwargs from .integrations.redis import RedisInstrumentKwargs from .integrations.sqlite3 import SQLite3Connection from .integrations.system_metrics import Base as SystemMetricsBase, Config as SystemMetricsConfig @@ -1697,18 +1701,36 @@ def instrument_aws_lambda( }, ) - def instrument_pymongo(self, **kwargs: Unpack[PymongoInstrumentKwargs]) -> None: + def instrument_pymongo( + self, + capture_statement: bool = False, + request_hook: PymongoRequestHook | None = None, + response_hook: PymongoResponseHook | None = None, + failed_hook: PymongoFailedHook | None = None, + **kwargs: Any, + ) -> None: """Instrument the `pymongo` module so that spans are automatically created for each operation. Uses the [OpenTelemetry pymongo Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/pymongo/pymongo.html) library, specifically `PymongoInstrumentor().instrument()`, to which it passes `**kwargs`. + + Args: + capture_statement: Set to `True` to capture the statement in the span attributes. + request_hook: A function called when a command is sent to the server. + response_hook: A function that is called when a command is successfully completed. + failed_hook: A function that is called when a command fails. + **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods for future compatibility. """ from .integrations.pymongo import instrument_pymongo self._warn_if_not_initialized_for_instrumentation() return instrument_pymongo( - **{ # type: ignore + capture_statement=capture_statement, + request_hook=request_hook, + response_hook=response_hook, + failed_hook=failed_hook, + **{ 'tracer_provider': self._config.get_tracer_provider(), 'meter_provider': self._config.get_meter_provider(), **kwargs, diff --git a/logfire/integrations/pymongo.py b/logfire/integrations/pymongo.py new file mode 100644 index 000000000..0b6a6acdf --- /dev/null +++ b/logfire/integrations/pymongo.py @@ -0,0 +1,13 @@ +from typing import Callable + +from opentelemetry.trace import Span +from pymongo.monitoring import CommandFailedEvent, CommandStartedEvent, CommandSucceededEvent + +RequestHook = Callable[[Span, CommandStartedEvent], None] +"""A hook that is called when a command is started.""" + +ResponseHook = Callable[[Span, CommandSucceededEvent], None] +"""A hook that is called when a command is succeeded.""" + +FailedHook = Callable[[Span, CommandFailedEvent], None] +"""A hook that is called when a command is failed."""