Acquire lock on singleton to handle parallel pod deployments (#10)

Co-authored-by: Ronny V <[email protected]>

CHANGES.md

@@ -2,6 +2,7 @@
 
 * *master* (unreleased)
   * Use ORM to reset `django_migrations` table
+  * Lock rows to enable parallel deployments
   * Dropped Python 3.8 support
   * Added multiple ruff linters
   * Updated GitHub actions

django_migration_zero/managers.py

@@ -1,35 +1,26 @@
+from django.core.exceptions import MultipleObjectsReturned, ObjectDoesNotExist
 from django.db import ProgrammingError, models
 
 from django_migration_zero.exceptions import MissingMigrationZeroConfigRecordError
 from django_migration_zero.helpers.logger import get_logger
 
 
-class MigrationZeroConfigurationQuerySet(models.QuerySet):
-    pass
-
-
 class MigrationZeroConfigurationManager(models.Manager):
     def fetch_singleton(self) -> None:
         logger = get_logger()
         try:
-            number_records = self.count()
+            config_singleton = self.select_for_update().get()
         except ProgrammingError:
             logger.warning(
                 "The migration zero table is missing. This might be ok for the first installation of "
                 '"django-migration-zero" but if you see this warning after that point, something went sideways.'
             )
-            return None
-
-        if number_records > 1:
+            config_singleton = None
+        except MultipleObjectsReturned as e:
             raise MissingMigrationZeroConfigRecordError(
                 "Too many configuration records detected. There can only be one."
-            )
-
-        config_singleton = self.all().first()
-        if not config_singleton:
-            raise MissingMigrationZeroConfigRecordError("No configuration record found in the database.")
+            ) from e
+        except ObjectDoesNotExist as e:
+            raise MissingMigrationZeroConfigRecordError("No configuration record found in the database.") from e
 
         return config_singleton
-
-
-MigrationZeroConfigurationManager = MigrationZeroConfigurationManager.from_queryset(MigrationZeroConfigurationQuerySet)

django_migration_zero/services/deployment.py

@@ -1,6 +1,7 @@
 from logging import Logger
 
 from django.core.management import call_command
+from django.db import transaction
 from django.db.migrations.recorder import MigrationRecorder
 
 from django_migration_zero.exceptions import InvalidMigrationTreeError
@@ -20,6 +21,7 @@ def __init__(self):
 
         self.logger = get_logger()
 
+    @transaction.atomic
     def process(self):
         self.logger.info("Starting migration zero database adjustments...")
 

tests/services/test_deployment.py

@@ -1,7 +1,9 @@
 from logging import Logger
+from threading import Thread
 from unittest import mock
+from unittest.mock import Mock, call
 
-from django.test import TestCase
+from django.test import TestCase, TransactionTestCase
 from django.utils import timezone
 from freezegun import freeze_time
 
@@ -81,3 +83,56 @@ def test_process_validate_all_commands_are_executed(self, mocked_call_command):
 
         self.assertEqual(calls[1].args, ("migrate",))
         self.assertEqual(calls[1].kwargs, {"check": True})
+
+
+class DatabasePreparationServiceTestParallelDeployment(TransactionTestCase):
+    @mock.patch("django_migration_zero.services.deployment.get_logger")
+    def test_process_multiple_threads(self, mock_get_logger):
+        """Test parallel deployments to multiple pods."""
+        # Setup
+        mock_logger_info = Mock(return_value=None)
+        mock_get_logger.return_value = Mock(info=mock_logger_info)
+        config, _ = MigrationZeroConfiguration.objects.update_or_create(
+            defaults={
+                "migration_imminent": True,
+                "migration_date": timezone.now().date(),
+            }
+        )
+
+        # Testable
+        number_of_pods = 1
+        threads = [Thread(target=DatabasePreparationService().process) for _ in range(number_of_pods)]
+
+        for thread in threads:
+            thread.start()
+
+        for thread in threads:
+            thread.join()
+
+        # Assertions
+        self.assertEqual(mock_logger_info.call_count, 6 + number_of_pods, msg=mock_logger_info.call_args_list)
+        mock_logger_info.assert_has_calls(
+            [
+                call("Starting migration zero database adjustments..."),
+                call("Resetting migration history for all apps..."),
+                call("Populating migration history."),
+                call("Checking migration integrity."),
+                call("All good."),
+                call("Deactivating migration zero switch in database."),
+                call("Process successfully finished."),
+            ],
+            any_order=True,
+        )
+        self.assertEqual(
+            len(
+                [
+                    mock_call
+                    for mock_call in mock_logger_info.call_args_list
+                    if mock_call == call("Starting migration zero database adjustments...")
+                ]
+            ),
+            number_of_pods,
+        )
+
+        config.refresh_from_db()
+        self.assertFalse(config.migration_imminent)

tests/test_managers.py

@@ -5,6 +5,7 @@
 from django.test import TestCase
 
 from django_migration_zero.exceptions import MissingMigrationZeroConfigRecordError
+from django_migration_zero.managers import MigrationZeroConfigurationManager
 from django_migration_zero.models import MigrationZeroConfiguration
 
 
@@ -18,7 +19,7 @@ def test_fetch_singleton_regular(self):
     def test_fetch_singleton_singleton_exists_via_migration(self):
         self.assertEqual(MigrationZeroConfiguration.objects.all().count(), 1)
 
-    @mock.patch.object(MigrationZeroConfiguration.objects, "count", side_effect=ProgrammingError)
+    @mock.patch.object(MigrationZeroConfigurationManager, "select_for_update", side_effect=ProgrammingError)
     def test_fetch_singleton_database_error(self, *args):
         self.assertIsNone(MigrationZeroConfiguration.objects.fetch_singleton())