Scripts for alembic to assert that the database is up to date, migrations are up to date and to run migrations programmatically

alembic_migrations.py

from alembic import config
from alembic import script
from alembic.autogenerate import compare_metadata
from alembic.runtime import migration
from alembic.runtime.environment import EnvironmentContext

# other imports ...

def assert_database_is_up_to_date():
    """ Database is up to date with migration head version """
    alembic_cfg = config.Config("alembic.ini")
    script_directory = script.ScriptDirectory.from_config(alembic_cfg)
    with engine.begin() as conn:
        context = migration.MigrationContext.configure(conn)
        if context.get_current_revision() != script_directory.get_current_head():
            logger.error("Database is not up to date")
            logger.warning("Upgrade the database")
            exit(1)


def assert_migrations_are_up_to_date():
    """ SQLAlchemy models are up to date with the database """
    with engine.begin() as conn:
        context = migration.MigrationContext.configure(conn)
        diff = compare_metadata(context, Base.metadata)
        if len(diff) > 0:
            logger.error(
                "Migrations are not up to date. The following changes have been detected:\n"
                + "\n".join(str(d) for d in diff)
            )
            logger.warning("Create a new revision")
            exit(1)


def run_migrations():
    """
    Run alembic migrations
    source: https://stackoverflow.com/questions/39021059/how-to-run-a-migration-with-python-alembic-by-code
    """
    alembic_cfg = config.Config("alembic.ini")
    alembic_script = script.ScriptDirectory.from_config(alembic_cfg)
    alembic_env = EnvironmentContext(alembic_cfg, alembic_script)

    def do_upgrade(revision, context):
        return alembic_script._upgrade_revs(alembic_script.get_heads(), revision)

    with engine.begin() as conn:
        alembic_env.configure(connection=conn, target_metadata=Base.metadata, fn=do_upgrade)
        with alembic_env.begin_transaction():
            alembic_env.run_migrations()