feat: use patched django class to replace all migrations

Author: madtibo

django_replace_migrations/management/commands/makemigrations.py

@@ -22,6 +22,10 @@
 from django.db.migrations.utils import get_migration_name_timestamp
 from django.db.migrations.writer import MigrationWriter
 
+from .replace_migration_autodetector import ReplaceMigrationAutodetector
+from .replace_migration_loader import ReplaceMigrationLoader
+from .replace_migration_writer import ReplaceMigrationWriter
+
 
 class Command(BaseCommand):
     help = "Creates new migration(s) for apps."
@@ -143,7 +147,10 @@ def handle(self, *app_labels, **options):
 
         # Load the current graph state. Pass in None for the connection so
         # the loader doesn't try to resolve replaced migrations from DB.
-        loader = MigrationLoader(None, ignore_no_migrations=True)
+        if self.replace_all:
+            loader = ReplaceMigrationLoader(None, ignore_no_migrations=True)
+        else:
+            loader = MigrationLoader(None, ignore_no_migrations=True)
 
         # Raise an error if any migrations are applied before their dependencies.
         consistency_check_labels = {config.label for config in apps.get_app_configs()}
@@ -224,19 +231,19 @@ def handle(self, *app_labels, **options):
                 k: v for (k, v) in loader.graph.nodes.items() if k[0] not in app_labels
             }
 
