fastapi-radar 0.1.9__py3-none-any.whl → 0.2.0__py3-none-any.whl

fastapi_radar/__init__.py

@@ -2,5 +2,5 @@
 
 from .radar import Radar
 
-__version__ = "0.1.9"
+__version__ = "0.2.0"
 __all__ = ["Radar"]

fastapi_radar/capture.py

@@ -1,13 +1,18 @@
 """SQLAlchemy query capture for FastAPI Radar."""
 
 import time
-from typing import Any, Callable, Dict, List, Union
-
+from typing import Any, Callable, Dict, List, Optional, Union
 from sqlalchemy import event
 from sqlalchemy.engine import Engine
 
+try:  # SQLAlchemy async support is optional
+    from sqlalchemy.ext.asyncio import AsyncEngine
+except Exception:  # pragma: no cover - module might not exist in older SQLAlchemy
+    AsyncEngine = None  # type: ignore[assignment]
 from .middleware import request_context
 from .models import CapturedQuery
+
+
 from .utils import format_sql
 from .tracing import get_current_trace_context
 
@@ -23,14 +28,20 @@ class QueryCapture:
         self.capture_bindings = capture_bindings
         self.slow_query_threshold = slow_query_threshold
         self._query_start_times = {}
+        self._registered_engines: Dict[int, Engine] = {}
 
     def register(self, engine: Engine) -> None:
-        event.listen(engine, "before_cursor_execute", self._before_cursor_execute)
-        event.listen(engine, "after_cursor_execute", self._after_cursor_execute)
+        sync_engine = self._resolve_engine(engine)
+        event.listen(sync_engine, "before_cursor_execute", self._before_cursor_execute)
+        event.listen(sync_engine, "after_cursor_execute", self._after_cursor_execute)
+        self._registered_engines[id(engine)] = sync_engine
 
     def unregister(self, engine: Engine) -> None:
-        event.remove(engine, "before_cursor_execute", self._before_cursor_execute)
-        event.remove(engine, "after_cursor_execute", self._after_cursor_execute)
+        sync_engine = self._registered_engines.pop(id(engine), None)
+        if not sync_engine:
+            sync_engine = self._resolve_engine(engine)
+        event.remove(sync_engine, "before_cursor_execute", self._before_cursor_execute)
+        event.remove(sync_engine, "after_cursor_execute", self._after_cursor_execute)
 
     def _before_cursor_execute(
         self,
@@ -44,14 +55,14 @@ class QueryCapture:
         request_id = request_context.get()
         if not request_id:
             return
-
         context_id = id(context)
         self._query_start_times[context_id] = time.time()
-
+        setattr(context, "_radar_request_id", request_id)
         trace_ctx = get_current_trace_context()
         if trace_ctx:
             formatted_sql = format_sql(statement)
             operation_type = self._get_operation_type(statement)
+            db_tags = self._get_db_tags(conn)
             span_id = trace_ctx.create_span(
                 operation_name=f"DB {operation_type}",
                 span_kind="client",
@@ -59,6 +70,7 @@ class QueryCapture:
                     "db.statement": formatted_sql[:500],  # limit SQL length
                     "db.operation_type": operation_type,
                     "component": "database",
+                    **db_tags,
                 },
             )
             setattr(context, "_radar_span_id", span_id)
@@ -74,8 +86,9 @@ class QueryCapture:
     ) -> None:
         request_id = request_context.get()
         if not request_id:
-            return
-
+            request_id = getattr(context, "_radar_request_id", None)
+            if not request_id:
+                return
         start_time = self._query_start_times.pop(id(context), None)
         if start_time is None:
             return
@@ -160,3 +173,18 @@ class QueryCapture:
             return {k: str(v) for k, v in list(parameters.items())[:100]}
 
         return [str(parameters)]
+
+    def _resolve_engine(self, engine: Engine) -> Engine:
+        if AsyncEngine is not None and isinstance(engine, AsyncEngine):
+            return engine.sync_engine
+        return engine
+
+    def _get_db_tags(self, conn: Any) -> Dict[str, Optional[str]]:
+        tags: Dict[str, Optional[str]] = {}
+        engine = getattr(conn, "engine", None)
+        if engine and getattr(engine, "dialect", None):
+            tags["db.system"] = engine.dialect.name
+            url = getattr(engine, "url", None)
+            if url is not None:
+                tags["db.name"] = getattr(url, "database", None)
+        return tags