Python 对Alembic升级进行试运行

有时，

alembic upgrade head

可能会在运行时对我的生产数据库失败，即使它对我的测试数据库运行良好。例如，迁移可能会将

notnull

约束添加到一个列中，该列以前在我的测试环境中不包含

NULL

s，但在生产环境中确实包含

NULL

s

在规划部署时，最好能够在运行迁移之前检查迁移是否能够干净地应用。对于不支持事务性DDL（在事务中进行模式更改）的数据库（如MySQL），这可能是不可能的，但对于支持事务性DDL的数据库（如PostgreSQL），这在原则上应该是可能的；Alembic可以尝试在事务中执行升级，然后回滚

（一个警告：这是一个不完美的解决方案，因为PostgreSQL允许延迟一些约束，这意味着在提交之前不会检查它们。我认为，如果不创建数据库副本，检查这些约束的试运行是不可能的。但是，执行DDL和回滚方法总比什么都不做要好。）

---

Alembic是否支持此类功能？如果没有，是否有一些黑客的方法来实现它？

您可以创建自己的alembic配置、脚本目录和环境上下文。关键的是，这应该有一个已经创建的连接传递给它，您可以在完成试运行后回滚该连接

engine = sa.create_engine(url)
session = Session(bind=engine)
try:
    alembic_config = AlembicConfig(config_args={...})
    alembic_config.attributes.update({"connection": session.connection()})
    script_directory = ScriptDirectory.from_config(alembic_config)
    env_context = EnvironmentContext(
        alembic_config,
        script_directory,
        fn=lambda rev, context: script_directory._upgrade_revs("head", rev),
        as_sql=False,
    )
    env_context.configure(connection=session.connection())
    env_context.run_migrations()

finally:
    session.rollback()
    session.close()

---

允许此操作的一个简单技巧是将条件回滚注入到

env.py

中的

run\u migrations\u online

函数中，该函数仅在出现指示我们需要干运行的标志时触发

如果您的已修改，请回想一下由

alembic init

创建的

run\u migrations\u online

函数的默认实现如下所示：

def run_migrations_online():
    """Run migrations in 'online' mode.

    In this scenario we need to create an Engine
    and associate a connection with the context.

    """
    connectable = engine_from_config(
        config.get_section(config.config_ini_section),
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )

    with connectable.connect() as connection:
        context.configure(
            connection=connection, target_metadata=target_metadata
        )

        with context.begin_transaction():
            context.run_migrations()

请注意：

• context.begin\u transaction（）

  的

  \uuuu enter\uuuu

  方法（我们已经在默认实现中调用了该方法）为我们提供了一个带有

  回滚（）方法的事务对象，以及

• 我们的

  context

  对象有一个方法，可以用来支持将自定义参数传递给

  alembic

  命令

因此，通过以下小改动（除了添加

as transaction

和最后三行之外，以下所有内容都是相同的），我们就可以拥有我们的试运行功能：

def run_migrations_online():
    """Run migrations in 'online' mode.

    In this scenario we need to create an Engine
    and associate a connection with the context.

    """
    connectable = engine_from_config(
        config.get_section(config.config_ini_section),
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )

    with connectable.connect() as connection:
        context.configure(
            connection=connection, target_metadata=target_metadata
        )
        with context.begin_transaction() as transaction:
            context.run_migrations()
            if 'dry-run' in context.get_x_argument():
                print('Dry-run succeeded; now rolling back transaction...')
                transaction.rollback()

现在，要进行试运行，请执行以下操作：

alembic -x dry-run upgrade head

要进行真正的跑步，只需执行以下操作：

alembic upgrade head

---

与以前一样。

这可能有效，但我有点不安，因为它没有读取alembic.ini配置，也没有按照运行真正迁移时的设置方式设置上下文对象。出于这个原因，我挠头了一会儿，直到我意识到一个更简单的解决方案是可行的——我将发布我自己的答案。