Skip to content

Commit 082f555

Browse files
committed
enhance QueryRecorder to track which database bind is used for each query
Replace flask_sqlalchemy's get_recorded_queries with SQLAlchemy event listeners that capture the bind_key for each engine. Add QueryInfo dataclass with statement, parameters, and bind_key fields. This enables tests to verify that queries are being routed to the correct database when using db.session_bulk. This is an alternative to the existing PR pallets-eco/flask-sqlalchemy#1403 for flask-sqlalchemy.
1 parent aae5d26 commit 082f555

File tree

1 file changed

+37
-5
lines changed

1 file changed

+37
-5
lines changed

tests/utils.py

Lines changed: 37 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,46 @@
1-
import flask_sqlalchemy
1+
from dataclasses import dataclass
2+
3+
from sqlalchemy import event
4+
5+
from app import db
6+
7+
8+
@dataclass
9+
class QueryInfo:
10+
statement: str
11+
parameters: tuple | dict | None
12+
bind_key: str | None
213

314

415
class QueryRecorder:
516
def __init__(self):
6-
self.queries = []
7-
self._count_on_enter = None
17+
self.queries: list[QueryInfo] = []
18+
self._listeners = []
819

920
def __enter__(self):
10-
self._count_on_enter = len(flask_sqlalchemy.record_queries.get_recorded_queries())
21+
# Register listeners for all engines to capture bind_key
22+
for bind_key, engine in db.engines.items():
23+
listener = self._listener(bind_key)
24+
event.listen(engine, "before_cursor_execute", listener)
25+
self._listeners.append((engine, listener))
1126
return self
1227

1328
def __exit__(self, exc_type, exc_val, exc_tb):
14-
self.queries = flask_sqlalchemy.record_queries.get_recorded_queries()[self._count_on_enter :]
29+
# Remove all listeners
30+
for engine, listener in self._listeners:
31+
event.remove(engine, "before_cursor_execute", listener)
32+
self._listeners.clear()
33+
34+
def _listener(self, bind_key):
35+
"""Create a listener function that captures the bind_key in its closure."""
36+
37+
def listener(conn, cursor, statement, parameters, context, executemany):
38+
self.queries.append(
39+
QueryInfo(
40+
statement=statement,
41+
parameters=parameters,
42+
bind_key=bind_key,
43+
)
44+
)
45+
46+
return listener

0 commit comments

Comments
 (0)