Add logfire.instrument_asyncpg() (#44)

Author: alexmojaki

docs/integrations/asyncpg.md

@@ -0,0 +1,76 @@
+# asyncpg
+
+The [`logfire.instrument_asyncpg()`][logfire.Logfire.instrument_asyncpg] function can be used to instrument the [asyncpg][asyncpg] PostgreSQL driver with **Logfire**.
+
+## Installation
+
+Install `logfire` with the `asyncpg` extra:
+
+{{ install_logfire(extras=['asyncpg']) }}
+
+## Usage
+
+Let's setup a PostgreSQL database using Docker and run a Python script that connects to the database using asyncpg to
+demonstrate how to use **Logfire** with asyncpg.
+
+### 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 asyncio
+
+import asyncpg
+
+import logfire
+
+logfire.configure()
+logfire.instrument_asyncpg()
+
+
+async def main():
+    connection: asyncpg.Connection = await asyncpg.connect(
+        user='user', password='secret', database='database', host='0.0.0.0', port=5432
+    )
+
+    with logfire.span('Create table and insert data'):
+        await connection.execute('CREATE TABLE IF NOT EXISTS test (id serial PRIMARY KEY, num integer, data varchar);')
+
+        # Insert some data
+        await connection.execute('INSERT INTO test (num, data) VALUES ($1, $2)', 100, 'abc')
+        await connection.execute('INSERT INTO test (num, data) VALUES ($1, $2)', 200, 'def')
+
+        # Query the data
+        for record in await connection.fetch('SELECT * FROM test'):
+            logfire.info('Retrieved {record=}', record=record)
+
+
+asyncio.run(main())
+```
+
+If you go to your project on the UI, you will see the span created by the script.
+
+[opentelemetry-asyncpg]: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/asyncpg/asyncpg.html
+[opentelemetry-asyncpg2]: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/asyncpg2/asyncpg2.html
+[asyncpg]: https://magicstack.github.io/asyncpg/

docs/integrations/index.md

@@ -22,6 +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               |
+| [Asyncpg](asyncpg.md)               | Databases               |
 | [Psycopg](psycopg.md)               | Databases               |
 | [PyMongo](pymongo.md)               | Databases               |
 | [Redis](redis.md)                   | Databases               |

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_asyncpg = DEFAULT_LOGFIRE_INSTANCE.instrument_asyncpg
 instrument_psycopg = DEFAULT_LOGFIRE_INSTANCE.instrument_psycopg
 shutdown = DEFAULT_LOGFIRE_INSTANCE.shutdown
 with_tags = DEFAULT_LOGFIRE_INSTANCE.with_tags

logfire/_internal/integrations/asyncpg.py

@@ -0,0 +1,9 @@
+from opentelemetry.instrumentation.asyncpg import AsyncPGInstrumentor
+
+
+def instrument_asyncpg():
+    """Instrument the `asyncpg` module so that spans are automatically created for each query.
+
+    See the `Logfire.instrument_asyncpg` method for details.
+    """
+    AsyncPGInstrumentor().instrument()

logfire/_internal/main.py

@@ -831,6 +831,12 @@ def instrument_openai(
 
         return instrument_openai(self, openai_client, suppress_other_instrumentation)
 
+    def instrument_asyncpg(self):
+        """Instrument the `asyncpg` module so that spans are automatically created for each query."""
+        from .integrations.asyncpg import instrument_asyncpg
+
+        return instrument_asyncpg()
+
     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.
 

mkdocs.yml

@@ -104,6 +104,7 @@ nav:
           - Requests: integrations/requests.md
           - AIOHTTP: integrations/aiohttp.md
           - SQLAlchemy: integrations/sqlalchemy.md
+          - Asyncpg: integrations/asyncpg.md
           - Psycopg: integrations/psycopg.md
           - PyMongo: integrations/pymongo.md
           - Redis: integrations/redis.md

pyproject.toml

@@ -59,6 +59,7 @@ flask = ["opentelemetry-instrumentation-flask >= 0.42b0"]
 httpx = ["opentelemetry-instrumentation-httpx >= 0.42b0"]
 starlette = ["opentelemetry-instrumentation-starlette >= 0.42b0"]
 sqlalchemy = ["opentelemetry-instrumentation-sqlalchemy >= 0.42b0"]
+asyncpg = ["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"]
@@ -104,6 +105,7 @@ dev-dependencies = [
     "opentelemetry-instrumentation-requests",
     "opentelemetry-instrumentation-sqlalchemy",
     "opentelemetry-instrumentation-system-metrics",
+    "opentelemetry-instrumentation-asyncpg",
     "opentelemetry-instrumentation-psycopg",
     "opentelemetry-instrumentation-psycopg2",
     "gitpython",
@@ -123,6 +125,7 @@ dev-dependencies = [
     "coverage[toml]>=7.5.0",
     "psycopg[binary]",
     "psycopg2-binary",
+    "asyncpg",
 ]
 
 [tool.rye.scripts]

requirements-dev.lock

@@ -20,6 +20,7 @@ asgiref==3.8.1
 asttokens==2.4.1
     # via devtools
     # via inline-snapshot
+asyncpg==0.29.0
 attrs==23.2.0
 babel==2.14.0
     # via mkdocs-material
@@ -144,6 +145,7 @@ opentelemetry-api==1.24.0
     # via opentelemetry-instrumentation
     # via opentelemetry-instrumentation-aiohttp-client
     # via opentelemetry-instrumentation-asgi
+    # via opentelemetry-instrumentation-asyncpg
     # via opentelemetry-instrumentation-dbapi
     # via opentelemetry-instrumentation-django
     # via opentelemetry-instrumentation-fastapi
@@ -165,6 +167,7 @@ opentelemetry-instrumentation==0.45b0
     # via logfire
     # via opentelemetry-instrumentation-aiohttp-client
     # via opentelemetry-instrumentation-asgi
+    # via opentelemetry-instrumentation-asyncpg
     # via opentelemetry-instrumentation-dbapi
     # via opentelemetry-instrumentation-django
     # via opentelemetry-instrumentation-fastapi
@@ -181,6 +184,7 @@ opentelemetry-instrumentation-aiohttp-client==0.45b0
 opentelemetry-instrumentation-asgi==0.45b0
     # via opentelemetry-instrumentation-fastapi
     # via opentelemetry-instrumentation-starlette
+opentelemetry-instrumentation-asyncpg==0.45b0
 opentelemetry-instrumentation-dbapi==0.45b0
     # via opentelemetry-instrumentation-psycopg
     # via opentelemetry-instrumentation-psycopg2
@@ -207,6 +211,7 @@ opentelemetry-sdk==1.24.0
 opentelemetry-semantic-conventions==0.45b0
     # via opentelemetry-instrumentation-aiohttp-client
     # via opentelemetry-instrumentation-asgi
+    # via opentelemetry-instrumentation-asyncpg
     # via opentelemetry-instrumentation-dbapi
     # via opentelemetry-instrumentation-django
     # via opentelemetry-instrumentation-fastapi

tests/otel_integrations/test_asyncpg.py

@@ -0,0 +1,12 @@
+import asyncpg
+from opentelemetry.instrumentation.asyncpg import AsyncPGInstrumentor
+
+import logfire
+
+
+def test_asyncpg():
+    original_execute = asyncpg.Connection.execute
+    logfire.instrument_asyncpg()
+    assert original_execute is not asyncpg.Connection.execute
+    AsyncPGInstrumentor().uninstrument()
+    assert original_execute is asyncpg.Connection.execute