#!/usr/bin/env python import dataclasses import importlib import os import pathlib import re import click from django.apps import apps from sentry.runner import configure configure() @dataclasses.dataclass class MigrationMeta: number: int name: str path: str @property def full_name(self) -> str: return f"{self.number:04}_{self.name}" @classmethod def from_filename(cls, filepath: str) -> "MigrationMeta": filename = os.path.basename(filepath) name, _ = filename.split(".", 2) number, name = name.split("_", 1) number = int(number) return cls(number=number, name=name, path=filepath) def load_module(filepath: str): module_name = os.path.basename(filepath) spec = importlib.util.spec_from_loader( module_name, importlib.machinery.SourceFileLoader(module_name, filepath) ) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module def find_migration(migrations_path: str, migration: str) -> MigrationMeta: matches = list(pathlib.Path(migrations_path).glob(f"*{migration}*")) if len(matches) > 1: click.echo(f"Found multiple migrations matching {migration}") for match in matches: click.echo(f"- {match}") click.echo("Try again with a more specific pattern") raise click.Abort() if len(matches) == 0: click.echo(f"Could not find migration matching {migration}") raise click.Abort() return MigrationMeta.from_filename(str(matches[0].resolve())) def find_highest_migration(migrations_path: str, renamed_migration: MigrationMeta) -> MigrationMeta: highest = 0 found = None matches = list(pathlib.Path(migrations_path).glob("[0-9]*")) for match in matches: meta = MigrationMeta.from_filename(str(match.resolve())) if meta.number > highest and meta.full_name != renamed_migration.full_name: highest = meta.number found = meta if not found: click.echo("Could not find the head migration") raise click.Abort() click.echo(f"> Current head migration is {found.full_name}") return found @click.command() @click.argument("migration", required=True) @click.argument("app_label", default="sentry") def main(migration: str, app_label: str): """ Update a migration to the top of the migration history. migration - The name or number of the migration. app_label - The name of the django app the migration is in. Defaults to sentry Will do the following: - Rename the migration file. - Update `dependencies` in the migration to point at the highest migration. - Update the name of the related tests/migration if it exists. - Update migrations_lockfile.txt. """ app_path = apps.get_app_config(app_label).path migrations_path = os.path.join(app_path, "migrations") migration_meta = find_migration(migrations_path, migration) current_head = find_highest_migration(migrations_path, migration_meta) # Create an instance of the migration so that we can read instance properties. module = load_module(migration_meta.path) migration_instance = module.Migration("", app_label) dependencies = migration_instance.dependencies new_name = migration_meta.full_name.replace( str(migration_meta.number), str(current_head.number + 1) ) click.echo(f"> Updating migration {migration_meta.full_name} to {new_name}") with open(migration_meta.path) as f: contents = f.read() for dep in dependencies: if dep[0] == app_label: contents = contents.replace(dep[1], current_head.full_name) with open(migration_meta.path, "w") as f: f.write(contents) # Rename the file. os.rename(migration_meta.path, os.path.join(migrations_path, f"{new_name}.py")) click.echo("> Migration file rename complete") # Update lockfile lockfile = os.path.join(app_path, "..", "..", "migrations_lockfile.txt") with open(lockfile) as f: contents = f.read() contents = contents.replace(current_head.full_name, new_name) with open(lockfile, "w") as f: f.write(contents) click.echo("> Updated migrations_lockfile.txt") # Rename test if it exists. test_file_prefix = f"*{migration_meta.number}*" migration_test_path = os.path.join(app_path, "..", "..", "tests", "sentry", "migrations") matches = list(pathlib.Path(migration_test_path).glob(test_file_prefix)) if len(matches) == 1: click.echo("> Updating migration test file to & from attributes") test_path = str(matches[0].resolve()) with open(test_path) as f: contents = f.read() contents = re.sub( r"(migrate_from\s+=\s+\")([^\"]+)(\")", lambda matches: f"{matches.group(1)}{current_head.full_name}{matches.group(3)}", contents, ) contents = re.sub( r"(migrate_to\s+=\s+\")([^\"]+)(\")", lambda matches: f"{matches.group(1)}{new_name}{matches.group(3)}", contents, ) with open(test_path, "w") as f: f.write(contents) click.echo("> Renaming test file") os.rename( str(matches[0].resolve()), os.path.join(migration_test_path, f"test_{new_name}.py") ) else: click.echo("> Could not find a migration test file") click.echo("All done!") if __name__ == "__main__": main()