I'd like to speed my integration tests a bit and execute raw SQL code equivalent to create_all()
My idea is to run create_all (in order to get it SQL equivalent) just once when the test session starts and use SQL code between the tests to migrate the tables.
Do you have any idea how it can be done?
Thanks in advance!
You can accomplish the task by hooking your code with sqlalachemy after_cursor_execute hook.
class QueryLogger:
"""Log query duration and SQL as a context manager."""
def __init__(self,
engine: sqlalchemy.engine.Engine,
f: io.StringIO):
"""
Initialize for an engine and file.
engine: The sqlalchemy engine for which events should be logged.
You can pass the class `sqlalchemy.engine.Engine` to capture all engines
f: file you want to write your output to
"""
self.engine = engine
self.file = f
def _after_cursor_execute(self, conn, cursor, statement, parameters, context, executemany):
"""Listen for the 'after_cursor_execute' event and log sqlstatement and time."""
# check if it's a ddl operation create_all execute a bunch of select statements
if context.isddl:
s = statement % parameters
self.file.write(f"{s};")
def __enter__(self, *args, **kwargs):
"""Context manager."""
if isinstance(self.engine, sqlalchemy.engine.Engine):
sqlalchemy.event.listen(self.engine, "after_cursor_execute", self._after_cursor_execute)
return self
def __exit__(self, *args, **kwargs) -> None:
"""Context manager."""
if isinstance(self.engine, sqlalchemy.engine.Engine):
sqlalchemy.event.remove(self.engine, "after_cursor_execute", self._after_cursor_execute)
And then you can use the context manager to log the queries to an in-memory file for write to SQL
with open("x.sql", "w") as f:
with QueryLogger(db.engine, f):
db.create_all()
A major part of the code is inspired by https://stackoverflow.com/a/67298123/3358570
