Implement BackgroundJob for running scripts

The independent implementations of interactive and background script
execution have been merged into a single BackgroundJob implementation.

netbox/extras/api/views.py

@@ -1,5 +1,6 @@
 from django.http import Http404
 from django.shortcuts import get_object_or_404
+from django.utils.module_loading import import_string
 from django_rq.queues import get_connection
 from rest_framework import status
 from rest_framework.decorators import action
@@ -14,7 +15,6 @@
 from core.models import Job, ObjectType
 from extras import filtersets
 from extras.models import *
-from extras.scripts import run_script
 from netbox.api.authentication import IsAuthenticatedOrLoginNotRequired
 from netbox.api.features import SyncedDataMixin
 from netbox.api.metadata import ContentTypeMetadata
@@ -252,8 +252,8 @@ def post(self, request, pk):
             raise RQWorkerNotRunningException()
 
         if input_serializer.is_valid():
-            Job.enqueue(
-                run_script,
+            ScriptJob = import_string("extras.jobs.ScriptJob")
+            ScriptJob.enqueue(
                 instance=script,
                 name=script.python_class.class_name,
                 user=request.user,

netbox/extras/events.py

@@ -9,7 +9,6 @@
 from django_rq import get_queue
 
 from core.choices import ObjectChangeActionChoices
-from core.models import Job
 from netbox.config import get_config
 from netbox.constants import RQ_QUEUE_DEFAULT
 from netbox.registry import registry
@@ -125,8 +124,8 @@ def process_event_rules(event_rules, model_name, event, data, username=None, sna
             script = event_rule.action_object.python_class()
 
             # Enqueue a Job to record the script's execution
-            Job.enqueue(
-                "extras.scripts.run_script",
+            ScriptJob = import_string("extras.jobs.ScriptJob")
+            ScriptJob.enqueue(
                 instance=event_rule.action_object,
                 name=script.name,
                 user=user,

netbox/extras/jobs.py

@@ -0,0 +1,105 @@
+import logging
+import traceback
+from contextlib import nullcontext
+
+from django.db import transaction
+from django.utils.translation import gettext as _
+
+from extras.models import Script as ScriptModel
+from extras.signals import clear_events
+from netbox.context_managers import event_tracking
+from utilities.exceptions import AbortScript, AbortTransaction
+from utilities.jobs import BackgroundJob
+from .utils import is_report
+
+
+class ScriptJob(BackgroundJob):
+    """
+    Script execution job.
+
+    A wrapper for calling Script.run(). This performs error handling and provides a hook for committing changes. It
+    exists outside the Script class to ensure it cannot be overridden by a script author.
+    """
+
+    @staticmethod
+    def run_script(script, job, request, data, commit):
+        """
+        Core script execution task. We capture this within a method to allow for conditionally wrapping it with the
+        event_tracking context manager (which is bypassed if commit == False).
+
+        Args:
+            job: The Job associated with this execution
+            request: The WSGI request associated with this execution (if any)
+            data: A dictionary of data to be passed to the script upon execution
+            commit: Passed through to Script.run()
+        """
+        try:
+            try:
+                with transaction.atomic():
+                    script.output = script.run(data, commit)
+                    if not commit:
+                        raise AbortTransaction()
+            except AbortTransaction:
+                script.log_info(message=_("Database changes have been reverted automatically."))
+                if script.failed:
+                    logger.warning(f"Script failed")
+                    raise
+
+        except Exception as e:
+            if type(e) is AbortScript:
+                msg = _("Script aborted with error: ") + str(e)
+                if is_report(type(script)):
+                    script.log_failure(message=msg)
+                else:
+                    script.log_failure(msg)
+                logger.error(f"Script aborted with error: {e}")
+
+            else:
+                stacktrace = traceback.format_exc()
+                script.log_failure(
+                    message=_("An exception occurred: ") + f"`{type(e).__name__}: {e}`\n```\n{stacktrace}\n```"
+                )
+                logger.error(f"Exception raised during script execution: {e}")
+
+            if type(e) is not AbortTransaction:
+                script.log_info(message=_("Database changes have been reverted due to error."))
+
+            # Clear all pending events. Job termination (including setting the status) is handled by the job framework.
+            if request:
+                clear_events.send(request)
+            raise
+
+        # Update the job data regardless of the execution status of the job. Successes should be reported as well as
+        # failures.
+        finally:
+            job.data = script.get_job_data()
+
+    @classmethod
+    def run(cls, job, data, request=None, commit=True, **kwargs):
+        """
+        Run the script.
+
+        Args:
+            job: The Job associated with this execution
+            data: A dictionary of data to be passed to the script upon execution
+            request: The WSGI request associated with this execution (if any)
+            commit: Passed through to Script.run()
+        """
+        script = ScriptModel.objects.get(pk=job.object_id).python_class()
+
+        logger = logging.getLogger(f"netbox.scripts.{script.full_name}")
+        logger.info(f"Running script (commit={commit})")
+
+        # Add files to form data
+        if request:
+            files = request.FILES
+            for field_name, fileobj in files.items():
+                data[field_name] = fileobj
+
+        # Add the current request as a property of the script
+        script.request = request
+
+        # Execute the script. If commit is True, wrap it with the event_tracking context manager to ensure we process
+        # change logging, event rules, etc.
+        with event_tracking(request) if commit else nullcontext():
+            cls.run_script(script, job, request, data, commit)

netbox/extras/management/commands/runscript.py

@@ -1,19 +1,13 @@
 import json
 import logging
 import sys
-import traceback
 import uuid
 
 from django.contrib.auth import get_user_model
 from django.core.management.base import BaseCommand, CommandError
-from django.db import transaction
+from django.utils.module_loading import import_string
 
-from core.choices import JobStatusChoices
-from core.models import Job
 from extras.scripts import get_module_and_script
-from extras.signals import clear_events
-from netbox.context_managers import event_tracking
-from utilities.exceptions import AbortTransaction
 from utilities.request import NetBoxFakeRequest
 
 
@@ -33,44 +27,6 @@ def add_arguments(self, parser):
         parser.add_argument('script', help="Script to run")
 
     def handle(self, *args, **options):
-
-        def _run_script():
-            """
-            Core script execution task. We capture this within a subfunction to allow for conditionally wrapping it with
-            the event_tracking context manager (which is bypassed if commit == False).
-            """
-            try:
-                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_events.send(request)
-                job.data = script.get_job_data()
-                job.terminate()
-            except Exception as e:
-                stacktrace = traceback.format_exc()
-                script.log_failure(
-                    f"An exception occurred: `{type(e).__name__}: {e}`\n```\n{stacktrace}\n```"
-                )
-                script.log_info("Database changes have been reverted due to error.")
-                logger.error(f"Exception raised during script execution: {e}")
-                clear_events.send(request)
-                job.data = script.get_job_data()
-                job.terminate(status=JobStatusChoices.STATUS_ERRORED, error=repr(e))
-
-            # Print any test method results
-            for test_name, attrs in job.data['tests'].items():
-                self.stdout.write(
-                    "\t{}: {} success, {} info, {} warning, {} failure".format(
-                        test_name, attrs['success'], attrs['info'], attrs['warning'], attrs['failure']
-                    )
-                )
-
-            logger.info(f"Script completed in {job.duration}")
-
         User = get_user_model()
 
         # Params
@@ -84,8 +40,8 @@ def _run_script():
             data = {}
 
         module_name, script_name = script.split('.', 1)
-        module, script = get_module_and_script(module_name, script_name)
-        script = script.python_class
+        module, script_obj = get_module_and_script(module_name, script_name)
+        script = script_obj.python_class
 
         # Take user from command line if provided and exists, other
         if options['user']:
@@ -120,40 +76,31 @@ def _run_script():
         # Initialize the script form
         script = script()
         form = script.as_form(data, None)
-
-        # Create the job
-        job = Job.objects.create(
-            object=module,
-            name=script.class_name,
-            user=User.objects.filter(is_superuser=True).order_by('pk')[0],
-            job_id=uuid.uuid4()
-        )
-
-        request = NetBoxFakeRequest({
-            'META': {},
-            'POST': data,
-            'GET': {},
-            'FILES': {},
-            'user': user,
-            'path': '',
-            'id': job.job_id
-        })
-
-        if form.is_valid():
-            job.status = JobStatusChoices.STATUS_RUNNING
-            job.save()
-
-            logger.info(f"Running script (commit={commit})")
-            script.request = request
-
-            # Execute the script. If commit is True, wrap it with the event_tracking context manager to ensure we process
-            # change logging, webhooks, etc.
-            with event_tracking(request):
-                _run_script()
-        else:
+        if not form.is_valid():
             logger.error('Data is not valid:')
             for field, errors in form.errors.get_json_data().items():
                 for error in errors:
                     logger.error(f'\t{field}: {error.get("message")}')
-            job.status = JobStatusChoices.STATUS_ERRORED
-            job.save()
+            raise CommandError()
+
+        # Execute the script.
+        ScriptJob = import_string("extras.jobs.ScriptJob")
+        job = ScriptJob.enqueue(
+            instance=script_obj,
+            name=script.name,
+            user=user,
+            run_now=True,
+            data=data,
+            request=NetBoxFakeRequest({
+                'META': {},
+                'POST': data,
+                'GET': {},
+                'FILES': {},
+                'user': user,
+                'path': '',
+                'id': uuid.uuid4()
+            }),
+            commit=commit,
+        )
+
+        logger.info(f"Script completed in {job.duration}")