Handle async engines in engines arg of instrument_sqlalchemy (#1425)

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

Author: dhruv-ahuja

logfire-api/logfire_api/_internal/integrations/sqlalchemy.pyi

@@ -1,9 +1,9 @@
 from logfire.integrations.sqlalchemy import CommenterOptions as CommenterOptions
 from sqlalchemy import Engine
 from sqlalchemy.ext.asyncio import AsyncEngine
-from typing import Any
+from typing import Any, Iterable
 
-def instrument_sqlalchemy(engine: AsyncEngine | Engine | None, enable_commenter: bool, commenter_options: CommenterOptions, **kwargs: Any) -> None:
+def instrument_sqlalchemy(engine: AsyncEngine | Engine | None, engines: Iterable[AsyncEngine | Engine] | None, enable_commenter: bool, commenter_options: CommenterOptions, **kwargs: Any) -> None:
     """Instrument the `sqlalchemy` module so that spans are automatically created for each query.
 
     See the `Logfire.instrument_sqlalchemy` method for details.

logfire-api/logfire_api/_internal/main.pyi

@@ -777,7 +777,7 @@ class Logfire:
         [OpenTelemetry aiohttp server Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/aiohttp_server/aiohttp_server.html)
         library, specifically `AioHttpServerInstrumentor().instrument()`, to which it passes `**kwargs`.
         """
-    def instrument_sqlalchemy(self, engine: AsyncEngine | Engine | None = None, enable_commenter: bool = False, commenter_options: SQLAlchemyCommenterOptions | None = None, **kwargs: Any) -> None:
+    def instrument_sqlalchemy(self, engine: AsyncEngine | Engine | None = None, engines: Iterable[AsyncEngine | Engine] | None = None, enable_commenter: bool = False, commenter_options: SQLAlchemyCommenterOptions | None = None, **kwargs: Any) -> None:
         """Instrument the `sqlalchemy` module so that spans are automatically created for each query.
 
         Uses the

logfire/_internal/integrations/sqlalchemy.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextlib
+from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any
 
 try:
@@ -19,8 +20,17 @@
     from sqlalchemy.ext.asyncio import AsyncEngine
 
 
+def _convert_to_sync_engine(engine: AsyncEngine | Engine | None) -> Engine | None:
+    from sqlalchemy.ext.asyncio import AsyncEngine
+
+    if isinstance(engine, AsyncEngine):
+        return engine.sync_engine
+    return engine
+
+
 def instrument_sqlalchemy(
     engine: AsyncEngine | Engine | None,
+    engines: Iterable[AsyncEngine | Engine] | None,
     enable_commenter: bool,
     commenter_options: CommenterOptions,
     **kwargs: Any,
@@ -30,10 +40,11 @@ def instrument_sqlalchemy(
     See the `Logfire.instrument_sqlalchemy` method for details.
     """
     with contextlib.suppress(ImportError):
-        from sqlalchemy.ext.asyncio import AsyncEngine
+        engine = _convert_to_sync_engine(engine)
+
+        if engines is not None:
+            engines = [_convert_to_sync_engine(engine_entry) for engine_entry in engines]  # type: ignore
 
-        if isinstance(engine, AsyncEngine):
-            engine = engine.sync_engine
     return SQLAlchemyInstrumentor().instrument(
-        engine=engine, enable_commenter=enable_commenter, commenter_options=commenter_options, **kwargs
+        engine=engine, engines=engines, enable_commenter=enable_commenter, commenter_options=commenter_options, **kwargs
     )

logfire/_internal/main.py

@@ -1816,6 +1816,7 @@ def instrument_aiohttp_server(self, **kwargs: Any) -> None:
     def instrument_sqlalchemy(
         self,
         engine: AsyncEngine | Engine | None = None,
+        engines: Iterable[AsyncEngine | Engine] | None = None,
         enable_commenter: bool = False,
         commenter_options: SQLAlchemyCommenterOptions | None = None,
         **kwargs: Any,
@@ -1827,7 +1828,8 @@ def instrument_sqlalchemy(
         library, specifically `SQLAlchemyInstrumentor().instrument()`, to which it passes `**kwargs`.
 
         Args:
-            engine: The `sqlalchemy` engine to instrument, or `None` to instrument all engines.
+            engine: The `sqlalchemy` engine to instrument.
+            engines: An iterable of `sqlalchemy` engines to instrument.
             enable_commenter: Adds comments to SQL queries performed by SQLAlchemy, 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.
@@ -1837,6 +1839,7 @@ def instrument_sqlalchemy(
         self._warn_if_not_initialized_for_instrumentation()
         return instrument_sqlalchemy(
             engine=engine,
+            engines=engines,
             enable_commenter=enable_commenter,
             commenter_options=commenter_options or {},
             **{

tests/otel_integrations/test_sqlalchemy.py

@@ -4,6 +4,7 @@
 from collections.abc import Iterator
 from contextlib import contextmanager
 from pathlib import Path
+from typing import Any
 from unittest import mock
 
 import pytest
@@ -210,9 +211,19 @@ def test_sqlalchemy_instrumentation(exporter: TestExporter):
     SQLAlchemyInstrumentor().uninstrument()
 
 
-def test_sqlalchemy_instrumentation_commenter(exporter: TestExporter):
+@pytest.mark.parametrize('parameter', ['engine', 'engines'])
+def test_sqlalchemy_instrumentation_commenter(parameter: str, exporter: TestExporter):
     with sqlite_engine(Path('example3.db')) as engine:
-        logfire.instrument_sqlalchemy(engine=engine, enable_commenter=True, commenter_options={'db_framework': False})
+        instrumentation_parameters: dict[str, Any] = {
+            'enable_commenter': True,
+            'commenter_options': {'db_framework': False},
+        }
+        if parameter == 'engine':
+            instrumentation_parameters['engine'] = engine
+        else:
+            instrumentation_parameters['engines'] = [engine]
+
+        logfire.instrument_sqlalchemy(**instrumentation_parameters)
         logfire.instrument_sqlite3()
 
         Base.metadata.create_all(engine)
@@ -377,9 +388,13 @@ def sqlite_async_engine(path: Path) -> Iterator[AsyncEngine]:
 
 
 @pytest.mark.anyio
-async def test_sqlalchemy_async_instrumentation(exporter: TestExporter):
+@pytest.mark.parametrize('parameter', ['engine', 'engines'])
+async def test_sqlalchemy_async_instrumentation(parameter: str, exporter: TestExporter):
     with sqlite_async_engine(Path('example2.db')) as engine:
-        logfire.instrument_sqlalchemy(engine=engine)
+        if parameter == 'engine':
+            logfire.instrument_sqlalchemy(engine=engine)
+        else:
+            logfire.instrument_sqlalchemy(engines=[engine])
 
         async with engine.begin() as conn:
             await conn.run_sync(Base.metadata.create_all)