Recreating Foreign Keys with Alembic

Alembic is a great tool for keeping track of schema changes in python applications. I am using it to manage DB migrations for braindump along with Flask SQL Alchemy as my ORM. One challenge is managing proper foreign key constraints. By default if you define a foreign key relationship in your schema definition it will not generate the proper migration code.

For example, in braindump we have a one to many relationship between users and notes.

class User(UserMixin, db.Model):
    __tablename__ = 'users'
    id = db.Column(db.Integer, primary_key=True)
    email = db.Column(db.String(254), unique=True, index=True)
    password_hash = db.Column(db.String(256))
    confirmed = db.Column(db.Boolean, default=False)
    avatar_hash = db.Column(db.String(32))
    created_date = db.Column(db.DateTime(), default=datetime.utcnow)
    updated_date = db.Column(db.DateTime(), default=datetime.utcnow)

    notes = db.relationship(
        'Note', backref='author',
        lazy='dynamic', cascade="all, delete-orphan")

Even though we define the cascade behavior using SQLAlchemy. When we generate the migration with alembic we get something like this:

sa.ForeignKeyConstraint(['author_id'], ['users.id'], ),

Notice how we are missing the ondelete action. What we actually want is something like this:

sa.ForeignKeyConstraint(['author_id'], ['users.id'], ondelete='CASCADE')

Running the default migration will not create the proper relationship in your database and in our case we are not able to delete a user until we have deleted all of the related records as well. There are two ways to fix this. If you catch this before running your migration adding ondelete='CASCADE' will create the proper relationship. If you are like me, and do not catch this, then you will need to run a second migration to remove and recreate these keys. The migration code to do this is shown below:

from alembic import op
import sqlalchemy as sa

def upgrade():
    with op.batch_alter_table("notes") as batch_op:
        batch_op.drop_constraint(
            "notes_author_id_fkey", type_="foreignkey")
    op.create_foreign_key(
        "notes_author_id_fkey", "notes", "users",
        ["author_id"], ["id"], ondelete="CASCADE")

Now you have the proper foreign key constraints and the CASCADE action exists in the DB.

If you made it this far, you should probably follow me on twitter. 🙂 Follow @levlaz