From 47e6bb2f673ccbd6f0ed4e1a2875b5a5453a1d53 Mon Sep 17 00:00:00 2001 From: Alex Hall Date: Mon, 29 Apr 2024 16:37:50 +0200 Subject: [PATCH] `logfire.instrument_psycopg()` function (#30) --- docs/integrations/index.md | 2 +- docs/integrations/psycopg.md | 114 ++++++++++++++++++++++ docs/integrations/psycopg2.md | 80 --------------- logfire/__init__.py | 1 + logfire/_internal/integrations/psycopg.py | 88 +++++++++++++++++ logfire/_internal/main.py | 23 +++++ mkdocs.yml | 2 +- pyproject.toml | 7 +- requirements-dev.lock | 20 ++++ tests/otel_integrations/test_psycopg.py | 111 +++++++++++++++++++++ 10 files changed, 365 insertions(+), 83 deletions(-) create mode 100644 docs/integrations/psycopg.md delete mode 100644 docs/integrations/psycopg2.md create mode 100644 logfire/_internal/integrations/psycopg.py create mode 100644 tests/otel_integrations/test_psycopg.py diff --git a/docs/integrations/index.md b/docs/integrations/index.md index 5bbe64a82..7ab864751 100644 --- a/docs/integrations/index.md +++ b/docs/integrations/index.md @@ -22,7 +22,7 @@ Below you can see more details on how to use Logfire with some of the most popul | [Requests](requests.md) | HTTP Client | | [AIOHTTP](aiohttp.md) | HTTP Client | | [SQLAlchemy](sqlalchemy.md) | Databases | -| [Psycopg2](psycopg2.md) | Databases | +| [Psycopg](psycopg.md) | Databases | | [PyMongo](pymongo.md) | Databases | | [Redis](redis.md) | Databases | | [Celery](celery.md) | Task Queue | diff --git a/docs/integrations/psycopg.md b/docs/integrations/psycopg.md new file mode 100644 index 000000000..947628ef5 --- /dev/null +++ b/docs/integrations/psycopg.md @@ -0,0 +1,114 @@ +# Psycopg + +The [`logfire.instrument_psycopg()`][logfire.Logfire.instrument_psycopg] function can be used to instrument the [Psycopg][psycopg] PostgreSQL driver with **Logfire**. It works with both the `psycopg2` and `psycopg` (i.e. Psycopg 3) packages. + +See the documentation for the [OpenTelemetry Psycopg Instrumentation][opentelemetry-psycopg] or the [OpenTelemetry Psycopg2 Instrumentation][opentelemetry-psycopg2] package for more details. + +## Installation + +Install `logfire` with the `psycopg` extra: + +{{ install_logfire(extras=['psycopg']) }} + +Or with the `psycopg2` extra: + +{{ install_logfire(extras=['psycopg2']) }} + +## Usage + +Let's setup a PostgreSQL database using Docker and run a Python script that connects to the database using Psycopg to +demonstrate how to use **Logfire** with Psycopg. + +### Setup a PostgreSQL Database Using Docker + +First, we need to initialize a PostgreSQL database. This can be easily done using Docker with the following command: + +```bash +docker run --name postgres \ + -e POSTGRES_USER=user \ + -e POSTGRES_PASSWORD=secret \ + -e POSTGRES_DB=database \ + -p 5432:5432 -d postgres +``` + +This command accomplishes the following: + +- `--name postgres`: This defines the name of the Docker container. +- `-e POSTGRES_USER=user`: This sets a user for the PostgreSQL server. +- `-e POSTGRES_PASSWORD=secret`: This sets a password for the PostgreSQL server. +- `-e POSTGRES_DB=database`: This creates a new database named "database", the same as the one used in your Python script. +- `-p 5432:5432`: This makes the PostgreSQL instance available on your local machine under port 5432. +- `-d postgres`: This denotes the Docker image to be used, in this case, "postgres". + +### Run the Python script + +The following Python script connects to the PostgreSQL database and executes some SQL queries: + +```py +import logfire +import psycopg + +logfire.configure() + +# To instrument the whole module: +logfire.instrument_psycopg(psycopg) +# or +logfire.instrument_psycopg('psycopg') +# or just instrument whichever modules (psycopg and/or psycopg2) are installed: +logfire.instrument_psycopg() + +connection = psycopg.connect( + 'dbname=database user=user password=secret host=0.0.0.0 port=5432' +) + +# Or instrument just the connection: +logfire.instrument_psycopg(connection) + +with logfire.span('Create table and insert data'), connection.cursor() as cursor: + cursor.execute( + 'CREATE TABLE IF NOT EXISTS test (id serial PRIMARY KEY, num integer, data varchar);' + ) + + # Insert some data + cursor.execute('INSERT INTO test (num, data) VALUES (%s, %s)', (100, 'abc')) + cursor.execute('INSERT INTO test (num, data) VALUES (%s, %s)', (200, 'def')) + + # Query the data + cursor.execute('SELECT * FROM test') +``` + +If you go to your project on the UI, you will see the span created by the script. + +## SQL Commenter + +To add SQL comments to the end of your queries to enrich your database logs with additional context, use the `enable_commenter` parameter: + +```python +import logfire + +logfire.instrument_psycopg(enable_commenter=True) +``` + +This can only be used when instrumenting the whole module, not individual connections. + +By default the SQL comments will include values for the following keys: + +- `db_driver` +- `dbapi_threadsafety` +- `dbapi_level` +- `libpq_version` +- `driver_paramstyle` +- `opentelemetry_values` + +You can exclude any of these keys by passing a dictionary with those keys and the value `False` to `commenter_options`, +e.g: + +```python +import logfire + +logfire.instrument_psycopg(enable_commenter=True, commenter_options={'db_driver': False, 'dbapi_threadsafety': False}) +``` + +[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/integrations/psycopg2.md b/docs/integrations/psycopg2.md deleted file mode 100644 index fbc630b20..000000000 --- a/docs/integrations/psycopg2.md +++ /dev/null @@ -1,80 +0,0 @@ -# Psycopg2 - -The [OpenTelemetry Instrumentation Psycopg2][opentelemetry-psycopg2] package can be used to instrument [Psycopg2][psycopg2]. - -## Installation - -Install `logfire` with the `psycopg2` extra: - -{{ install_logfire(extras=['psycopg2']) }} - -## Usage - -Let's setup a PostgreSQL database using Docker and run a Python script that connects to the database using Psycopg2 to -demonstrate how to use **Logfire** with Psycopg2. - -### Setup a PostgreSQL Database Using Docker - -First, we need to initialize a PostgreSQL database. This can be easily done using Docker with the following command: - -```bash -docker run --name postgres \ - -e POSTGRES_USER=user \ - -e POSTGRES_PASSWORD=secret \ - -e POSTGRES_DB=database \ - -p 5432:5432 -d postgres -``` - -This command accomplishes the following: - -- `--name postgres`: This defines the name of the Docker container. -- `-e POSTGRES_USER=user`: This sets a user for the PostgreSQL server. -- `-e POSTGRES_PASSWORD=secret`: This sets a password for the PostgreSQL server. -- `-e POSTGRES_DB=database`: This creates a new database named "database", the same as the one used in your Python script. -- `-p 5432:5432`: This makes the PostgreSQL instance available on your local machine under port 5432. -- `-d postgres`: This denotes the Docker image to be used, in this case, "postgres". - -### Run the Python script - -The following Python script connects to the PostgreSQL database and executes some SQL queries: - -```py -import logfire -import psycopg2 -from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor - -logfire.configure() -Psycopg2Instrumentor().instrument() - -conn = psycopg2.connect(database='database', user='user', password='secret', host='0.0.0.0', port='5433') - -with logfire.span('Create table and insert data'), conn.cursor() as cursor: - cursor.execute('CREATE TABLE IF NOT EXISTS test (id serial PRIMARY KEY, num integer, data varchar);') - - # Insert some data - cursor.execute('INSERT INTO test (num, data) VALUES (%s, %s)', (100, 'abc')) - cursor.execute('INSERT INTO test (num, data) VALUES (%s, %s)', (200, 'def')) - - # Query the data - cursor.execute('SELECT * FROM test') -``` - -If you go to your project on the UI, you will see the span created by the script. - -Feel free to read more about the Psycopg2 OpenTelemetry package on the official [documentation][opentelemetry-psycopg2]. - -!!! bug - A bug occurs when `opentelemetry-instrumentation-psycopg2` is used with `psycopg2-binary` instead of `psycopg2`. - More details on the issue can be found [here][psycopg2-binary-issue]. - - A workaround is to include `skip_dep_check` in `instrument` method: - - ```py - from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor - - Psycopg2Instrumentor().instrument(skip_dep_check=True) - ``` - -[opentelemetry-psycopg2]: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/psycopg2/psycopg2.html -[psycopg2]: https://www.psycopg.org/ -[psycopg2-binary-issue]: https://github.com/open-telemetry/opentelemetry-python-contrib/issues/610 diff --git a/logfire/__init__.py b/logfire/__init__.py index 14b6bcf2c..93c67c179 100644 --- a/logfire/__init__.py +++ b/logfire/__init__.py @@ -21,6 +21,7 @@ install_auto_tracing = DEFAULT_LOGFIRE_INSTANCE.install_auto_tracing instrument_fastapi = DEFAULT_LOGFIRE_INSTANCE.instrument_fastapi instrument_openai = DEFAULT_LOGFIRE_INSTANCE.instrument_openai +instrument_psycopg = DEFAULT_LOGFIRE_INSTANCE.instrument_psycopg shutdown = DEFAULT_LOGFIRE_INSTANCE.shutdown with_tags = DEFAULT_LOGFIRE_INSTANCE.with_tags # with_trace_sample_rate = DEFAULT_LOGFIRE_INSTANCE.with_trace_sample_rate diff --git a/logfire/_internal/integrations/psycopg.py b/logfire/_internal/integrations/psycopg.py new file mode 100644 index 000000000..ea1336581 --- /dev/null +++ b/logfire/_internal/integrations/psycopg.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +import contextlib +import importlib +from importlib.util import find_spec +from types import ModuleType +from typing import TYPE_CHECKING, Any + +from packaging.requirements import Requirement + +if TYPE_CHECKING: # pragma: no cover + from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor + from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor + + Instrumentor = PsycopgInstrumentor | Psycopg2Instrumentor + +PACKAGE_NAMES = ('psycopg', 'psycopg2') + + +def instrument_psycopg(conn_or_module: Any = None, **kwargs: Any): + """Instrument a `psycopg` connection or module so that spans are automatically created for each query. + + See the `Logfire.instrument_psycopg` method for details. + """ + if conn_or_module is None: + # By default, instrument whichever libraries are installed. + for package in PACKAGE_NAMES: + if find_spec(package): # pragma: no branch + instrument_psycopg(package, **kwargs) + return + elif conn_or_module in PACKAGE_NAMES: + _instrument_psycopg(conn_or_module, **kwargs) + return + elif isinstance(conn_or_module, ModuleType): + instrument_psycopg(conn_or_module.__name__, **kwargs) + return + else: + # Given an object that's not a module or string, + # and whose class (or an ancestor) is defined in one of the packages, assume it's a connection object. + for cls in conn_or_module.__class__.__mro__: + package = cls.__module__.split('.')[0] + if package in PACKAGE_NAMES: + if kwargs: + raise TypeError( + f'Extra keyword arguments are only supported when instrumenting the {package} module, not a connection.' + ) + _instrument_psycopg(package, conn_or_module, **kwargs) + return + + raise ValueError(f"Don't know how to instrument {conn_or_module!r}") + + +def _instrument_psycopg(name: str, conn: Any = None, **kwargs: Any): + 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')() + 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) + # then they can `import psycopg` but OTEL doesn't recognise this correctly. + # So we implement an alternative strategy, which is to import `psycopg(2)` and check `__version__`. + # If it matches, we can tell OTEL to skip the check so that it still works and doesn't log a warning. + mod = importlib.import_module(name) + skip_dep_check = check_version(name, mod.__version__, instrumentor) + + if kwargs.get('enable_commenter') and name == 'psycopg': + import psycopg.pq + + # OTEL looks for __libpq_version__ which only exists in psycopg2. + mod.__libpq_version__ = psycopg.pq.version() + + instrumentor.instrument(skip_dep_check=skip_dep_check, **kwargs) + else: + # instrument_connection doesn't have a skip_dep_check argument. + instrumentor.instrument_connection(conn) + + +def check_version(name: str, version: str, instrumentor: Instrumentor): + 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' + # The module __version__ can be something like '2.9.9 (dt dec pq3 ext lo64)', hence the split. + if req.name == name and req.specifier.contains(version.split()[0]): + return True + return False diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index 389db5cb7..44adc30b6 100644 --- a/logfire/_internal/main.py +++ b/logfire/_internal/main.py @@ -831,6 +831,29 @@ def instrument_openai( return instrument_openai(self, openai_client, suppress_other_instrumentation) + def instrument_psycopg(self, conn_or_module: Any = None, **kwargs: Any): + """Instrument a `psycopg` connection or module so that spans are automatically created for each query. + + Uses the OpenTelemetry instrumentation libraries for + [`psycopg`](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/psycopg/psycopg.html) + and + [`psycopg2`](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/psycopg2/psycopg2.html). + + Args: + conn_or_module: Can be: + + - The `psycopg` (version 3) or `psycopg2` module. + - The string `'psycopg'` or `'psycopg2'` to instrument the module. + - `None` (the default) to instrument whichever module(s) are installed. + - A `psycopg` or `psycopg2` connection. + + **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods, + particularly `enable_commenter` and `commenter_options`. + """ + from .integrations.psycopg import instrument_psycopg + + return instrument_psycopg(conn_or_module, **kwargs) + def metric_counter(self, name: str, *, unit: str = '', description: str = '') -> Counter: """Create a counter metric. diff --git a/mkdocs.yml b/mkdocs.yml index d5c252d0c..ed124f638 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -104,7 +104,7 @@ nav: - Requests: integrations/requests.md - AIOHTTP: integrations/aiohttp.md - SQLAlchemy: integrations/sqlalchemy.md - - Psycopg2: integrations/psycopg2.md + - Psycopg: integrations/psycopg.md - PyMongo: integrations/pymongo.md - Redis: integrations/redis.md - Celery: integrations/celery.md diff --git a/pyproject.toml b/pyproject.toml index 105936bc6..ae7e11075 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,7 +59,8 @@ flask = ["opentelemetry-instrumentation-flask >= 0.42b0"] httpx = ["opentelemetry-instrumentation-httpx >= 0.42b0"] starlette = ["opentelemetry-instrumentation-starlette >= 0.42b0"] sqlalchemy = ["opentelemetry-instrumentation-sqlalchemy >= 0.42b0"] -psycopg2 = ["opentelemetry-instrumentation-psycopg2 >= 0.42b0"] +psycopg = ["opentelemetry-instrumentation-psycopg2 >= 0.42b0", "packaging"] +psycopg2 = ["opentelemetry-instrumentation-psycopg2 >= 0.42b0", "packaging"] pymongo = ["opentelemetry-instrumentation-pymongo >= 0.42b0"] redis = ["opentelemetry-instrumentation-redis >= 0.42b0"] requests = ["opentelemetry-instrumentation-requests >= 0.42b0"] @@ -103,6 +104,8 @@ dev-dependencies = [ "opentelemetry-instrumentation-requests", "opentelemetry-instrumentation-sqlalchemy", "opentelemetry-instrumentation-system-metrics", + "opentelemetry-instrumentation-psycopg", + "opentelemetry-instrumentation-psycopg2", "gitpython", "devtools", "eval-type-backport", @@ -118,6 +121,8 @@ dev-dependencies = [ "mkdocs-glightbox>=0.3.7", "mkdocstrings-python>=1.8.0", "coverage[toml]>=7.5.0", + "psycopg[binary]", + "psycopg2-binary", ] [tool.rye.scripts] diff --git a/requirements-dev.lock b/requirements-dev.lock index f31a5463d..d8241c205 100644 --- a/requirements-dev.lock +++ b/requirements-dev.lock @@ -70,6 +70,8 @@ gitdb==4.0.11 gitpython==3.1.43 googleapis-common-protos==1.63.0 # via opentelemetry-exporter-otlp-proto-http +greenlet==3.0.3 + # via sqlalchemy griffe==0.44.0 # via mkdocstrings-python h11==0.14.0 @@ -142,10 +144,13 @@ opentelemetry-api==1.24.0 # via opentelemetry-instrumentation # via opentelemetry-instrumentation-aiohttp-client # via opentelemetry-instrumentation-asgi + # via opentelemetry-instrumentation-dbapi # via opentelemetry-instrumentation-django # via opentelemetry-instrumentation-fastapi # via opentelemetry-instrumentation-flask # via opentelemetry-instrumentation-httpx + # via opentelemetry-instrumentation-psycopg + # via opentelemetry-instrumentation-psycopg2 # via opentelemetry-instrumentation-requests # via opentelemetry-instrumentation-sqlalchemy # via opentelemetry-instrumentation-starlette @@ -160,10 +165,13 @@ opentelemetry-instrumentation==0.45b0 # via logfire # via opentelemetry-instrumentation-aiohttp-client # via opentelemetry-instrumentation-asgi + # via opentelemetry-instrumentation-dbapi # via opentelemetry-instrumentation-django # via opentelemetry-instrumentation-fastapi # via opentelemetry-instrumentation-flask # via opentelemetry-instrumentation-httpx + # via opentelemetry-instrumentation-psycopg + # via opentelemetry-instrumentation-psycopg2 # via opentelemetry-instrumentation-requests # via opentelemetry-instrumentation-sqlalchemy # via opentelemetry-instrumentation-starlette @@ -173,10 +181,15 @@ opentelemetry-instrumentation-aiohttp-client==0.45b0 opentelemetry-instrumentation-asgi==0.45b0 # via opentelemetry-instrumentation-fastapi # via opentelemetry-instrumentation-starlette +opentelemetry-instrumentation-dbapi==0.45b0 + # via opentelemetry-instrumentation-psycopg + # via opentelemetry-instrumentation-psycopg2 opentelemetry-instrumentation-django==0.45b0 opentelemetry-instrumentation-fastapi==0.45b0 opentelemetry-instrumentation-flask==0.45b0 opentelemetry-instrumentation-httpx==0.45b0 +opentelemetry-instrumentation-psycopg==0.45b0 +opentelemetry-instrumentation-psycopg2==0.45b0 opentelemetry-instrumentation-requests==0.45b0 opentelemetry-instrumentation-sqlalchemy==0.45b0 opentelemetry-instrumentation-starlette==0.45b0 @@ -194,6 +207,7 @@ opentelemetry-sdk==1.24.0 opentelemetry-semantic-conventions==0.45b0 # via opentelemetry-instrumentation-aiohttp-client # via opentelemetry-instrumentation-asgi + # via opentelemetry-instrumentation-dbapi # via opentelemetry-instrumentation-django # via opentelemetry-instrumentation-fastapi # via opentelemetry-instrumentation-flask @@ -239,6 +253,10 @@ protobuf==4.25.3 # via opentelemetry-proto psutil==5.9.8 # via opentelemetry-instrumentation-system-metrics +psycopg==3.1.18 +psycopg-binary==3.1.18 + # via psycopg +psycopg2-binary==2.9.9 pydantic==2.7.1 # via fastapi # via openai @@ -310,6 +328,7 @@ typing-extensions==4.11.0 # via logfire # via openai # via opentelemetry-sdk + # via psycopg # via pydantic # via pydantic-core # via sqlalchemy @@ -327,6 +346,7 @@ wrapt==1.16.0 # via deprecated # via opentelemetry-instrumentation # via opentelemetry-instrumentation-aiohttp-client + # via opentelemetry-instrumentation-dbapi # via opentelemetry-instrumentation-sqlalchemy zipp==3.18.1 # via importlib-metadata diff --git a/tests/otel_integrations/test_psycopg.py b/tests/otel_integrations/test_psycopg.py new file mode 100644 index 000000000..9a3b5e2cb --- /dev/null +++ b/tests/otel_integrations/test_psycopg.py @@ -0,0 +1,111 @@ +import sys +from unittest import mock + +import psycopg +import psycopg.pq +import psycopg2 +import pytest +from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor +from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor + +from logfire import instrument_psycopg +from logfire._internal.integrations.psycopg import check_version + + +def test_check_version(): + assert check_version('psycopg2', '2.7.3.1', Psycopg2Instrumentor()) + assert not check_version('psycopg2', '2.7.3.0', Psycopg2Instrumentor()) + assert check_version('psycopg', '3.1.0', PsycopgInstrumentor()) + assert not check_version('psycopg', '3.0.1', PsycopgInstrumentor()) + + assert check_version(psycopg.__name__, psycopg.__version__, PsycopgInstrumentor()) + assert check_version(psycopg2.__name__, psycopg2.__version__, Psycopg2Instrumentor()) # type: ignore + + +def test_instrument_psycopg(): + original_connect = psycopg.connect # type: ignore + + instrument_psycopg(psycopg) + assert original_connect is not psycopg.connect # type: ignore + PsycopgInstrumentor().uninstrument() # type: ignore + assert original_connect is psycopg.connect # type: ignore + + instrument_psycopg('psycopg') + assert original_connect is not psycopg.connect # type: ignore + PsycopgInstrumentor().uninstrument() # type: ignore + assert original_connect is psycopg.connect # type: ignore + + +def test_instrument_psycopg2(): + original_connect = psycopg2.connect + + instrument_psycopg(psycopg2) + assert original_connect is not psycopg2.connect + Psycopg2Instrumentor().uninstrument() # type: ignore + assert original_connect is psycopg2.connect + + instrument_psycopg('psycopg2') + assert original_connect is not psycopg2.connect + Psycopg2Instrumentor().uninstrument() # type: ignore + assert original_connect is psycopg2.connect + + +def test_instrument_both(): + original_connect = psycopg.connect # type: ignore + original_connect2 = psycopg2.connect + + instrument_psycopg() + assert original_connect is not psycopg.connect # type: ignore + assert original_connect2 is not psycopg2.connect + PsycopgInstrumentor().uninstrument() # type: ignore + Psycopg2Instrumentor().uninstrument() # type: ignore + assert original_connect is psycopg.connect # type: ignore + assert original_connect2 is psycopg2.connect + + +def test_instrument_psycopg_connection(): + pgconn = mock.Mock() + conn = psycopg.Connection(pgconn) + original_cursor_factory = conn.cursor_factory + + instrument_psycopg(conn) + assert original_cursor_factory is not conn.cursor_factory + assert conn._is_instrumented_by_opentelemetry # type: ignore + PsycopgInstrumentor().uninstrument_connection(conn) # type: ignore + assert original_cursor_factory is conn.cursor_factory + + pgconn.status = psycopg.pq.ConnStatus.BAD + conn.close() + + +def test_instrument_unknown(): + with pytest.raises(ValueError): + instrument_psycopg('unknown') + + +def test_instrument_missing_otel_package(): + sys.modules['opentelemetry.instrumentation.psycopg'] = None # type: ignore + with pytest.raises( + ImportError, match=r"Run `pip install 'logfire\[psycopg\]'` to install psycopg instrumentation." + ): + instrument_psycopg(psycopg) + del sys.modules['opentelemetry.instrumentation.psycopg'] + + +def test_instrument_connection_kwargs(): + pgconn = mock.Mock() + conn = psycopg.Connection(pgconn) + + with pytest.raises( + TypeError, + match=r'Extra keyword arguments are only supported when instrumenting the psycopg module, not a connection.', + ): + instrument_psycopg(conn, enable_commenter=True) + + pgconn.status = psycopg.pq.ConnStatus.BAD + conn.close() + + +def test_sql_commenter(): + instrument_psycopg(psycopg, enable_commenter=True) + assert psycopg.__libpq_version__ >= 110000 # type: ignore