-            autodetector = MigrationAutodetector(
-                loader.project_state(),
-                ProjectState.from_apps(apps),
-                questioner,
+            autodetector = ReplaceMigrationAutodetector(
+                from_state=loader.project_state(),
+                to_state=ProjectState.from_apps(apps),
+                questioner=questioner,
             )
 
             loader.graph.nodes = temp_nodes
 
         else:
             autodetector = MigrationAutodetector(
-                loader.project_state(),
-                ProjectState.from_apps(apps),
-                questioner,
+                from_state=loader.project_state(),
+                to_state=ProjectState.from_apps(apps),
+                questioner=questioner,
             )
 
         # If they want to make an empty migration, make one for each app
@@ -378,7 +385,7 @@ def write_migration_files(self, changes, update_previous_migration_paths=None):
                 self.log(self.style.MIGRATE_HEADING("Migrations for '%s':" % app_label))
             for migration in app_migrations:
                 # Describe the migration
-                writer = MigrationWriter(migration, self.include_header)
+                writer = ReplaceMigrationWriter(migration, self.include_header)
                 if self.verbosity >= 1:
                     # Display a relative path if it's below the current working
                     # directory, or an absolute path otherwise.

django_replace_migrations/management/commands/replace_migration_autodetector.py

@@ -0,0 +1,16 @@
+import re
+
+from django.db.migrations.autodetector import MigrationAutodetector
+
+
+class ReplaceMigrationAutodetector(MigrationAutodetector):
+    @classmethod
+    def parse_number(cls, name):
+        """
+        Given a migration name, try to extract a number from the beginning of
+        it. For a squashed migration such as '0001_squashed_0004…', return the
+        second number. If no number is found, return None.
+        """
+        if squashed_match := re.search(r"(\d+)_squashed_.*", name):
+            return int(squashed_match[1])
+        return None

django_replace_migrations/management/commands/replace_migration_graph.py

@@ -0,0 +1,26 @@
+from django.db.migrations.graph import MigrationGraph
+from django.db.migrations.state import ProjectState
+
+
+class ReplaceMigrationGraph(MigrationGraph):
+    def make_state(self, nodes=None, at_end=True, real_apps=None):
+        """
+        Given a migration node or nodes, return a complete ProjectState for it.
+        If at_end is False, return the state before the migration has run.
+        If nodes is not provided, return the overall most current project state.
+        """
+        if nodes is None:
+            nodes = list(self.leaf_nodes())
+        if not nodes:
+            return ProjectState()
+        if not isinstance(nodes[0], tuple):
+            nodes = [nodes]
+        plan = self._generate_plan(nodes, at_end)
+        project_state = ProjectState(real_apps=real_apps)
+        for node in plan:
+            # We have dependencies between the contrib and our migrations
+            # if a node is not found, do not link / use it
+            if node not in self.nodes:
+                continue
+            project_state = self.nodes[node].mutate_state(project_state, preserve=False)
+        return project_state

django_replace_migrations/management/commands/replace_migration_loader.py

@@ -0,0 +1,94 @@
+from django.db.migrations.exceptions import NodeNotFoundError
+from django.db.migrations.loader import MigrationLoader
+from django.db.migrations.recorder import MigrationRecorder
+
+from .replace_migration_graph import ReplaceMigrationGraph
+
+
+class ReplaceMigrationLoader(MigrationLoader):
+    def build_graph(self):
+        """
+        Build a migration dependency graph using both the disk and database.
+        You'll need to rebuild the graph if you apply migrations. This isn't
+        usually a problem as generally migration stuff runs in a one-shot process.
+        """
+        # Load disk data
+        self.load_disk()
+        # Load database data
+        if self.connection is None:
+            self.applied_migrations = {}
+        else:
+            recorder = MigrationRecorder(self.connection)
+            self.applied_migrations = recorder.applied_migrations()
+        # To start, populate the migration graph with nodes for ALL migrations
+        # and their dependencies. Also make note of replacing migrations at this step.
+        self.graph = ReplaceMigrationGraph()
+        self.replacements = {}
+        for key, migration in self.disk_migrations.items():
+            self.graph.add_node(key, migration)
+            # Replacing migrations.
+            if migration.replaces:
+                self.replacements[key] = migration
+        for key, migration in self.disk_migrations.items():
+            # Internal (same app) dependencies.
+            self.add_internal_dependencies(key, migration)
+        # Add external dependencies now that the internal ones have been resolved.
+        for key, migration in self.disk_migrations.items():
+            self.add_external_dependencies(key, migration)
+        # Carry out replacements where possible and if enabled.
+        if self.replace_migrations:
+            for key, migration in self.replacements.items():
+                # Get applied status of each of this migration's replacement
+                # targets.
+                applied_statuses = [
+                    (target in self.applied_migrations) for target in migration.replaces
+                ]
+                # The replacing migration is only marked as applied if all of
+                # its replacement targets are.
+                if all(applied_statuses):
+                    self.applied_migrations[key] = migration
+                else:
+                    self.applied_migrations.pop(key, None)
+                # A replacing migration can be used if either all or none of
+                # its replacement targets have been applied.
+                if all(applied_statuses) or (not any(applied_statuses)):
+                    self.graph.remove_replaced_nodes(key, migration.replaces)
+                else:
+                    # This replacing migration cannot be used because it is
+                    # partially applied. Remove it from the graph and remap
+                    # dependencies to it (#25945).
+                    self.graph.remove_replacement_node(key, migration.replaces)
+        # Ensure the graph is consistent.
+        try:
+            self.graph.validate_consistency()
+        except NodeNotFoundError as exc:
+            # Check if the missing node could have been replaced by any squash
+            # migration but wasn't because the squash migration was partially
+            # applied before. In that case raise a more understandable exception
+            # (#23556).
+            # Get reverse replacements.
+            reverse_replacements = {}
+            for key, migration in self.replacements.items():
+                for replaced in migration.replaces:
+                    reverse_replacements.setdefault(replaced, set()).add(key)
+            # Try to reraise exception with more detail.
+            if exc.node in reverse_replacements:
+                candidates = reverse_replacements.get(exc.node, set())
+                is_replaced = any(
+                    candidate in self.graph.nodes for candidate in candidates
+                )
+                if not is_replaced:
+                    tries = ", ".join(
+                        f"{key}.{migration}" for key, migration in candidates
+                    )
+                    raise NodeNotFoundError(
+                        "Migration {0} depends on nonexistent node ('{1}', '{2}'). "
+                        "Django tried to replace migration {1}.{2} with any of [{3}] "
+                        "but wasn't able to because some of the replaced migrations "
+                        "are already applied.".format(
+                            exc.origin, exc.node[0], exc.node[1], tries
+                        ),
+                        exc.node,
+                    ) from exc
+            raise
+        self.graph.ensure_not_cyclic()

django_replace_migrations/management/commands/replace_migration_writer.py

@@ -0,0 +1,149 @@
+import re
+
+from django import get_version
+from django.db.migrations.operations.special import (
+    RunPython,
+    RunSQL,
+    SeparateDatabaseAndState,
+)
+from django.db.migrations.writer import MigrationWriter, OperationWriter
+from django.utils.timezone import now
+
+
+class ReplaceMigrationWriter(MigrationWriter):
+    def as_string(self):
+        """Return a string of the file contents."""
+        items = {
+            "replaces_str": "",
+            "initial_str": "",
+        }
+
+        imports = set()
+
+        # Deconstruct operations
+        operations = []
+        for operation in self.migration.operations:
+            if isinstance(operation, (RunPython, RunSQL, SeparateDatabaseAndState)):
+                continue
+            operation_string, operation_imports = OperationWriter(operation).serialize()
+            imports.update(operation_imports)
+            operations.append(operation_string)
+        items["operations"] = "\n".join(operations) + "\n" if operations else ""
+
+        # Deconstruct special_operations
+        special_operations = []
+        for special_op in self.migration.operations:
+            if (
+                isinstance(special_op, (RunPython, RunSQL, SeparateDatabaseAndState))
+                and not special_op.elidable
+            ):
+                operation_string, operation_imports = OperationWriter(
+                    special_op
+                ).serialize()
+                imports.update(operation_imports)
+                special_operations.append(operation_string)
+        special_ops = (
+            "\n".join(special_operations) + "\n" if special_operations else None
+        )
+        items["special_operations"] = (
+            "\n    # /!\\ PRINT ALL THE SPECIAL OPERATIONS\n"
+            + "    # /!\\ MUST BE MANUALLY REVIEWED\n\n"
+            + "   special_operations = [\n"
+            + special_ops
+            + "   ]\n"
+            if special_ops
+            else ""
+        )
+
+        # Format dependencies and write out swappable dependencies right
+        dependencies = []
+        for dependency in self.migration.dependencies:
+            if dependency[0] == "__setting__":
+                dependencies.append(
+                    f"        migrations.swappable_dependency(settings.{dependency[1]}),"
+                )
+                imports.add("from django.conf import settings")
+            else:
+                dependencies.append(f"        {self.serialize(dependency)[0]},")
+        items["dependencies"] = (
+            "\n".join(sorted(dependencies)) + "\n" if dependencies else ""
+        )
+
+        # Format imports nicely, swapping imports of functions from migration files
+        # for comments
+        migration_imports = set()
+        for line in list(imports):
+            if re.match(r"^import (.*)\.\d+[^\s]*$", line):
+                migration_imports.add(line.split("import")[1].strip())
+                imports.remove(line)
+                self.needs_manual_porting = True
+
+        # django.db.migrations is always used, but models import may not be.
+        # If models import exists, merge it with migrations import.
+        if "from django.db import models" in imports:
+            imports.discard("from django.db import models")
+            imports.add("from django.db import migrations, models")
+        else:
+            imports.add("from django.db import migrations")
+
+        # Sort imports by the package / module to be imported (the part after
+        # "from" in "from ... import ..." or after "import" in "import ...").
+        sorted_imports = sorted(imports, key=lambda i: i.split()[1])
+        items["imports"] = "\n".join(sorted_imports) + "\n" if imports else ""
+        if migration_imports:
+            items["imports"] += (
+                "\n\n# Functions from the following migrations need manual "
+                "copying.\n# Move them and any dependencies into this file, "
+                "then update the\n# RunPython operations to refer to the local "
+                "versions:\n# %s"
+            ) % "\n# ".join(sorted(migration_imports))
+        # If there's a replaces, make a string for it
+        if self.migration.replaces:
+            items[
+                "replaces_str"
+            ] = f"\n    replaces = {self.serialize(sorted(self.migration.replaces))[0]}\n"
+        # Hinting that goes into comment
+        if self.include_header:
+            items["migration_header"] = MIGRATION_HEADER_TEMPLATE % {
+                "version": get_version(),
+                "timestamp": now().strftime("%Y-%m-%d %H:%M"),
+            }
+        else:
+            items["migration_header"] = ""
+
+        if self.migration.initial:
+            items["initial_str"] = "\n    initial = True\n"
+
+        return MIGRATION_TEMPLATE % items
+
+
+MIGRATION_HEADER_TEMPLATE = """\
+# Generated by Django %(version)s on %(timestamp)s
+
+"""
+
+
+MIGRATION_TEMPLATE = """\
+%(migration_header)s%(imports)s
+
+from phased_migrations.constants import DeployPhase
+
+
+class Migration(migrations.Migration):
+    # Note: deploy_phase was added to ensure consistency with no down time
+    # it is possible that this migration in not really compatible with pre-deploy
+    deploy_phase = DeployPhase.pre_deploy
+
+    squashed_with_gg_script = True
+
+%(replaces_str)s%(initial_str)s
+    dependencies = [
+%(dependencies)s\
+    ]
+
+    operations = [
+%(operations)s\
+    ]
+
+%(special_operations)s
+"""