Files
notifications-api/scripts/fix_migrations.py
Kenneth Kehl 1ecb747c6d reformat
2023-08-29 14:54:30 -07:00

103 lines
2.7 KiB
Python
Executable File

#!/usr/bin/env python
# encoding: utf-8
import os
import sys
from alembic.script import ScriptDirectory
sys.path.append(".")
def get_branch_points(migrations):
return [m for m in migrations.walk_revisions() if m.is_branch_point]
def get_branches(migrations, branch_point, heads):
return [
list(migrations.iterate_revisions(m, branch_point.revision))[::-1]
for m in heads
]
def choice(prompt, options, option_fmt=lambda x: x):
print("{}:\n".format(prompt))
for i, option in enumerate(options):
print("{}. {}".format(i + 1, option_fmt(option)))
print()
choice = input("Option> ")
return options[int(choice) - 1]
def rename_revision(current_revision, new_base):
new_id = int(new_base[:4]) + 1
return "{:04d}{}".format(new_id, current_revision[4:])
def reorder_revisions(revisions, old_base, new_base):
if not revisions:
return
head, *tail = revisions
new_revision_id = rename_revision(head.revision, new_base)
print("Moving {} to {}".format(head.revision, new_revision_id))
with open(head.path, "r") as rev_file:
file_data = rev_file.read()
file_data = file_data.replace(head.revision, new_revision_id).replace(
old_base, new_base
)
new_filename = head.path.replace(head.revision, new_revision_id)
assert (
head.path != new_filename
), "Old filename not same as revision id, please rename file before continuing"
with open(new_filename, "w") as rev_file:
rev_file.write(file_data)
print("Removing {}".format(head.path))
os.remove(head.path)
reorder_revisions(tail, head.revision, new_revision_id)
def fix_branch_point(migrations, branch_point, heads):
print("Migrations directory has a branch point at {}".format(branch_point.revision))
branches = get_branches(migrations, branch_point, heads)
move_branch = choice(
"Select migrations to move",
branches,
lambda x: " -> ".join(m.revision for m in x),
)
branches.remove(move_branch)
reorder_revisions(move_branch, branch_point.revision, branches[0][-1].revision)
def main(migrations_path):
migrations = ScriptDirectory(migrations_path)
branch_points = get_branch_points(migrations)
heads = migrations.get_heads()
if not branch_points:
print("Migrations are ordered")
elif len(branch_points) == 1 and len(heads) == 2:
fix_branch_point(migrations, branch_points[0], heads)
else:
print(
"Found {} branch points and {} heads, can't fix automatically".format(
[bp.revision for bp in branch_points], heads
)
)
sys.exit(1)
if __name__ == "__main__":
main("migrations/")