Drop psycopg kwargs (#768)

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

docs/integrations/databases/psycopg.md

@@ -109,6 +109,23 @@ logfire.configure()
 logfire.instrument_psycopg(enable_commenter=True, commenter_options={'db_driver': False, 'dbapi_threadsafety': False})
 ```
 
+## API Reference
+
+::: logfire.Logfire.instrument_psycopg
+    options:
+        heading_level: 4
+        show_source: false
+        show_root_doc_entry: true
+        show_root_heading: true
+        show_root_full_path: false
+
+::: logfire.integrations.psycopg.CommenterOptions
+    options:
+        heading_level: 4
+        show_root_heading: true
+        show_source: false
+        filters: []
+
 [opentelemetry-psycopg]: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/psycopg/psycopg.html
 [opentelemetry-psycopg2]: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/psycopg2/psycopg2.html
 [psycopg]: https://www.psycopg.org/

docs/reference/api/logfire.md

@@ -12,6 +12,7 @@
         filters:
             - "!instrument_redis"
             - "!instrument_pymongo"
+            - "!instrument_psycopg"
             - "!^with_trace_sample_rate$"
             - "!^_[^_]"
 

logfire/_internal/integrations/psycopg.py

@@ -4,7 +4,7 @@
 import importlib
 from importlib.util import find_spec
 from types import ModuleType
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal, cast, overload
 
 from packaging.requirements import Requirement
 
@@ -13,25 +13,22 @@
 if TYPE_CHECKING:  # pragma: no cover
     from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor
     from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
-    from typing_extensions import TypedDict, Unpack
+    from psycopg import AsyncConnection, Connection
+    from psycopg2._psycopg import connection as Psycopg2Connection
+    from typing_extensions import TypeVar
 
-    Instrumentor = PsycopgInstrumentor | Psycopg2Instrumentor
-
-    class CommenterOptions(TypedDict, total=False):
-        db_driver: bool
-        db_framework: bool
-        opentelemetry_values: bool
+    PsycopgConnection = TypeVar('PsycopgConnection', Connection[Any], AsyncConnection[Any], Psycopg2Connection)
 
-    class PsycopgInstrumentKwargs(TypedDict, total=False):
-        enable_commenter: bool
-        commenter_options: CommenterOptions
+    Instrumentor = PsycopgInstrumentor | Psycopg2Instrumentor
 
 
-PACKAGE_NAMES = ('psycopg', 'psycopg2')
+PACKAGE_NAMES: tuple[Literal['psycopg'], Literal['psycopg2']] = ('psycopg', 'psycopg2')
 
 
 def instrument_psycopg(
-    logfire_instance: Logfire, conn_or_module: Any = None, **kwargs: Unpack[PsycopgInstrumentKwargs]
+    logfire_instance: Logfire,
+    conn_or_module: ModuleType | Literal['psycopg', 'psycopg2'] | None | PsycopgConnection | Psycopg2Connection,
+    **kwargs: Any,
 ) -> None:
     """Instrument a `psycopg` connection or module so that spans are automatically created for each query.
 
@@ -44,10 +41,11 @@ def instrument_psycopg(
                 instrument_psycopg(logfire_instance, package, **kwargs)
         return
     elif conn_or_module in PACKAGE_NAMES:
-        _instrument_psycopg(logfire_instance, conn_or_module, **kwargs)
+        _instrument_psycopg(logfire_instance, name=conn_or_module, **kwargs)
         return
     elif isinstance(conn_or_module, ModuleType):
-        instrument_psycopg(logfire_instance, conn_or_module.__name__, **kwargs)
+        module_name = cast(Literal['psycopg', 'psycopg2'], conn_or_module.__name__)
+        instrument_psycopg(logfire_instance, module_name, **kwargs)
         return
     else:
         # Given an object that's not a module or string,
@@ -66,14 +64,12 @@ def instrument_psycopg(
 
 
 def _instrument_psycopg(
-    logfire_instance: Logfire, name: str, conn: Any = None, **kwargs: Unpack[PsycopgInstrumentKwargs]
+    logfire_instance: Logfire,
+    name: Literal['psycopg', 'psycopg2'],
+    conn: Any = None,
+    **kwargs: Any,
 ) -> None:
-    try:
-        instrumentor_module = importlib.import_module(f'opentelemetry.instrumentation.{name}')
-    except ImportError:
-        raise ImportError(f"Run `pip install 'logfire[{name}]'` to install {name} instrumentation.")
-
-    instrumentor = getattr(instrumentor_module, f'{name.capitalize()}Instrumentor')()
+    instrumentor = _get_instrumentor(name)
     if conn is None:
         # OTEL looks at the installed packages to determine if the correct dependencies are installed.
         # This means that if a user installs `psycopg-binary` (which is commonly recommended)
@@ -99,7 +95,7 @@ def _instrument_psycopg(
         )
     else:
         # instrument_connection doesn't have a skip_dep_check argument.
-        instrumentor.instrument_connection(conn, tracer_provider=logfire_instance.config.get_tracer_provider())
+        instrumentor.instrument_connection(conn, tracer_provider=logfire_instance.config.get_tracer_provider())  # type: ignore[reportUnknownMemberType]
 
 
 def check_version(name: str, version: str, instrumentor: Instrumentor) -> bool:
@@ -110,3 +106,20 @@ def check_version(name: str, version: str, instrumentor: Instrumentor) -> bool:
             if req.name == name and req.specifier.contains(version.split()[0]):
                 return True
     return False
+
+
+@overload
+def _get_instrumentor(name: Literal['psycopg']) -> PsycopgInstrumentor: ...
+
+
+@overload
+def _get_instrumentor(name: Literal['psycopg2']) -> Psycopg2Instrumentor: ...
+
+
+def _get_instrumentor(name: str) -> Instrumentor:
+    try:
+        instrumentor_module = importlib.import_module(f'opentelemetry.instrumentation.{name}')
+    except ImportError:
+        raise ImportError(f"Run `pip install 'logfire[{name}]'` to install {name} instrumentation.")
+
+    return getattr(instrumentor_module, f'{name.capitalize()}Instrumentor')()

logfire/_internal/main.py

@@ -65,6 +65,7 @@
 from .utils import get_version, handle_internal_errors, log_internal_error, uniquify_sequence
 
 if TYPE_CHECKING:
+    from types import ModuleType
     from wsgiref.types import WSGIApplication
 
     import anthropic
@@ -95,13 +96,14 @@
         RequestHook as HttpxRequestHook,
         ResponseHook as HttpxResponseHook,
     )
+    from ..integrations.psycopg import CommenterOptions as PsycopgCommenterOptions
     from ..integrations.redis import RequestHook as RedisRequestHook, ResponseHook as RedisResponseHook
     from ..integrations.sqlalchemy import CommenterOptions as SQLAlchemyCommenterOptions
     from ..integrations.wsgi import RequestHook as WSGIRequestHook, ResponseHook as WSGIResponseHook
     from .integrations.asgi import ASGIApp, ASGIInstrumentKwargs
     from .integrations.aws_lambda import LambdaEvent, LambdaHandler
     from .integrations.mysql import MySQLConnection
-    from .integrations.psycopg import PsycopgInstrumentKwargs
+    from .integrations.psycopg import Psycopg2Connection, PsycopgConnection
     from .integrations.sqlite3 import SQLite3Connection
     from .integrations.system_metrics import Base as SystemMetricsBase, Config as SystemMetricsConfig
     from .utils import SysExcInfo
@@ -1374,7 +1376,25 @@ def instrument_requests(
             },
         )
 
-    def instrument_psycopg(self, conn_or_module: Any = None, **kwargs: Unpack[PsycopgInstrumentKwargs]) -> None:
+    @overload
+    def instrument_psycopg(self, conn_or_module: PsycopgConnection | Psycopg2Connection, **kwargs: Any) -> None: ...
+
+    @overload
+    def instrument_psycopg(
+        self,
+        conn_or_module: None | Literal['psycopg', 'psycopg2'] | ModuleType = None,
+        enable_commenter: bool = False,
+        commenter_options: PsycopgCommenterOptions | None = None,
+        **kwargs: Any,
+    ) -> None: ...
+
+    def instrument_psycopg(
+        self,
+        conn_or_module: Any = None,
+        enable_commenter: bool = False,
+        commenter_options: PsycopgCommenterOptions | None = None,
+        **kwargs: Any,
+    ) -> None:
         """Instrument a `psycopg` connection or module so that spans are automatically created for each query.
 
         Uses the OpenTelemetry instrumentation libraries for
@@ -1390,13 +1410,17 @@ def instrument_psycopg(self, conn_or_module: Any = None, **kwargs: Unpack[Psycop
                 - `None` (the default) to instrument whichever module(s) are installed.
                 - A `psycopg` or `psycopg2` connection.
 
+            enable_commenter: Adds comments to SQL queries performed by Psycopg, so that database logs have additional context.
+            commenter_options: Configure the tags to be added to the SQL comments.
             **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods,
-                particularly `enable_commenter` and `commenter_options`.
+                for future compatibility.
         """
         from .integrations.psycopg import instrument_psycopg
 
         self._warn_if_not_initialized_for_instrumentation()
-        return instrument_psycopg(self, conn_or_module, **kwargs)
+        if enable_commenter:
+            kwargs.update({'enable_commenter': True, 'commenter_options': commenter_options or {}})
+        return instrument_psycopg(self, conn_or_module=conn_or_module, **kwargs)
 
     def instrument_flask(
         self,

logfire/integrations/psycopg.py

@@ -0,0 +1,23 @@
+from typing import TypedDict
+
+
+class CommenterOptions(TypedDict, total=False):
+    """The `commenter_options` parameter for `instrument_psycopg`."""
+
+    db_driver: bool
+    """Include the database driver name in the comment e.g. 'psycopg2'."""
+
+    dbapi_threadsafety: bool
+    """Include the DB-API threadsafety value in the comment."""
+
+    dbapi_level: bool
+    """Include the DB-API level in the comment."""
+
+    libpq_version: bool
+    """Include the libpq version in the comment."""
+
+    driver_paramstyle: bool
+    """Include the driver paramstyle in the comment e.g. 'driver_paramstyle=pyformat'"""
+
+    opentelemetry_values: bool
+    """Enabling this flag will add traceparent values to the comment."""

tests/otel_integrations/test_psycopg.py

@@ -80,7 +80,7 @@ def test_instrument_psycopg_connection():
 
 def test_instrument_unknown():
     with pytest.raises(ValueError):
-        instrument_psycopg('unknown')
+        instrument_psycopg('unknown')  # type: ignore[reportCallIssue]
 
 
 def test_instrument_missing_otel_package():