首页 > 解决方案 > 获取`create_all`的编译SQL

问题描述

我想加快我的集成测试速度并执行相当于create_all()

我的想法是create_all在测试会话开始时只运行一次(以使其与 SQL 等效)并在测试之间使用 SQL 代码来迁移表。

你知道怎么做吗?

提前致谢!

标签: pythonsqlalchemyflask-sqlalchemyalembic

解决方案


您可以通过使用 sqlalachemyafter_cursor_execute挂钩挂钩您的代码来完成任务。

https://docs.sqlalchemy.org/en/13/core/events.html#sqlalchemy.events.ConnectionEvents.after_cursor_execute

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)

然后您可以使用上下文管理器将查询记录到内存文件以写入 SQL

with open("x.sql", "w") as f:
    with QueryLogger(db.engine, f):
        db.create_all()

代码的主要部分灵感来自https://stackoverflow.com/a/67298123/3358570


推荐阅读