reversible_migrate.py

from django.conf import settings
from django.core.management import CommandError
from django.core.management.commands.migrate import Command as CoreCommand
from django.db import connections
from django.db.migrations.executor import MigrationExecutor
from django.db.migrations.recorder import MigrationRecorder


class Command(CoreCommand):
    def handle(self, *args, **options):
        db = options['database']
        connection = connections[db]
        old_state = self.get_current_migrations(connection)
        executor = MigrationExecutor(connection, self.migration_progress_callback)
        targets = executor.loader.graph.leaf_nodes()
        migration_plan = executor.migration_plan(list(targets))

        for (migration, is_backwards) in migration_plan:
            for operation in migration.operations:
                if not operation.reversible:
                    raise CommandError("Operation {} at {} is not reversible!".format(operation, migration))

        try:
            super().handle(*args, **options)
        except Exception:
            self.stdout.write(self.style.ERROR(" FAILURE "))
            self.stdout.write(self.style.WARNING("Failure on migrations. Reversing..."))
            self.reverse_migrations(connection, old_state)
            self.stdout.write(self.style.WARNING("Migrations reversed."))
            self.stdout.write("Error on migration:")
            raise

    def reverse_migrations(self, connection, old_state):
        MigrationExecutor(connection, self.migration_progress_callback).migrate(old_state)

    def get_current_migrations(self, connection):
        """
        Show a list of all migrations on the system, or only those of
        some named apps.
        """
        migrations = MigrationRecorder(connection).migration_qs.order_by('app', '-id').values_list(
            'app', 'name').distinct('app')

        targets = []

        for app_label, migration_name in migrations:
            if app_label in settings.INSTALLED_APPS:
                targets.append((app_label, migration_name))

        return targets