profiling-db

services/db.py

@@ -1,11 +1,12 @@
 import math
+import pstats
 import time
 from functools import wraps
+from sqlalchemy import event, Engine
 from typing import Any, Callable, Dict, TypeVar
 
 from dogpile.cache import make_region
-from sqlalchemy import Column, Integer, create_engine, event
-from sqlalchemy.engine import Engine
+from sqlalchemy import Column, Integer, create_engine
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.schema import Table
@@ -21,26 +22,46 @@ T = TypeVar("T")
 REGISTRY: Dict[str, type] = {}
 Base = declarative_base()
 
+def profile_sqlalchemy_queries(threshold=0.1):
+    def decorator(fn):
+        @wraps(fn)
+        def wrapper(*args, **kw):
+            elapsed, stat_loader, result = _profile(fn, threshold, *args, **kw)
+            if elapsed is not None:
+                print(f"Query took {elapsed:.3f} seconds to execute.")
+                stats = stat_loader()
+                stats.sort_stats('cumulative')
+                stats.print_stats()
+            return result
+        return wrapper
+    return decorator
+
+def _profile(fn, threshold, *args, **kw):
+    began = time.time()
+    result = fn(*args, **kw)
+    ended = time.time()
+
+    if ended - began > threshold:
+        return ended - began, pstats.Stats, result
+    else:
+        return None, None, result
 
 # Перехватчики для журнала запросов SQLAlchemy
 @event.listens_for(Engine, "before_cursor_execute")
 def before_cursor_execute(conn, cursor, statement, parameters, context, executemany):
-    query_id = id(cursor)  # Уникальный идентификатор запроса
-    conn.info.setdefault(f"query_start_time:{query_id}", time.time())
-
+    conn._query_start_time = time.time()
 
 @event.listens_for(Engine, "after_cursor_execute")
 def after_cursor_execute(conn, cursor, statement, parameters, context, executemany):
-    query_id = id(cursor)  # Уникальный идентификатор запроса
-    if f"query_start_time:{query_id}" in conn.info:
-        total = time.time() - conn.info.get(f"query_start_time:{query_id}", time.time())
-        try:
-            del conn.info[f"query_start_time:{query_id}"]
-            stars = "*" * math.floor(total * 1000)
-            if stars:
-                logger.debug(f"\n{statement}\n {stars} {total*1000} s\n")
-        except Exception:
-            pass
+    if hasattr(conn, '_query_start_time'):
+        elapsed = time.time() - conn._query_start_time
+        del conn._query_start_time
+        if elapsed > 0.2:  # Adjust threshold as needed
+            print(f"{'*' * math.floor(elapsed)} {elapsed:.3f} seconds to execute.")
+            # Profile the query if execution time exceeds the threshold
+            profiler = profile_sqlalchemy_queries(threshold=0.2)(cursor.execute)
+            profiler(statement, parameters)
+
 
 
 def local_session(src=""):