From 51e770fa9c7a5e5702ac84ce4b90c7be6e194e27 Mon Sep 17 00:00:00 2001
From: Arthur Hanson <worldnomad@gmail.com>
Date: Mon, 2 Dec 2024 12:53:10 -0800
Subject: [PATCH] 18136 Allow Script running within a Branch

---
 docs/configuration/system.md |  8 ++++++++
 netbox/extras/jobs.py        | 23 +++++++++++++++++++++--
 netbox/netbox/settings.py    |  1 +
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/docs/configuration/system.md b/docs/configuration/system.md
index 25c724bc9..39c945e7a 100644
--- a/docs/configuration/system.md
+++ b/docs/configuration/system.md
@@ -219,3 +219,11 @@ The time zone NetBox will use when dealing with dates and times. It is recommend
 Default: True
 
 Enables language translation for the user interface. (This parameter maps to Django's [USE_I18N](https://docs.djangoproject.com/en/stable/ref/settings/#std-setting-USE_I18N) setting.)
+
+---
+
+## BRANCHING_BACKEND
+
+Default: None
+
+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`
diff --git a/netbox/extras/jobs.py b/netbox/extras/jobs.py
index 190166b5b..b3bf44587 100644
--- a/netbox/extras/jobs.py
+++ b/netbox/extras/jobs.py
@@ -2,7 +2,10 @@ import logging
 import traceback
 from contextlib import nullcontext
 
+from django.conf import settings
+from django.core.exceptions import ImproperlyConfigured
 from django.db import transaction
+from django.utils.module_loading import import_string
 from django.utils.translation import gettext as _
 
 from core.signals import clear_events
@@ -36,6 +39,7 @@ class ScriptJob(JobRunner):
         """
         logger = logging.getLogger(f"netbox.scripts.{script.full_name}")
         logger.info(f"Running script (commit={commit})")
+        print("running script")
 
         try:
             try:
@@ -49,6 +53,7 @@ class ScriptJob(JobRunner):
                     logger.warning("Script failed")
 
         except Exception as e:
+            print(f"exception: {e}")
             if type(e) is AbortScript:
                 msg = _("Script aborted with error: ") + str(e)
                 if is_report(type(script)):
@@ -100,5 +105,19 @@ class ScriptJob(JobRunner):
 
         # 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():
-            self.run_script(script, request, data, commit)
+        branch = None
+        if settings.BRANCHING_BACKEND:
+            try:
+                branching_cls = import_string(settings.BRANCHING_BACKEND)
+            except AttributeError:
+                logger = logging.getLogger(f"netbox.scripts.{script.full_name}")
+                message = _("Failed to import configured BRANCHING_BACKEND: ") + settings.BRANCHING_BACKEND
+                logger.error(message)
+                raise ImproperlyConfigured(message)
+
+            branching_backend = branching_cls()
+            branch = branching_backend.get_active_branch(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)
diff --git a/netbox/netbox/settings.py b/netbox/netbox/settings.py
index a8ac68d4d..af0fc727e 100644
--- a/netbox/netbox/settings.py
+++ b/netbox/netbox/settings.py
@@ -179,6 +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)
 
 # Load any dynamic configuration parameters which have been hard-coded in the configuration file
 for param in CONFIG_PARAMS: