From a971f551645db32c474a5fd151b0da21d8410690 Mon Sep 17 00:00:00 2001
From: Sidney Richards <sidney@maykinmedia.nl>
Date: Wed, 18 Dec 2024 10:43:53 +0100
Subject: [PATCH] [#37] Explicitly handle transactions in the runner

---
 .../commands/setup_configuration.py           |  2 -
 django_setup_configuration/runner.py          | 25 ++++++++-
 tests/test_runner.py                          |  2 +
 tests/test_test_utils.py                      |  2 +
 tests/test_transactions.py                    | 56 +++++++++++++++++++
 5 files changed, 82 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_transactions.py

diff --git a/django_setup_configuration/management/commands/setup_configuration.py b/django_setup_configuration/management/commands/setup_configuration.py
index 118c23b..fb4f918 100644
--- a/django_setup_configuration/management/commands/setup_configuration.py
+++ b/django_setup_configuration/management/commands/setup_configuration.py
@@ -1,7 +1,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
@@ -33,7 +32,6 @@ def add_arguments(self, parser):
             help="Path to YAML file containing the configurations",
         )
 
-    @transaction.atomic
     def handle(self, **options):
         yaml_file = Path(options["yaml_file"]).resolve()
         if not yaml_file.exists():
diff --git a/django_setup_configuration/runner.py b/django_setup_configuration/runner.py
index 9f6e7e1..9bb9ddf 100644
--- a/django_setup_configuration/runner.py
+++ b/django_setup_configuration/runner.py
@@ -6,6 +6,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
@@ -157,7 +158,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:
@@ -207,8 +209,25 @@ 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)
+                    yield result
+
+                if any(result.run_exception for result in results):
+                    raise Rollback  # Trigger the rollback
+        except Rollback:
+            pass
 
     def execute_all(self) -> list[StepExecutionResult]:
         """
diff --git a/tests/test_runner.py b/tests/test_runner.py
index 40eb93c..402215d 100644
--- a/tests/test_runner.py
+++ b/tests/test_runner.py
@@ -12,6 +12,8 @@
 from django_setup_configuration.test_utils import execute_single_step
 from tests.conftest import TestStep, TestStepConfig
 
+pytestmark = pytest.mark.django_db
+
 
 def test_runner_raises_on_non_existent_step_module_path(test_step_yaml_path):
     with pytest.raises(ConfigurationException):
diff --git a/tests/test_test_utils.py b/tests/test_test_utils.py
index 8efb192..1f85978 100644
--- a/tests/test_test_utils.py
+++ b/tests/test_test_utils.py
@@ -4,6 +4,8 @@
 from django_setup_configuration.test_utils import execute_single_step
 from tests.conftest import TestStep
 
+pytestmark = pytest.mark.django_db
+
 
 def test_exception_during_execute_step_is_immediately_raised(
     step_execute_mock,
diff --git a/tests/test_transactions.py b/tests/test_transactions.py
new file mode 100644
index 0000000..5aa81d2
--- /dev/null
+++ b/tests/test_transactions.py
@@ -0,0 +1,56 @@
+from django.contrib.auth.models import User
+
+import pytest
+
+from testapp.configuration import UserConfigurationStep
+from tests.conftest import TestStep
+
+pytestmark = pytest.mark.django_db
+
+
+@pytest.fixture()
+def yaml_file_with_valid_configuration(yaml_file_factory, test_step_valid_config):
+    yaml_path = yaml_file_factory(
+        {
+            "user_configuration_enabled": True,
+            "user_configuration": {
+                "username": "demo",
+                "password": "secret",
+            },
+            "some_extra_attrs": "should be allowed",
+        }
+        | test_step_valid_config
+    )
+
+    return yaml_path
+
+
+def test_runner_rolls_back_all_on_failing_step(
+    runner_factory, yaml_file_with_valid_configuration, step_execute_mock
+):
+    runner = runner_factory(
+        steps=[UserConfigurationStep, TestStep],
+        yaml_source=yaml_file_with_valid_configuration,
+    )
+    exc = Exception()
+    step_execute_mock.side_effect = exc
+
+    user_configuration_step_result, test_step_result = runner.execute_all()
+
+    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
+
+    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