logfire.instrument_psycopg() function (#30)

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              |

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/

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

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

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.
 

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

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]