Do a dry-run of an Alembic upgrade

Sometimes alembic upgrade head may fail at runtime against my production database even though it worked fine against my testing databases. For instance, a migration might add a NOT NULL constraint to a column that didn’t previously contain NULLs in my testing environment, but did contain NULLs in production.

When planning deployments, it’d be nice to able to check in advance of running a migration whether it will manage to apply cleanly. This is presumably impossible for databases that don’t support transactional DDL (making schema changes in a transaction) like MySQL, but ought to be possible in principle for databases that do support transactional DDL, like PostgreSQL; Alembic could just try performing the upgrade in a transaction, then roll it back.

(One caveat: this is an imperfect solution, since PostgreSQL permits some constraints to be DEFERRED, meaning they aren’t checked until you commit. A dry-run that checks these is, I suppose, impossible without creating a copy of the database. But still, doing the DDL-and-rollback approach is better than nothing.)

Does Alembic have any support for such functionality? If not, is there some hacky way to achieve it?

Answer

A simple trick to allow this is to inject a conditional rollback into the run_migrations_online function in env.py that fires only when some flag is present indicating that we want a dry-run.

In case yours is already modified, recall the default implementation of the run_migrations_online function created by alembic init looks like this:

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()

Note that:

• the __enter__ method of context.begin_transaction() – which we’re already calling in the default implementation – gives us a transaction object with a rollback() method, if the backend uses transactional DDL, or if transactional ddl is forced on using the transactional_ddl flag, and
• our context object has a get_x_argument method we can use to support passing custom arguments to the alembic command.

Thus, with the following small change (everything below is the same besides the addition of as transaction plus the final three lines) we can have our dry-run functionality:

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,

        # Ensure the context will create a transaction
        # for backends that don't normally use transactional DDL.
        # Note that ROLLBACK will not roll back DDL structures
        # on databases such as MySQL, as well as with the SQLite
        # Python driver's default settings.
        transactional_ddl=True,
    )

    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()

Now, to do a dry-run, do:

alembic -x dry-run upgrade head

and to do a real run, just do:

alembic upgrade head

like before.