Bash script to relink alembic migrations

Jun 12th, 2017 in  by Michael Cho

Sick of `alembic merge heads` slowing down your builds? Here's a short bash script I wrote to rename your migration files and update its down revision.

I had written previously about using `alembic merge heads` for fixing the dreaded `Multiple head revisions are present` error. This usually when you have written a migration on your branch, but when it's time to merge to master - somebody else has committed another migration before you, meaning you are both pointing at the same downstream migration.

Using `merge heads` is fine, but if you keep doing this and you need to keep running all your migrations (eg for continuous integration tests!) this starts to slow down your build. I suspect this is because alembic has to parse through all the different heads and resolve them at runtime. 

So I wanted to write a short bash script with 2 objectives:

  1. Resolve these down revision conflicts with alembic
  2. Also rename the timestamped migration file, so that my migrations are listed on the filesystem in the same order they run

Here is the script I came up with, feel free to use in your own projects:


#!/usr/bin/env bash

###########
# A script which finds any migrations not yet merged to master, and
# ensures it is timestamped last in the directory and with an updated revision number.

# Similar in effect to `alembic merge heads`, but neater...
###########


# Common
Green='\033[0;32m'
Red='\033[0;31m'
Color_Off='\033[0m'

# eg pretty_print(color, colorised_text, normal_text)
function pretty_print {
    echo -e "$1 $2 $Color_Off $3"

}

# Step 1 : Find any migrations not yet in master
UNMERGED_MIGRATIONS=(`git diff --name-only --diff-filter=A origin/master | grep alembic/versions`)
if [ -z "$UNMERGED_MIGRATIONS" ]; then
    pretty_print $Red "No unmerged migrations!"
    exit
fi

for i in "${UNMERGED_MIGRATIONS[@]}"
do
   pretty_print $Green "Found unmerged migration: " $i
done

# Step 2 : Find the last migration merged to master
LAST_MERGED_MIGRATION=`ls -1 alembic/versions/*.py |
{
    while read N
        do
        if [[ ! ${UNMERGED_MIGRATIONS[@]} =~ "${N}" ]]; then
            echo $N
        fi
        done

} | tail -n1`

pretty_print $Green "Last merged migration: " $LAST_MERGED_MIGRATION

# Step 3 : Parse the last migration file and find the line which says "revision = 'abc123'"
LAST_REVISION=`cat $LAST_MERGED_MIGRATION | grep "revision = " | head -n1 | sed -n " s,[^']*'\([^']*\).*,\1,p "`
pretty_print $Green "Last revision: " $LAST_REVISION

LAST_MERGED_TIMESTAMP=`echo $LAST_MERGED_MIGRATION | cut -d'/' -f3 | cut -d'_' -f1`
pretty_print $Green "Last merged timestamp: " $LAST_MERGED_TIMESTAMP

# Step 4 : Update the unmerged migrations to link to the last merged
for migration in "${UNMERGED_MIGRATIONS[@]}"
do
    pretty_print $Red "------ Relinking: " $migration

    CURRENT_DOWN_REVISION=`cat $migration | grep "down_revision = " | head -n1 | sed -n " s,[^']*'\([^']*\).*,\1,p "`
    CURRENT_REVISION=`cat $migration | grep "revision = " | head -n1 | sed -n " s,[^']*'\([^']*\).*,\1,p "`
    pretty_print $Red "Modifying unmerged revision: " $CURRENT_DOWN_REVISION
    sed -i'.bak' "s/down_revision = '$CURRENT_DOWN_REVISION'/down_revision = '$LAST_REVISION'/" $migration
    rm "$migration.bak"

    # Step 5 : Move unmerged migrations to last file in directory
    LAST_UNMERGED_TIMESTAMP=`echo $migration | cut -d'/' -f3 | cut -d'_' -f1`
    LAST_UNMERGED_FILENAME=`echo $migration | cut -d'/' -f3`
    NEW_MIGRATION_TIMESTAMP=`echo $(($LAST_MERGED_TIMESTAMP + 1))`
    NEW_MIGRATION_FILENAME="${LAST_UNMERGED_FILENAME/$LAST_UNMERGED_TIMESTAMP/$NEW_MIGRATION_TIMESTAMP}"

    # Store this migrations filename and timestamp so next migration in loop is also updated correctly
    LAST_REVISION=$CURRENT_REVISION
    LAST_MERGED_TIMESTAMP=NEW_MIGRATION_TIMESTAMP

    pretty_print $Red "Renaming file: " $NEW_MIGRATION_TIMESTAMP
    mv $migration "alembic/versions/$NEW_MIGRATION_FILENAME"
done

Assuming you name this file `relink_migrations.sh`, you can just run it with ./relink_migrations.sh.

 


Other articles you may like

SQLAlchemy commit(), flush(), expire(), refresh(), merge() - what's the difference?
Nov 2nd, 2017
Prioritized Code Review Checklist - what to look for first, second, and last
Sep 21st, 2017
Many to many relationships in SQLAlchemy models (Flask)
Jul 28th, 2017
Using model callbacks in SQLAlchemy to generate slugs
Feb 7th, 2017
Overriding default Jinja dictionary attributes in Python view templates
Jan 29th, 2017