Python

When a Flask DB Upgrade Fails on a Batch of Scripts

What happens when 1 script in a "flask db upgrade" batch fails, how does the roll back behave?

Luke Merrett

13 Nov 2021 • 5 min read

On a client project recently we wanted to understand what happens when a flask db upgrade operation against a Postgres database fails.

More specifically we wanted to find out what happens when a batch of scripts is being applied, and 1 script in the middle of the batch fails, e.g: in this scenario:

Given a migration set consisting of [X, Y, Z]
When the migration X succeeds
And the migration Y fails
Then ?

Flask could do one of two things:

Leave script X applied to the database, roll back script Y
Roll back both scripts X and Y

However in the documentation it wasn't wholly clear what it would do.

We theorised the following would happen:

Scenario: Single migration in set succeeds
    Given a migration set consisting of [X]
    When the migration X succeeds
    Then the migration X is applied to the database
    
Scenario: Single migration in set fails
    Given a migration set consisting of [X]
    When the migration X fails
    Then the migration X is rolled back
    
Scenario: Multiple migrations in set, 1 fails in the set
    Given a migration set consisting of [X, Y, Z]
    When the migration X succeeds
    And the migration Y fails
    Then the migration Y is rolled back
    And the migration X is rolled back
    And the migration Z is never applied

How would we test this though?

Validating the Theory

All the below code can be found in this repository along with instructions on how to set this up locally: https://github.com/lukemerrett/flask-migration-rollback-research

Environment

Firstly we'll provision an environment with a Postgres database running using Docker compose, hosted on port 6456.

# docker-compose.yml

version: "3.9"

services:
  db:
    image: postgres
    restart: always
    environment:
      - POSTGRES_DB=postgres
      - POSTGRES_USER=postgres
      - POSTGRES_PASSWORD=postgres
    ports:
      - "6456:5432"

Then spin it up using docker compose up -d

Application

Next we have an application that is used to generate 4 dummy database migration scripts that looks like this:

# app.py

from flask import Flask
from flask_sqlalchemy import SQLAlchemy
from flask_migrate import Migrate

app = Flask(__name__)
app.config["SQLALCHEMY_DATABASE_URI"] = "postgresql://postgres:postgres@localhost:6456/postgres"
app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False

db = SQLAlchemy(app)
migrate = Migrate(app, db)


class User(db.Model):
    name = db.Column(db.String(128), primary_key=True)
    date = db.Column(db.DateTime())
    department = db.Column(db.String(128))

Modifying the User model and running flask db migrate will generate a new script applying those changes. The above code is the "final state" of the migration.

Migration Scripts

Through modifying the User and running flask db migrate a few times we end up with the below 4 migration scripts:

Script 1: dd7c69423ad9_.py

Creates the user table with an id integer primary key and a name string field.

"""empty message
Revision ID: dd7c69423ad9
Revises: 
Create Date: 2021-08-20 09:41:57.938972
"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = 'dd7c69423ad9'
down_revision = None
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('user',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sa.String(length=128), nullable=True),
    sa.PrimaryKeyConstraint('id')
    )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('user')
    # ### end Alembic commands ###

Script 2: 5a77262c502b_.py

Adds the date column of type date time to the user table:

"""empty message
Revision ID: 5a77262c502b
Revises: dd7c69423ad9
Create Date: 2021-08-20 11:57:30.389344
"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '5a77262c502b'
down_revision = 'dd7c69423ad9'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('user', sa.Column('date', sa.DateTime(), nullable=True))
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('user', 'date')
    # ### end Alembic commands ###

Script 3: 653d68312991_.py - Broken Script

This is a modified script that will fail, it tries to drop a column on a table that doesn't exist. We've had to manually add that line in to break it.

When flask db upgrade reaches this script, it will error.

"""empty message
Revision ID: 653d68312991
Revises: 5a77262c502b
Create Date: 2021-08-20 12:04:14.107041
"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '653d68312991'
down_revision = '5a77262c502b'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column('user', 'name',
               existing_type=sa.VARCHAR(length=128),
               nullable=False)
    op.drop_column('jim', 'id')
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('user', sa.Column('id', sa.INTEGER(), autoincrement=True, nullable=False))
    op.alter_column('user', 'name',
               existing_type=sa.VARCHAR(length=128),
               nullable=True)
    # ### end Alembic commands ###

Script 4: 5929dc9c624a_.py

The final script, one we expect to never run due to the broken script before it, adds a department field to the user table:

"""empty message
Revision ID: 5929dc9c624a
Revises: 653d68312991
Create Date: 2021-08-20 12:26:57.567531
"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '5929dc9c624a'
down_revision = '653d68312991'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('user', sa.Column('department', sa.String(length=128), nullable=True))
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('user', 'department')
    # ### end Alembic commands ###

Running the Upgrade Tests

Now we need to test our 3 scenarios with the 4 scripts created above.

Scenario: Single migration in set succeeds
    Given a migration set consisting of [X]
    When the migration X succeeds
    Then the migration X is applied to the database

To test this we run flask db upgrade dd7c69423ad9 (the first script in our batch) which applies a valid script to the database. This succeeds as expected and the database is brought up to date with the new user table.

Scenario: Single migration in set fails
    Given a migration set consisting of [X]
    When the migration X fails
    Then the migration X is rolled back

To test this we first run flask db upgrade 5a77262c502b, which is the second script that we also know works correctly. The database now has the first 2 scripts applied and is ready for the broken script to be run,

Running flask db upgrade 653d68312991 (the broken script) correctly fails and any changes are rolled back. Because we've run scripts dd7c69423ad9 and 5a77262c502b in separate upgrade operations they are still applied to the database, in this scenario only script 653d68312991 is rolled back.

Scenario: Multiple migrations in set, 1 fails in the set
    Given a migration set consisting of [X, Y, Z]
    When the migration X succeeds
    And the migration Y fails
    Then the migration Y is rolled back
    And the migration X is rolled back
    And the migration Z is never applied

Firstly we need to get our database into the correct state so there is a batch of upgrades that starts with 1 valid script, then a broken one, then a valid one again.

To do this we run:

# Ensure we've at least applied the stable first revision
flask db upgrade dd7c69423ad9
# Ensure if other revisions have been applied since, we return to the stable revision
flask db downgrade dd7c69423ad9

Leaving us with scripts 5a77262c502b (working), 653d68312991 (broken) and 5929dc9c624a (working) in a batch to be applied on the next upgrade.

Now running flask db upgrade will do the following:

Apply 5a77262c502b, succeeded
Apply 653d68312991, failed
Roll back 653d68312991
Roll back 5a77262c502b
Never attempt to run 5929dc9c624a

Findings

If any of the scripts in a batch fail during a flask db upgrade, then none of the scripts in that batch are applied.

This means the database isn't left in a state where only partial migrations have been applied.

This proves all our theorised scenarios are correct.

This works exactly the same in multidb mode; if 1 script across any of the databases fails on upgrade, every script in the batch across all databases is rolled back (see the behaviour-on-multi-db branch in the aforementioned repo).