Merge pull request #38 from maykinmedia/runner-managed-db-transactions

[#37] Explicitly handle transactions in the runner

django_setup_configuration/management/commands/setup_configuration.py

@@ -3,7 +3,6 @@
 from pathlib import Path
 
 from django.core.management import BaseCommand, CommandError
-from django.db import transaction
 
 from django_setup_configuration.exceptions import ValidateRequirementsFailure
 from django_setup_configuration.runner import SetupConfigurationRunner
@@ -34,7 +33,6 @@ def add_arguments(self, parser):
             "from source, without actually executing the steps.",
         )
 
-    @transaction.atomic
     def handle(self, **options):
         validate_only = options["validate_only"]
         yaml_file = Path(options["yaml_file"]).resolve()

django_setup_configuration/runner.py

@@ -7,6 +7,7 @@
 from typing import Any, Generator
 
 from django.conf import settings
+from django.db import transaction
 from django.utils.module_loading import import_string
 
 from pydantic import ValidationError
@@ -160,7 +161,8 @@ def _execute_step(
         step_exc = None
 
         try:
-            step.execute(config_model)
+            with transaction.atomic():
+                step.execute(config_model)
         except BaseException as exc:
             step_exc = exc
         finally:
@@ -216,8 +218,27 @@ def execute_all_iter(self) -> Generator[StepExecutionResult, Any, None]:
             Generator[StepExecutionResult, Any, None]: The results of each step's
                 execution.
         """
-        for step in self.enabled_steps:
-            yield self._execute_step(step)
+
+        # Not the most elegant approach to rollbacks, but it's preferable to the
+        # pitfalls of manual transaction management. We want all steps to run and only
+        # rollback at the end, hence intra-step exceptions are caught and persisted.
+        class Rollback(BaseException):
+            pass
+
+        try:
+            with transaction.atomic():
+                results = []
+                for step in self.enabled_steps:
+                    result = self._execute_step(step)
+                    results.append(result)
+
+                if any(result.run_exception for result in results):
+                    raise Rollback  # Trigger the rollback
+        except Rollback:
+            pass
+
+        for result in results:
+            yield result
 
     def execute_all(self) -> list[StepExecutionResult]:
         """

tests/test_command.py

@@ -482,3 +482,50 @@ def test_command_with_disabled_and_enabled_steps_lists_the_disabled_steps(
 
     assert User.objects.count() == 1
     step_execute_mock.assert_not_called()
+
+
+@pytest.fixture()
+def valid_config_object(test_step_valid_config):
+    return {
+        "transaction_test_configuration_enabled": True,
+        "transaction_test_configuration": {"username": "alice"},
+    } | test_step_valid_config
+
+
+def test_command_rolls_back_all_on_failing_step(
+    yaml_file_with_valid_configuration, step_execute_mock
+):
+    exc = Exception()
+    step_execute_mock.side_effect = exc
+
+    assert User.objects.count() == 0
+    stdout, stderr = StringIO(), StringIO()
+
+    with pytest.raises(CommandError) as excinfo:
+        call_command(
+            "setup_configuration",
+            yaml_file=yaml_file_with_valid_configuration,
+            stdout=stdout,
+            stderr=stderr,
+        )
+
+    assert (
+        str(excinfo.value)
+        == "Aborting run due to a failed step. All database changes have been rolled back."
+    )
+    step_execute_mock.assert_called_once()
+
+    # Initial run is rolled back, so no objects created
+    assert User.objects.count() == 0
+
+    # Subsequent run does not raise, so the objects are created
+    step_execute_mock.side_effect = None
+
+    call_command(
+        "setup_configuration",
+        yaml_file=yaml_file_with_valid_configuration,
+        stdout=stdout,
+        stderr=stderr,
+    )
+    assert User.objects.count() == 1
+    assert step_execute_mock.call_count == 2

tests/test_runner.py

@@ -14,6 +14,8 @@
 from django_setup_configuration.test_utils import execute_single_step
 from tests.conftest import ConfigModel, ConfigStep
 
+pytestmark = pytest.mark.django_db
+
 
 def test_runner_raises_on_non_existent_step_module_path(test_step_yaml_path):
     with pytest.raises(ConfigurationException):

tests/test_test_utils.py

@@ -4,6 +4,8 @@
 from django_setup_configuration.test_utils import execute_single_step
 from tests.conftest import ConfigStep
 
+pytestmark = pytest.mark.django_db
+
 
 def test_exception_during_execute_step_is_immediately_raised(
     step_execute_mock,

tests/test_transactions.py

@@ -0,0 +1,105 @@
+from unittest import mock
+
+from django.contrib.auth.models import User
+
+import pytest
+
+from django_setup_configuration.configuration import BaseConfigurationStep
+from django_setup_configuration.models import ConfigurationModel
+from tests.conftest import TestStep
+
+pytestmark = pytest.mark.django_db
+
+
+def side_effect_test_func():
+    pass
+
+
+class TransactionTestConfigurationModel(ConfigurationModel):
+
+    username: str
+
+
+class TransactionTestConfigurationStep(
+    BaseConfigurationStep[TransactionTestConfigurationModel]
+):
+
+    config_model = TransactionTestConfigurationModel
+    enable_setting = "transaction_test_configuration_enabled"
+    namespace = "transaction_test_configuration"
+    verbose_name = "Transaction Test Configuration"
+
+    def execute(self, model) -> None:
+        User.objects.create_user(
+            username=model.username,
+            password="secret",
+        )
+
+        side_effect_test_func()
+
+
+@pytest.fixture()
+def valid_config_object(test_step_valid_config):
+    return {
+        "transaction_test_configuration_enabled": True,
+        "transaction_test_configuration": {"username": "alice"},
+    } | test_step_valid_config
+
+
+def test_runner_rolls_back_all_on_failing_step(
+    runner_factory, valid_config_object, step_execute_mock
+):
+    runner = runner_factory(
+        steps=[TransactionTestConfigurationStep, TestStep],
+        object_source=valid_config_object,
+    )
+    exc = Exception()
+    step_execute_mock.side_effect = exc
+
+    user_configuration_step_result, test_step_result = runner.execute_all()
+
+    # Initial run is rolled back, so no objects created
+    assert test_step_result.has_run
+    assert test_step_result.run_exception is exc
+
+    assert user_configuration_step_result.has_run
+    assert user_configuration_step_result.run_exception is None
+    assert User.objects.count() == 0
+
+    # Subsequent run does not raise, so the objects are created
+    step_execute_mock.side_effect = None
+
+    user_configuration_step_result, test_step_result = runner.execute_all()
+
+    assert test_step_result.has_run
+    assert test_step_result.run_exception is None
+
+    assert user_configuration_step_result.has_run
+    assert user_configuration_step_result.run_exception is None
+    assert User.objects.count() == 1
+
+
+def test_runner_rolls_back_on_executing_single_step(
+    runner_factory, valid_config_object
+):
+    runner = runner_factory(
+        steps=[TransactionTestConfigurationStep, TestStep],
+        object_source=valid_config_object,
+    )
+    with mock.patch("tests.test_transactions.side_effect_test_func") as m:
+        exc = Exception()
+        m.side_effect = exc
+
+        user_configuration_step_result = runner._execute_step(
+            runner.configured_steps[0]
+        )
+
+        assert user_configuration_step_result.has_run
+        assert user_configuration_step_result.run_exception is exc
+        assert User.objects.count() == 0
+
+    user_configuration_step_result = runner._execute_step(runner.configured_steps[0])
+
+    assert user_configuration_step_result.has_run
+    assert user_configuration_step_result.run_exception is None
+    assert User.objects.count() == 1