From 03a769e8ec6f5238804c38eb96141871ce3cffc2 Mon Sep 17 00:00:00 2001
From: Arthur Hanson <worldnomad@gmail.com>
Date: Mon, 9 Dec 2024 09:33:13 -0800
Subject: [PATCH] 18136 update script runner to use SCRIPT_CONTEXT_MANAGERS

---
 docs/configuration/system.md |  6 +++---
 netbox/extras/jobs.py        | 27 +++++++++++++++------------
 netbox/netbox/settings.py    |  2 +-
 3 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/docs/configuration/system.md b/docs/configuration/system.md
index 39c945e7a..3cc277660 100644
--- a/docs/configuration/system.md
+++ b/docs/configuration/system.md
@@ -222,8 +222,8 @@ Enables language translation for the user interface. (This parameter maps to Dja
 
 ---
 
-## BRANCHING_BACKEND
+## SCRIPT_CONTEXT_MANAGERS
 
-Default: None
+Default: []
 
-The dotted path to the desired branching class. If using [netboxlabs-netbox-branching](https://github.com/netboxlabs/netbox-branching) set this to `netbox_branching.backends.BranchingBackend`
+The dotted path to any additional context managers to use when running scripts.
diff --git a/netbox/extras/jobs.py b/netbox/extras/jobs.py
index da6810aac..ee4b00ca4 100644
--- a/netbox/extras/jobs.py
+++ b/netbox/extras/jobs.py
@@ -1,6 +1,6 @@
 import logging
 import traceback
-from contextlib import nullcontext
+from contextlib import ExitStack, nullcontext
 
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
@@ -91,6 +91,7 @@ class ScriptJob(JobRunner):
             commit: Passed through to Script.run()
         """
         script = ScriptModel.objects.get(pk=self.job.object_id).python_class()
+        logger = logging.getLogger(f"netbox.scripts.{script.full_name}")
 
         # Add files to form data
         if request:
@@ -101,21 +102,23 @@ class ScriptJob(JobRunner):
         # 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.
-        branch = None
-        if settings.BRANCHING_BACKEND:
+        context_managers = []
+        for context_manager_str in settings.SCRIPT_CONTEXT_MANAGERS:
             try:
-                branching_cls = import_string(settings.BRANCHING_BACKEND)
+                context_managers.append(import_string(context_manager_str))
             except AttributeError:
-                logger = logging.getLogger(f"netbox.scripts.{script.full_name}")
-                message = _("Failed to import configured BRANCHING_BACKEND: ") + settings.BRANCHING_BACKEND
+                message = _("Failed to import configured SCRIPT_CONTEXT_MANAGERS: ") + context_manager_str
                 logger.error(message)
                 raise ImproperlyConfigured(message)
 
-            branching_backend = branching_cls()
-            branch = branching_backend.get_active_branch(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():
+            if context_managers:
+                with ExitStack() as stack:
+                    for cm in context_managers:
+                        stack.enter_context(cm(request))
 
-        with branching_backend.activate_branch(branch) if branch else nullcontext():
-            with event_tracking(request) if commit else nullcontext():
+                    self.run_script(script, request, data, commit)
+            else:
                 self.run_script(script, request, data, commit)
diff --git a/netbox/netbox/settings.py b/netbox/netbox/settings.py
index af0fc727e..83450efc3 100644
--- a/netbox/netbox/settings.py
+++ b/netbox/netbox/settings.py
@@ -179,7 +179,7 @@ STORAGE_BACKEND = getattr(configuration, 'STORAGE_BACKEND', None)
 STORAGE_CONFIG = getattr(configuration, 'STORAGE_CONFIG', {})
 TIME_ZONE = getattr(configuration, 'TIME_ZONE', 'UTC')
 TRANSLATION_ENABLED = getattr(configuration, 'TRANSLATION_ENABLED', True)
-BRANCHING_BACKEND = getattr(configuration, 'BRANCHING_BACKEND', None)
+SCRIPT_CONTEXT_MANAGERS = getattr(configuration, 'SCRIPT_CONTEXT_MANAGERS', [])
 
 # Load any dynamic configuration parameters which have been hard-coded in the configuration file
 for param in CONFIG_PARAMS: