Просмотр исходного кода

Implemented run_script() wrapper

Jeremy Stretch 6 лет назад
Родитель
Сommit
3e75da4307
2 измененных файлов с 20 добавлено и 5 удалено
  1. 18 0
      netbox/extras/scripts.py
  2. 2 5
      netbox/extras/views.py

+ 18 - 0
netbox/extras/scripts.py

@@ -5,6 +5,7 @@ import pkgutil
 from django import forms
 from django.conf import settings
 from django.core.validators import RegexValidator
+from django.db import transaction
 
 from .constants import LOG_DEFAULT, LOG_FAILURE, LOG_INFO, LOG_SUCCESS, LOG_WARNING
 from .forms import ScriptForm
@@ -196,6 +197,23 @@ def is_variable(obj):
     return isinstance(obj, ScriptVariable)
 
 
+def run_script(script, data=None):
+    """
+    A wrapper for calling Script.run(). This performs error handling. It exists outside of the Script class to ensure
+    it cannot be overridden by a script author.
+    """
+    try:
+        with transaction.atomic():
+            return script.run(data)
+    except Exception as e:
+        script.log_failure(
+            "An exception occurred: {}".format(e)
+        )
+        script.log_info(
+            "Database changes have been reverted automatically."
+        )
+
+
 def get_scripts():
     scripts = OrderedDict()
 

+ 2 - 5
netbox/extras/views.py

@@ -3,7 +3,6 @@ from django.conf import settings
 from django.contrib import messages
 from django.contrib.auth.mixins import PermissionRequiredMixin
 from django.contrib.contenttypes.models import ContentType
-from django.db import transaction
 from django.db.models import Count, Q
 from django.http import Http404, HttpResponseForbidden
 from django.shortcuts import get_object_or_404, redirect, render
@@ -21,7 +20,7 @@ from .forms import (
 )
 from .models import ConfigContext, ImageAttachment, ObjectChange, ReportResult, Tag, TaggedItem
 from .reports import get_report, get_reports
-from .scripts import get_scripts
+from .scripts import get_scripts, run_script
 from .tables import ConfigContextTable, ObjectChangeTable, TagTable, TaggedItemTable
 
 
@@ -405,9 +404,7 @@ class ScriptView(PermissionRequiredMixin, View):
         output = None
 
         if form.is_valid():
-
-            with transaction.atomic():
-                output = script.run(form.cleaned_data)
+            run_script(script)
 
         return render(request, 'extras/script.html', {
             'module': module,