From 5f9dd90ad14e20be9ff53ee1b56c889b3d3bfd5e Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Mon, 6 Jan 2025 17:28:31 +0100 Subject: [PATCH] Drop psycopg kwargs (#768) Co-authored-by: Alex Hall --- docs/integrations/databases/psycopg.md | 17 +++++++ docs/reference/api/logfire.md | 1 + logfire/_internal/integrations/psycopg.py | 59 ++++++++++++++--------- logfire/_internal/main.py | 32 ++++++++++-- logfire/integrations/psycopg.py | 23 +++++++++ tests/otel_integrations/test_psycopg.py | 2 +- 6 files changed, 106 insertions(+), 28 deletions(-) create mode 100644 logfire/integrations/psycopg.py diff --git a/docs/integrations/databases/psycopg.md b/docs/integrations/databases/psycopg.md index 1403bbb29..2bcfce1b4 100644 --- a/docs/integrations/databases/psycopg.md +++ b/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/ diff --git a/docs/reference/api/logfire.md b/docs/reference/api/logfire.md index 8b0491520..8ce55ed37 100644 --- a/docs/reference/api/logfire.md +++ b/docs/reference/api/logfire.md @@ -12,6 +12,7 @@ filters: - "!instrument_redis" - "!instrument_pymongo" + - "!instrument_psycopg" - "!^with_trace_sample_rate$" - "!^_[^_]" diff --git a/logfire/_internal/integrations/psycopg.py b/logfire/_internal/integrations/psycopg.py index 95a58c160..797624507 100644 --- a/logfire/_internal/integrations/psycopg.py +++ b/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')() diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index ab5e91743..21b15c6ad 100644 --- a/logfire/_internal/main.py +++ b/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, diff --git a/logfire/integrations/psycopg.py b/logfire/integrations/psycopg.py new file mode 100644 index 000000000..21ea1d406 --- /dev/null +++ b/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.""" diff --git a/tests/otel_integrations/test_psycopg.py b/tests/otel_integrations/test_psycopg.py index b7e639f12..fd4d23026 100644 --- a/tests/otel_integrations/test_psycopg.py +++ b/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():