Drop PyMongo kwargs (#758)

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... 🚗"

logfire/_internal/integrations/pymongo.py

@@ -1,40 +1,38 @@
 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'
         'You can install this with:\n'
         "    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,
+    )

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,

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."""