Replace JobResult.set_status() with terminate()

netbox/core/jobs.py

@@ -22,8 +22,9 @@ def sync_datasource(job_result, *args, **kwargs):
         # Update the search cache for DataFiles belonging to this source
         search_backend.cache(datasource.datafiles.iterator())
 
+        job_result.terminate()
+
     except SyncError as e:
-        job_result.set_status(JobResultStatusChoices.STATUS_ERRORED)
-        job_result.save()
+        job_result.terminate(status=JobResultStatusChoices.STATUS_ERRORED)
         DataSource.objects.filter(pk=datasource.pk).update(status=DataSourceStatusChoices.FAILED)
         logging.error(e)

netbox/extras/management/commands/runscript.py

@@ -41,16 +41,16 @@ class Command(BaseCommand):
             the change_logging context manager (which is bypassed if commit == False).
             """
             try:
-                with transaction.atomic():
-                    script.output = script.run(data=data, commit=commit)
-                    job_result.set_status(JobResultStatusChoices.STATUS_COMPLETED)
-
-                    if not commit:
-                        raise AbortTransaction()
-
-            except AbortTransaction:
-                script.log_info("Database changes have been reverted automatically.")
-                clear_webhooks.send(request)
+                try:
+                    with transaction.atomic():
+                        script.output = script.run(data=data, commit=commit)
+                        if not commit:
+                            raise AbortTransaction()
+                except AbortTransaction:
+                    script.log_info("Database changes have been reverted automatically.")
+                    clear_webhooks.send(request)
+                job_result.data = ScriptOutputSerializer(script).data
+                job_result.terminate()
             except Exception as e:
                 stacktrace = traceback.format_exc()
                 script.log_failure(
@@ -58,11 +58,9 @@ class Command(BaseCommand):
                 )
                 script.log_info("Database changes have been reverted due to error.")
                 logger.error(f"Exception raised during script execution: {e}")
-                job_result.set_status(JobResultStatusChoices.STATUS_ERRORED)
                 clear_webhooks.send(request)
-            finally:
                 job_result.data = ScriptOutputSerializer(script).data
-                job_result.save()
+                job_result.terminate(status=JobResultStatusChoices.STATUS_ERRORED)
 
             logger.info(f"Script completed in {job_result.duration}")
 

netbox/extras/models/models.py

@@ -694,14 +694,16 @@ class JobResult(models.Model):
             self.status = JobResultStatusChoices.STATUS_RUNNING
             JobResult.objects.filter(pk=self.pk).update(started=self.started, status=self.status)
 
-    def set_status(self, status):
+    def terminate(self, status=JobResultStatusChoices.STATUS_COMPLETED):
         """
-        Helper method to change the status of the job result. If the target status is terminal, the completion
-        time is also set.
+        Mark the job as completed, optionally specifying a particular termination status.
         """
+        valid_statuses = JobResultStatusChoices.TERMINAL_STATE_CHOICES
+        if status not in valid_statuses:
+            raise ValueError(f"Invalid status for job termination. Choices are: {', '.join(valid_statuses)}")
         self.status = status
-        if status in JobResultStatusChoices.TERMINAL_STATE_CHOICES:
-            self.completed = timezone.now()
+        self.completed = timezone.now()
+        JobResult.objects.filter(pk=self.pk).update(status=self.status, completed=self.completed)
 
     @classmethod
     def enqueue_job(cls, func, name, obj_type, user, schedule_at=None, interval=None, *args, **kwargs):

netbox/extras/reports.py

@@ -85,8 +85,7 @@ def run_report(job_result, *args, **kwargs):
         job_result.start()
         report.run(job_result)
     except Exception:
-        job_result.set_status(JobResultStatusChoices.STATUS_ERRORED)
-        job_result.save()
+        job_result.terminate(status=JobResultStatusChoices.STATUS_ERRORED)
         logging.error(f"Error during execution of report {job_result.name}")
     finally:
         # Schedule the next job if an interval has been set
@@ -241,28 +240,23 @@ class Report(object):
         self.pre_run()
 
         try:
-
             for method_name in self.test_methods:
                 self.active_test = method_name
                 test_method = getattr(self, method_name)
                 test_method()
-
             if self.failed:
                 self.logger.warning("Report failed")
                 job_result.status = JobResultStatusChoices.STATUS_FAILED
             else:
                 self.logger.info("Report completed successfully")
                 job_result.status = JobResultStatusChoices.STATUS_COMPLETED
-
         except Exception as e:
             stacktrace = traceback.format_exc()
             self.log_failure(None, f"An exception occurred: {type(e).__name__}: {e} <pre>{stacktrace}</pre>")
             logger.error(f"Exception raised during report execution: {e}")
-            job_result.set_status(JobResultStatusChoices.STATUS_ERRORED)
-
-        job_result.data = self._results
-        job_result.completed = timezone.now()
-        job_result.save()
+            job_result.terminate(status=JobResultStatusChoices.STATUS_ERRORED)
+        finally:
+            job_result.terminate()
 
         # Perform any post-run tasks
         self.post_run()

netbox/extras/scripts.py

@@ -460,36 +460,28 @@ def run_script(data, request, commit=True, *args, **kwargs):
         the change_logging context manager (which is bypassed if commit == False).
         """
         try:
-            with transaction.atomic():
-                script.output = script.run(data=data, commit=commit)
-                job_result.set_status(JobResultStatusChoices.STATUS_COMPLETED)
-
-                if not commit:
-                    raise AbortTransaction()
-
-        except AbortTransaction:
-            script.log_info("Database changes have been reverted automatically.")
-            clear_webhooks.send(request)
-        except AbortScript as e:
-            script.log_failure(
-                f"Script aborted with error: {e}"
-            )
-            script.log_info("Database changes have been reverted due to error.")
-            logger.error(f"Script aborted with error: {e}")
-            job_result.set_status(JobResultStatusChoices.STATUS_ERRORED)
-            clear_webhooks.send(request)
+            try:
+                with transaction.atomic():
+                    script.output = script.run(data=data, commit=commit)
+                    if not commit:
+                        raise AbortTransaction()
+            except AbortTransaction:
+                script.log_info("Database changes have been reverted automatically.")
+                clear_webhooks.send(request)
+            job_result.data = ScriptOutputSerializer(script).data
+            job_result.terminate()
         except Exception as e:
-            stacktrace = traceback.format_exc()
-            script.log_failure(
-                f"An exception occurred: `{type(e).__name__}: {e}`\n```\n{stacktrace}\n```"
-            )
+            if type(e) is AbortScript:
+                script.log_failure(f"Script aborted with error: {e}")
+                logger.error(f"Script aborted with error: {e}")
+            else:
+                stacktrace = traceback.format_exc()
+                script.log_failure(f"An exception occurred: `{type(e).__name__}: {e}`\n```\n{stacktrace}\n```")
+                logger.error(f"Exception raised during script execution: {e}")
             script.log_info("Database changes have been reverted due to error.")
-            logger.error(f"Exception raised during script execution: {e}")
-            job_result.set_status(JobResultStatusChoices.STATUS_ERRORED)
-            clear_webhooks.send(request)
-        finally:
             job_result.data = ScriptOutputSerializer(script).data
-            job_result.save()
+            job_result.terminate(status=JobResultStatusChoices.STATUS_ERRORED)
+            clear_webhooks.send(request)
 
         logger.info(f"Script completed in {job_result.duration}")