Merge pull request #6 from charettes/remove-pre_migrate-hack

Stop relying on the pre_migrate signal to alter migration plan.
charettes · Apr 21, 2020 · 462b8b3 · 462b8b3
2 parents 58d086e + 902bda6
commit 462b8b3
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 34 deletions.
diff --git a/syzygy/management/commands/migrate.py b/syzygy/management/commands/migrate.py
@@ -1,14 +1,22 @@
-from django.apps import apps
 from django.core.management import CommandError
-from django.core.management.commands.migrate import (  # type: ignore
-    Command as MigrateCommand,
-)
-from django.db.models.signals import pre_migrate
+from django.core.management.commands import migrate  # type: ignore
+from django.db.migrations.executor import MigrationExecutor
 
 from syzygy.plan import get_pre_deploy_plan
 
 
-class Command(MigrateCommand):
+class PreDeployMigrationExecutor(MigrationExecutor):
+    def migration_plan(self, targets, clean_start=False):
+        plan = super().migration_plan(targets, clean_start=clean_start)
+        if not clean_start:
+            try:
+                plan = get_pre_deploy_plan(plan)
+            except ValueError as exc:
+                raise CommandError(str(exc)) from exc
+        return plan
+
+
+class Command(migrate.Command):
     def add_arguments(self, parser):
         super().add_arguments(parser)
         parser.add_argument(
@@ -17,13 +25,12 @@ def add_arguments(self, parser):
             help="Only run migrations staged for pre-deployment.",
         )
 
-    def migrate_pre_deploy(self, plan, **kwargs):
-        try:
-            plan[:] = get_pre_deploy_plan(plan)
-        except ValueError as exc:
-            raise CommandError(str(exc)) from exc
-
     def handle(self, *args, pre_deploy, **options):
         if pre_deploy:
-            pre_migrate.connect(self.migrate_pre_deploy, sender=apps.get_app_config('syzygy'))
-        super().handle(*args, **options)
+            # Monkey-patch migrate.MigrationExecutor since the command doesn't
+            # allow it to be overridden in any other way.
+            migrate.MigrationExecutor = PreDeployMigrationExecutor
+        try:
+            super().handle(*args, **options)
+        finally:
+            migrate.MigrationExecutor = MigrationExecutor
diff --git a/syzygy/models.py b/syzygy/models.py
diff --git a/tests/test_commands.py b/tests/test_commands.py
@@ -1,10 +1,15 @@
+from io import StringIO
+
 from django.core.management import CommandError, call_command
 from django.db import connection
 from django.db.migrations.recorder import MigrationRecorder
-from django.test import TransactionTestCase
+from django.test import TransactionTestCase, override_settings
 
 
 class MigrateTests(TransactionTestCase):
+    def tearDown(self):
+        MigrationRecorder(connection).flush()
+
     def get_applied_migrations(self):
         return {
             name
@@ -15,27 +20,29 @@ def get_applied_migrations(self):
     def assert_not_applied(self, name):
         self.assertIn(("tests", name), self.recorder.applied_migrations())
 
+    @override_settings(MIGRATION_MODULES={"tests": "tests.test_migrations.functional"})
     def test_pre_deploy_forward(self):
-        with self.settings(
-            MIGRATION_MODULES={"tests": "tests.test_migrations.functional"}
-        ):
-            call_command("migrate", "tests", pre_deploy=True, verbosity=0)
+        stdout = StringIO()
+        call_command("migrate", "tests", plan=True, no_color=True, pre_deploy=True, stdout=stdout)
+        self.assertIn('tests.0001_pre_deploy', stdout.getvalue())
+        self.assertNotIn('tests.0002_post_deploy', stdout.getvalue())
+        call_command("migrate", "tests", pre_deploy=True, verbosity=0)
         self.assertEqual(self.get_applied_migrations(), {"0001_pre_deploy"})
 
+    @override_settings(MIGRATION_MODULES={"tests": "tests.test_migrations.functional"})
     def test_pre_deploy_backward(self):
-        with self.settings(
-            MIGRATION_MODULES={"tests": "tests.test_migrations.functional"}
-        ):
-            call_command("migrate", "tests", verbosity=0)
-            self.assertEqual(
-                self.get_applied_migrations(), {"0001_pre_deploy", "0002_post_deploy"}
-            )
-            call_command("migrate", "tests", "zero", pre_deploy=True, verbosity=0)
-            self.assertEqual(self.get_applied_migrations(), {"0001_pre_deploy"})
+        call_command("migrate", "tests", verbosity=0)
+        self.assertEqual(
+            self.get_applied_migrations(), {"0001_pre_deploy", "0002_post_deploy"}
+        )
+        stdout = StringIO()
+        call_command("migrate", "tests", "zero", plan=True, no_color=True, pre_deploy=True, stdout=stdout)
+        self.assertIn('tests.0002_post_deploy', stdout.getvalue())
+        self.assertNotIn('tests.0001_pre_deploy', stdout.getvalue())
+        call_command("migrate", "tests", "zero", pre_deploy=True, verbosity=0)
+        self.assertEqual(self.get_applied_migrations(), {"0001_pre_deploy"})
 
+    @override_settings(MIGRATION_MODULES={"tests": "tests.test_migrations.ambiguous"})
     def test_ambiguous(self):
-        with self.settings(
-            MIGRATION_MODULES={"tests": "tests.test_migrations.ambiguous"}
-        ):
-            with self.assertRaisesMessage(CommandError, 'Cannot automatically determine stage of tests.0001_initial.'):
-                call_command("migrate", "tests", pre_deploy=True, verbosity=0)
+        with self.assertRaisesMessage(CommandError, 'Cannot automatically determine stage of tests.0001_initial.'):
+            call_command("migrate", "tests", pre_deploy=True, verbosity=0)