Drop SQLAlchemy kwargs (#757)

pydantic · Dec 31, 2024 · 18fe86c · 18fe86c
1 parent f41a427
commit 18fe86c
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 20 deletions.
diff --git a/docs/integrations/databases/sqlalchemy.md b/docs/integrations/databases/sqlalchemy.md
@@ -2,8 +2,6 @@
 integration: otel
 ---
 
-# SQLAlchemy
-
 The [`logfire.instrument_sqlalchemy()`][logfire.Logfire.instrument_sqlalchemy] method will create a span for every query executed by a [SQLAlchemy][sqlalchemy] engine.
 
 ## Installation

diff --git a/logfire/_internal/integrations/sqlalchemy.py b/logfire/_internal/integrations/sqlalchemy.py
@@ -1,10 +1,12 @@
 from __future__ import annotations
 
 import contextlib
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 try:
     from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+
+    from logfire.integrations.sqlalchemy import CommenterOptions
 except ImportError:
     raise RuntimeError(
         '`logfire.instrument_sqlalchemy()` requires the `opentelemetry-instrumentation-sqlalchemy` package.\n'
@@ -15,20 +17,14 @@
 if TYPE_CHECKING:
     from sqlalchemy import Engine
     from sqlalchemy.ext.asyncio import AsyncEngine
-    from typing_extensions import TypedDict, Unpack
-
-    class CommenterOptions(TypedDict, total=False):
-        db_driver: bool
-        db_framework: bool
-        opentelemetry_values: bool
-
-    class SQLAlchemyInstrumentKwargs(TypedDict, total=False):
-        enable_commenter: bool | None
-        commenter_options: CommenterOptions | None
-        skip_dep_check: bool
 
 
-def instrument_sqlalchemy(engine: AsyncEngine | Engine | None, **kwargs: Unpack[SQLAlchemyInstrumentKwargs]) -> None:
+def instrument_sqlalchemy(
+    engine: 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.

diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py
@@ -84,7 +84,7 @@
     from typing_extensions import Unpack
 
     from ..integrations.flask import (
-        CommenterOptions,
+        CommenterOptions as FlaskCommenterOptions,
         RequestHook as FlaskRequestHook,
         ResponseHook as FlaskResponseHook,
     )
@@ -94,6 +94,7 @@
         RequestHook as HttpxRequestHook,
         ResponseHook as HttpxResponseHook,
     )
+    from ..integrations.sqlalchemy import CommenterOptions as SQLAlchemyCommenterOptions
     from ..integrations.wsgi import (
         RequestHook as WSGIRequestHook,
         ResponseHook as WSGIResponseHook,
@@ -104,7 +105,6 @@
     from .integrations.psycopg import PsycopgInstrumentKwargs
     from .integrations.pymongo import PymongoInstrumentKwargs
     from .integrations.redis import RedisInstrumentKwargs
-    from .integrations.sqlalchemy import SQLAlchemyInstrumentKwargs
     from .integrations.sqlite3 import SQLite3Connection
     from .integrations.system_metrics import Base as SystemMetricsBase, Config as SystemMetricsConfig
     from .utils import SysExcInfo
@@ -1438,7 +1438,7 @@ def instrument_flask(
         *,
         capture_headers: bool = False,
         enable_commenter: bool = True,
-        commenter_options: CommenterOptions | None = None,
+        commenter_options: FlaskCommenterOptions | None = None,
         exclude_urls: str | None = None,
         request_hook: FlaskRequestHook | None = None,
         response_hook: FlaskResponseHook | None = None,
@@ -1618,7 +1618,9 @@ def instrument_aiohttp_client(self, **kwargs: Any) -> None:
     def instrument_sqlalchemy(
         self,
         engine: AsyncEngine | Engine | None = None,
-        **kwargs: Unpack[SQLAlchemyInstrumentKwargs],
+        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.
 
@@ -1628,14 +1630,18 @@ def instrument_sqlalchemy(
 
         Args:
             engine: The `sqlalchemy` engine to instrument, or `None` to instrument all engines.
+            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.
         """
         from .integrations.sqlalchemy import instrument_sqlalchemy
 
         self._warn_if_not_initialized_for_instrumentation()
         return instrument_sqlalchemy(
             engine=engine,
-            **{  # type: ignore
+            enable_commenter=enable_commenter,
+            commenter_options=commenter_options or {},
+            **{
                 'tracer_provider': self._config.get_tracer_provider(),
                 'meter_provider': self._config.get_meter_provider(),
                 **kwargs,

diff --git a/logfire/integrations/sqlalchemy.py b/logfire/integrations/sqlalchemy.py
@@ -0,0 +1,12 @@
+from typing import TypedDict
+
+
+class CommenterOptions(TypedDict, total=False):
+    """The `commenter_options` parameter for `instrument_sqlalchemy`."""
+
+    db_driver: bool
+    """Include the database driver name in the comment e.g. 'psycopg2'."""
+    db_framework: bool
+    """Enabling this flag will add the database framework name and version to the comment e.g. 'sqlalchemy:1.4.0'."""
+    opentelemetry_values: bool
+    """Enabling this flag will add traceparent values to the comment."""