Improved sqlalchemy panel

Author: mongkok

debug_toolbar/panels/sqlalchemy.py

@@ -5,8 +5,7 @@
 from fastapi.concurrency import AsyncExitStack
 from fastapi.dependencies.utils import solve_dependencies
 from sqlalchemy import event
-from sqlalchemy.engine import Connection, Engine
-from sqlalchemy.engine.default import DefaultExecutionContext
+from sqlalchemy.engine import Connection, Engine, ExecutionContext
 from sqlalchemy.orm import Session
 
 from debug_toolbar.panels.sql import SQLPanel
@@ -20,40 +19,26 @@ def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
         self.engines: t.Set[Engine] = set()
 
     def register(self, engine: Engine) -> None:
-        event.listen(engine, "before_cursor_execute", self.before_execute)
-        event.listen(engine, "after_cursor_execute", self.after_execute)
+        event.listen(engine, "before_cursor_execute", self.before_execute, named=True)
+        event.listen(engine, "after_cursor_execute", self.after_execute, named=True)
 
     def unregister(self, engine: Engine) -> None:
         event.remove(engine, "before_cursor_execute", self.before_execute)
         event.remove(engine, "after_cursor_execute", self.after_execute)
 
-    def before_execute(
-        self,
-        conn: Connection,
-        cursor: t.Any,
-        statement: str,
-        parameters: t.Union[t.Sequence, t.Dict],
-        context: DefaultExecutionContext,
-        executemany: bool,
-    ) -> None:
-        conn.info.setdefault("start_time", []).append(perf_counter())
+    def before_execute(self, context: ExecutionContext, **kwargs: t.Any) -> None:
+        context._start_time = perf_counter()  # type: ignore[attr-defined]
 
-    def after_execute(
-        self,
-        conn: Connection,
-        cursor: t.Any,
-        statement: str,
-        parameters: t.Union[t.Sequence, t.Dict],
-        context: DefaultExecutionContext,
-        executemany: bool,
-    ) -> None:
+    def after_execute(self, context: ExecutionContext, **kwargs: t.Any) -> None:
         query = {
-            "duration": (perf_counter() - conn.info["start_time"].pop(-1)) * 1000,
-            "sql": statement,
-            "params": parameters,
-            "is_select": context.invoked_statement.is_select,
+            "duration": (
+                perf_counter() - context._start_time  # type: ignore[attr-defined]
+            )
+            * 1000,
+            "sql": context.statement,
+            "params": context.parameters,
         }
-        self.add_query(str(conn.engine.url), query)
+        self.add_query(str(context.engine.url), query)
 
     async def add_engines(self, request: Request):
         route = request["route"]
@@ -70,7 +55,12 @@ async def add_engines(self, request: Request):
             )
             for value in solved_result[0].values():
                 if isinstance(value, Session):
-                    self.engines.add(value.get_bind())
+                    bind = value.get_bind()
+
+                    if isinstance(bind, Connection):
+                        self.engines.add(bind.engine)
+                    else:
+                        self.engines.add(bind)
 
     async def process_request(self, request: Request) -> Response:
         await self.add_engines(request)