Fix transaction

nautobot_design_builder/design_job.py

@@ -48,6 +48,7 @@ def __init__(self, *args, **kwargs):
         self.designs = {}
         # TODO: Remove this when we no longer support Nautobot 1.x
         self.rendered = None
+        self.rendered_design = None
         self.failed = False
         self.report = None
 
@@ -119,21 +120,14 @@ def render_design(self, context, design_file):
             context (Context object): a tree of variables that can include templates for values
             design_file (str): Filename of the design file to render.
         """
+        self.rendered_design = design_file
         self.rendered = self.render(context, design_file)
-        # Save the rendered result for later examination from
-        # the job result/additional data tab.
-        output_file = path.basename(design_file)
-        # this should remove the .j2
-        output_file, _ = path.splitext(output_file)
-        if not output_file.endswith(".yaml") and not output_file.endswith(".yml"):
-            output_file = f"{output_file}.yaml"
-        self.save_design_file(output_file, self.rendered)
-
         design = yaml.safe_load(self.rendered)
         self.designs[design_file] = design
 
         # no need to save the rendered content if yaml loaded
         # it okay
+        self.rendered_design = None
         self.rendered = None
         return design
 
@@ -160,8 +154,28 @@ def implement_design(self, context, design_file, commit):
         design = self.render_design(context, design_file)
         self.environment.implement_design(design, commit)
 
-    def run(self, **kwargs):  # pylint: disable=arguments-differ,too-many-branches
+    def run(self, **kwargs):  # pylint: disable=arguments-differ
         """Render the design and implement it within a build Environment object."""
+        try:
+            return self._run_in_transaction(**kwargs)
+        finally:
+            if self.rendered:
+                self.save_design_file(self.rendered_design, self.rendered)
+            for design_file, design in self.designs.items():
+                output_file = path.basename(design_file)
+                # this should remove the .j2
+                output_file, _ = path.splitext(output_file)
+                if not output_file.endswith(".yaml") and not output_file.endswith(".yml"):
+                    output_file = f"{output_file}.yaml"
+                self.save_design_file(output_file, yaml.safe_dump(design))
+
+    @transaction.atomic
+    def _run_in_transaction(self, **kwargs):  # pylint: disable=too-many-branches
+        """Render the design and implement it within a build Environment object.
+
+        This version of `run` is wrapped in a transaction and will roll back database changes
+        on error. In general, this method should only be called by the `run` method.
+        """
         self.log_info(message=f"Building {getattr(self.Meta, 'name')}")
         extensions = getattr(self.Meta, "extensions", [])
         self.environment = Environment(job_result=self.job_result, extensions=extensions)

nautobot_design_builder/tests/__init__.py

@@ -39,6 +39,7 @@ def get_mocked_job(self, design_class: Type[DesignJob]):
         if nautobot_version < "2.0.0":
             job.request = mock.Mock()
         else:
+            # TODO: Remove this when we no longer support Nautobot 1.x
             job.job_result.data = {}
             old_run = job.run
 

nautobot_design_builder/tests/designs/templates/simple_design_3.yaml.j2

@@ -0,0 +1,4 @@
+---
+manufacturers:
+    name: "Test Manufacturer 1"
+    name: "Test Manufacturer"

nautobot_design_builder/tests/designs/test_designs.py

@@ -13,6 +13,14 @@ class Meta:  # pylint: disable=too-few-public-methods
         design_file = "templates/simple_design.yaml.j2"
 
 
+class SimpleDesign3(DesignJob):
+    """Simple design job with extra manufacturer."""
+
+    class Meta:  # pylint: disable=too-few-public-methods
+        name = "Simple Design 3"
+        design_file = "templates/simple_design_3.yaml.j2"
+
+
 class SimpleDesignReport(DesignJob):
     """Simple design job that includes a post-implementation report."""
 

nautobot_design_builder/tests/test_design_job.py

@@ -27,6 +27,19 @@ def test_simple_design_commit(self, environment: Mock):
         )
         environment.return_value.roll_back.assert_not_called()
 
+    def test_simple_design_rollback(self):
+        job1 = self.get_mocked_job(test_designs.SimpleDesign)
+        job1.run(data={}, commit=True)
+        self.assertFalse(job1.failed)
+        self.assertEqual(1, Manufacturer.objects.all().count())
+        job2 = self.get_mocked_job(test_designs.SimpleDesign3)
+        if nautobot_version < "2":
+            job2.run(data={}, commit=True)
+        else:
+            self.assertRaises(DesignValidationError, job2.run, data={}, commit=True)
+        self.assertTrue(job2.failed)
+        self.assertEqual(1, Manufacturer.objects.all().count())
+
     def test_simple_design_report(self):
         job = self.get_mocked_job(test_designs.SimpleDesignReport)
         job.run(data={}, commit=True)