Flat alembic multi-branches in one linear branch

flat_alembic_versions.py

# -*- coding: utf-8 -*-
"""
Flat alembic versions

A______
|  |   |
B  B1  ...
|  |
C  C1
|
D

into

A
|
B
|
C
|
D
|
B1
|
C1
|
...

Asumming B1's branch is not in the integration git branch (ie "develop")
"""
import argparse
import re
import sys
import logging
from alembic.config import Config
from alembic.script import ScriptDirectory
from alembic import command
import subprocess

logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)


def main():

    parser = argparse.ArgumentParser(description='Flat alembic migrations')
    parser.add_argument('ini', help='PasteDeploy .ini file')
    parser.add_argument('--git-integration-branch', help='default: %(default)s', default='origin/develop')
    parser.add_argument('--apply', action='store_true', help='Apply downgrade until branch point and upgrade the flat version')
    args = parser.parse_args()

    root = subprocess.check_output(['git', 'rev-parse', '--show-toplevel'])[:-1] + '/'

    alembic_cfg = Config(args.ini, 'app:main')
    script = ScriptDirectory.from_config(alembic_cfg)

    def branch_detail(rev):
            return list(script.walk_revisions(rev))

    def is_in_branch(rev_sc, branch='develop'):
        path = rev_sc.path.replace(root, '')
        try:
            subprocess.check_call(
                ['git', 'show', '{}:{}'.format(branch, path)],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            logging.info('%s is in %s', path, branch)
            return True
        except subprocess.CalledProcessError:
            logging.info('%s is not in %s', path, branch)
            return False

    def move(rev_sc, to):
        with open(rev_sc.path, 'r') as i:
            code = i.read()
        code = re.sub(
            r'down_revision = [\da-f\'"]+',
            "down_revision = '{}'".format(to.revision),
            code
        )
        with open(rev_sc.path, 'w') as o:
            o.write(code)

    branches = []
    for sc in script.walk_revisions():
        if sc.is_branch_point:
            branches = [branch_detail(v) for v in sc.nextrev]     # B and B1 in the example tree
            break
    else:
        # no branch point
        return

    if args.apply:
        logging.info('Applying downgrade until %s (branching point)', sc.revision)
        command.downgrade(alembic_cfg, sc.revision)

    # branches is a list representing each branch's path,
    # sorting so the first is the one that are already in git "develòp"
    branches = sorted(branches, key=lambda l: not is_in_branch(l[0], args.git_integration_branch))
    current_head = branches[0][0]    # this would be D in our
    for h in branches[1:]:
        move(h[-1], current_head)
        current_head = h[0]

    if args.apply:
        logging.info('Applying upgrade to head')
        command.upgrade(alembic_cfg, 'head')


if __name__ == '__main__':
    main()

Enjoy!