Using the SQLAlchemy ORM inside an Alembic migration: how do I?

After a bit of experimentation using @velochy’s answer, I settled on something like the following pattern for using SqlAlchemy inside Alembic. This worked great for me and could probably serve as a general solution for the OP’s question:

from sqlalchemy.orm.session import Session
from alembic import op

# Copy the model definitions into the migration script if
# you want the migration script to be robust against later
# changes to the models. Also, if your migration includes
# deleting an existing column that you want to access as 
# part of the migration, then you'll want to leave that 
# column defined in the model copies here.
class Model1(Base): ...
class Model2(Base): ...

def upgrade():
    # Attach a sqlalchemy Session to the env connection
    session = Session(bind=op.get_bind())

    # Perform arbitrarily-complex ORM logic
    instance1 = Model1(foo='bar')
    instance2 = Model2(monkey='banana')

    # Add models to Session so they're tracked
    session.add(instance1)
    session.add(instance2)

    # Apply a transform to existing data
    m1s = session.query(Model1).all()
    for m1 in m1s:
        m1.foo = transform(m1.foo)
    session.commit()

def downgrade():
    # Attach a sqlalchemy Session to the env connection
    session = Session(bind=op.get_bind())

    # Perform ORM logic in downgrade (e.g. clear tables)
    session.query(Model2).delete()
    session.query(Model1).delete()

    # Revert transform of existing data
    m1s = session.query(Model1).all()
    for m1 in m1s:
        m1.foo = un_transform(m1.foo)
    session.commit()

This approach appears to handle transactions properly. Frequently while working on this, I would generate DB exceptions and they would roll things back as expected.

Leave a Comment

Hata!: SQLSTATE[HY000] [1045] Access denied for user 'divattrend_liink'@'localhost' (using password: YES)