Browse Source

Closes #21751: Enable toggling user notifications when executing custom scripts (#21923)

Jeremy Stretch 1 tháng trước cách đây
mục cha
commit
b68b0c6d78

+ 14 - 0
docs/customization/custom-scripts.md

@@ -115,6 +115,20 @@ commit_default = False
 
 
 By default, a script can be scheduled for execution at a later time. Setting `scheduling_enabled` to False disables this ability: Only immediate execution will be possible. (This also disables the ability to set a recurring execution interval.)
 By default, a script can be scheduled for execution at a later time. Setting `scheduling_enabled` to False disables this ability: Only immediate execution will be possible. (This also disables the ability to set a recurring execution interval.)
 
 
+### `notifications_default`
+
+By default, a notification is generated for the requesting user each time a script finishes running. This attribute sets the initial value for the notifications field when running a script. Valid values are `always` (default), `on_failure`, and `never`.
+
+```python
+notifications_default = 'on_failure'
+```
+
+| Value | Behavior |
+|-------|----------|
+| `always` | Notify on every completion (default) |
+| `on_failure` | Notify only when the job fails or errors |
+| `never` | Never send a notification |
+
 ### `job_timeout`
 ### `job_timeout`
 
 
 Set the maximum allowed runtime for the script. If not set, `RQ_DEFAULT_TIMEOUT` will be used.
 Set the maximum allowed runtime for the script. If not set, `RQ_DEFAULT_TIMEOUT` will be used.

+ 2 - 1
netbox/core/api/serializers_/jobs.py

@@ -26,13 +26,14 @@ class JobSerializer(BaseModelSerializer):
     object = serializers.SerializerMethodField(
     object = serializers.SerializerMethodField(
         read_only=True
         read_only=True
     )
     )
+    notifications = ChoiceField(choices=JobNotificationChoices, read_only=True)
 
 
     class Meta:
     class Meta:
         model = Job
         model = Job
         fields = [
         fields = [
             'id', 'url', 'display_url', 'display', 'object_type', 'object_id', 'object', 'name', 'status', 'created',
             'id', 'url', 'display_url', 'display', 'object_type', 'object_id', 'object', 'name', 'status', 'created',
             'scheduled', 'interval', 'started', 'completed', 'user', 'data', 'error', 'job_id', 'queue_name',
             'scheduled', 'interval', 'started', 'completed', 'user', 'data', 'error', 'job_id', 'queue_name',
-            'log_entries',
+            'notifications', 'log_entries',
         ]
         ]
         brief_fields = ('url', 'created', 'completed', 'user', 'status')
         brief_fields = ('url', 'created', 'completed', 'user', 'status')
 
 

+ 12 - 0
netbox/core/choices.py

@@ -72,6 +72,18 @@ class JobStatusChoices(ChoiceSet):
     )
     )
 
 
 
 
+class JobNotificationChoices(ChoiceSet):
+    NOTIFICATION_ALWAYS = 'always'
+    NOTIFICATION_ON_FAILURE = 'on_failure'
+    NOTIFICATION_NEVER = 'never'
+
+    CHOICES = (
+        (NOTIFICATION_ALWAYS, _('Always')),
+        (NOTIFICATION_ON_FAILURE, _('On failure')),
+        (NOTIFICATION_NEVER, _('Never')),
+    )
+
+
 class JobIntervalChoices(ChoiceSet):
 class JobIntervalChoices(ChoiceSet):
     INTERVAL_MINUTELY = 1
     INTERVAL_MINUTELY = 1
     INTERVAL_HOURLY = 60
     INTERVAL_HOURLY = 60

+ 16 - 0
netbox/core/migrations/0024_job_notifications.py

@@ -0,0 +1,16 @@
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('core', '0023_datasource_sync_permission'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='job',
+            name='notifications',
+            field=models.CharField(default='always', max_length=30),
+        ),
+    ]

+ 21 - 8
netbox/core/models/jobs.py

@@ -16,7 +16,7 @@ from django.utils import timezone
 from django.utils.translation import gettext as _
 from django.utils.translation import gettext as _
 from rq.exceptions import InvalidJobOperation
 from rq.exceptions import InvalidJobOperation
 
 
-from core.choices import JobStatusChoices
+from core.choices import JobNotificationChoices, JobStatusChoices
 from core.dataclasses import JobLogEntry
 from core.dataclasses import JobLogEntry
 from core.events import JOB_COMPLETED, JOB_ERRORED, JOB_FAILED
 from core.events import JOB_COMPLETED, JOB_ERRORED, JOB_FAILED
 from core.models import ObjectType
 from core.models import ObjectType
@@ -118,6 +118,12 @@ class Job(models.Model):
         blank=True,
         blank=True,
         help_text=_('Name of the queue in which this job was enqueued')
         help_text=_('Name of the queue in which this job was enqueued')
     )
     )
+    notifications = models.CharField(
+        verbose_name=_('notifications'),
+        max_length=30,
+        choices=JobNotificationChoices,
+        default=JobNotificationChoices.NOTIFICATION_ALWAYS
+    )
     log_entries = ArrayField(
     log_entries = ArrayField(
         verbose_name=_('log entries'),
         verbose_name=_('log entries'),
         base_field=models.JSONField(
         base_field=models.JSONField(
@@ -238,12 +244,16 @@ class Job(models.Model):
         self.save()
         self.save()
 
 
         # Notify the user (if any) of completion
         # Notify the user (if any) of completion
-        if self.user:
-            Notification(
-                user=self.user,
-                object=self,
-                event_type=self.get_event_type(),
-            ).save()
+        if self.user and self.notifications != JobNotificationChoices.NOTIFICATION_NEVER:
+            if (
+                self.notifications == JobNotificationChoices.NOTIFICATION_ALWAYS or
+                status != JobStatusChoices.STATUS_COMPLETED
+            ):
+                Notification(
+                    user=self.user,
+                    object=self,
+                    event_type=self.get_event_type(),
+                ).save()
 
 
         # Send signal
         # Send signal
         job_end.send(self)
         job_end.send(self)
@@ -267,6 +277,7 @@ class Job(models.Model):
             interval=None,
             interval=None,
             immediate=False,
             immediate=False,
             queue_name=None,
             queue_name=None,
+            notifications=None,
             **kwargs
             **kwargs
     ):
     ):
         """
         """
@@ -281,6 +292,7 @@ class Job(models.Model):
             interval: Recurrence interval (in minutes)
             interval: Recurrence interval (in minutes)
             immediate: Run the job immediately without scheduling it in the background. Should be used for interactive
             immediate: Run the job immediately without scheduling it in the background. Should be used for interactive
                 management commands only.
                 management commands only.
+            notifications: Notification behavior on job completion (always, on_failure, or never)
         """
         """
         if schedule_at and immediate:
         if schedule_at and immediate:
             raise ValueError(_("enqueue() cannot be called with values for both schedule_at and immediate."))
             raise ValueError(_("enqueue() cannot be called with values for both schedule_at and immediate."))
@@ -302,7 +314,8 @@ class Job(models.Model):
             interval=interval,
             interval=interval,
             user=user,
             user=user,
             job_id=uuid.uuid4(),
             job_id=uuid.uuid4(),
-            queue_name=rq_queue_name
+            queue_name=rq_queue_name,
+            notifications=notifications if notifications is not None else JobNotificationChoices.NOTIFICATION_ALWAYS
         )
         )
         job.full_clean()
         job.full_clean()
         job.save()
         job.save()

+ 88 - 1
netbox/core/tests/test_models.py

@@ -1,13 +1,16 @@
+import uuid
 from unittest.mock import MagicMock, patch
 from unittest.mock import MagicMock, patch
 
 
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ObjectDoesNotExist
 from django.core.exceptions import ObjectDoesNotExist
 from django.test import TestCase
 from django.test import TestCase
 
 
-from core.choices import ObjectChangeActionChoices
+from core.choices import JobNotificationChoices, JobStatusChoices, ObjectChangeActionChoices
 from core.models import DataSource, Job, ObjectType
 from core.models import DataSource, Job, ObjectType
 from dcim.models import Device, Location, Site
 from dcim.models import Device, Location, Site
+from extras.models import Notification
 from netbox.constants import CENSOR_TOKEN, CENSOR_TOKEN_CHANGED
 from netbox.constants import CENSOR_TOKEN, CENSOR_TOKEN_CHANGED
+from users.models import User
 
 
 
 
 class DataSourceIgnoreRulesTestCase(TestCase):
 class DataSourceIgnoreRulesTestCase(TestCase):
@@ -226,6 +229,18 @@ class ObjectTypeTest(TestCase):
 
 
 class JobTest(TestCase):
 class JobTest(TestCase):
 
 
+    def _make_job(self, user, notifications):
+        """
+        Create and return a persisted Job with the given user and notifications setting.
+        """
+        return Job.objects.create(
+            name='Test Job',
+            job_id=uuid.uuid4(),
+            user=user,
+            notifications=notifications,
+            status=JobStatusChoices.STATUS_RUNNING,
+        )
+
     @patch('core.models.jobs.django_rq.get_queue')
     @patch('core.models.jobs.django_rq.get_queue')
     def test_delete_cancels_job_from_correct_queue(self, mock_get_queue):
     def test_delete_cancels_job_from_correct_queue(self, mock_get_queue):
         """
         """
@@ -257,3 +272,75 @@ class JobTest(TestCase):
         mock_get_queue.assert_called_with(custom_queue)
         mock_get_queue.assert_called_with(custom_queue)
         mock_queue.fetch_job.assert_called_with(str(job.job_id))
         mock_queue.fetch_job.assert_called_with(str(job.job_id))
         mock_rq_job.cancel.assert_called_once()
         mock_rq_job.cancel.assert_called_once()
+
+    @patch('core.models.jobs.job_end')
+    def test_terminate_notification_always(self, mock_job_end):
+        """
+        With notifications=always, a Notification should be created for every
+        terminal status (completed, failed, errored).
+        """
+        user = User.objects.create_user(username='notification-always')
+
+        for status in (
+            JobStatusChoices.STATUS_COMPLETED,
+            JobStatusChoices.STATUS_FAILED,
+            JobStatusChoices.STATUS_ERRORED,
+        ):
+            with self.subTest(status=status):
+                job = self._make_job(user, JobNotificationChoices.NOTIFICATION_ALWAYS)
+                job.terminate(status=status)
+                self.assertEqual(
+                    Notification.objects.filter(user=user, object_id=job.pk).count(),
+                    1,
+                    msg=f"Expected a notification for status={status} with notifications=always",
+                )
+
+    @patch('core.models.jobs.job_end')
+    def test_terminate_notification_on_failure(self, mock_job_end):
+        """
+        With notifications=on_failure, a Notification should be created only for
+        non-completed terminal statuses (failed, errored), not for completed.
+        """
+        user = User.objects.create_user(username='notification-on-failure')
+
+        # No notification on successful completion
+        job = self._make_job(user, JobNotificationChoices.NOTIFICATION_ON_FAILURE)
+        job.terminate(status=JobStatusChoices.STATUS_COMPLETED)
+        self.assertEqual(
+            Notification.objects.filter(user=user, object_id=job.pk).count(),
+            0,
+            msg="Expected no notification for status=completed with notifications=on_failure",
+        )
+
+        # Notification on failure/error
+        for status in (JobStatusChoices.STATUS_FAILED, JobStatusChoices.STATUS_ERRORED):
+            with self.subTest(status=status):
+                job = self._make_job(user, JobNotificationChoices.NOTIFICATION_ON_FAILURE)
+                job.terminate(status=status)
+                self.assertEqual(
+                    Notification.objects.filter(user=user, object_id=job.pk).count(),
+                    1,
+                    msg=f"Expected a notification for status={status} with notifications=on_failure",
+                )
+
+    @patch('core.models.jobs.job_end')
+    def test_terminate_notification_never(self, mock_job_end):
+        """
+        With notifications=never, no Notification should be created regardless
+        of terminal status.
+        """
+        user = User.objects.create_user(username='notification-never')
+
+        for status in (
+            JobStatusChoices.STATUS_COMPLETED,
+            JobStatusChoices.STATUS_FAILED,
+            JobStatusChoices.STATUS_ERRORED,
+        ):
+            with self.subTest(status=status):
+                job = self._make_job(user, JobNotificationChoices.NOTIFICATION_NEVER)
+                job.terminate(status=status)
+                self.assertEqual(
+                    Notification.objects.filter(user=user, object_id=job.pk).count(),
+                    0,
+                    msg=f"Expected no notification for status={status} with notifications=never",
+                )

+ 14 - 1
netbox/extras/api/serializers_/scripts.py

@@ -7,7 +7,7 @@ from drf_spectacular.utils import extend_schema_field
 from rest_framework import serializers
 from rest_framework import serializers
 
 
 from core.api.serializers_.jobs import JobSerializer
 from core.api.serializers_.jobs import JobSerializer
-from core.choices import ManagedFileRootPathChoices
+from core.choices import JobNotificationChoices, ManagedFileRootPathChoices
 from extras.models import Script, ScriptModule
 from extras.models import Script, ScriptModule
 from netbox.api.serializers import ValidatedModelSerializer
 from netbox.api.serializers import ValidatedModelSerializer
 from utilities.datetime import local_now
 from utilities.datetime import local_now
@@ -114,6 +114,19 @@ class ScriptInputSerializer(serializers.Serializer):
     commit = serializers.BooleanField()
     commit = serializers.BooleanField()
     schedule_at = serializers.DateTimeField(required=False, allow_null=True)
     schedule_at = serializers.DateTimeField(required=False, allow_null=True)
     interval = serializers.IntegerField(required=False, allow_null=True)
     interval = serializers.IntegerField(required=False, allow_null=True)
+    notifications = serializers.ChoiceField(
+        choices=JobNotificationChoices,
+        required=False,
+        default=JobNotificationChoices.NOTIFICATION_ALWAYS,
+    )
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Default to script's Meta.notifications_default if set
+        script = self.context.get('script')
+        if script and script.python_class:
+            self.fields['notifications'].default = script.python_class.notifications_default
 
 
     def validate_schedule_at(self, value):
     def validate_schedule_at(self, value):
         """
         """

+ 2 - 1
netbox/extras/api/views.py

@@ -338,7 +338,8 @@ class ScriptViewSet(ModelViewSet):
                 commit=input_serializer.data['commit'],
                 commit=input_serializer.data['commit'],
                 job_timeout=script.python_class.job_timeout,
                 job_timeout=script.python_class.job_timeout,
                 schedule_at=input_serializer.validated_data.get('schedule_at'),
                 schedule_at=input_serializer.validated_data.get('schedule_at'),
-                interval=input_serializer.validated_data.get('interval')
+                interval=input_serializer.validated_data.get('interval'),
+                notifications=input_serializer.validated_data.get('notifications'),
             )
             )
             serializer = serializers.ScriptDetailSerializer(script, context={'request': request})
             serializer = serializers.ScriptDetailSerializer(script, context={'request': request})
 
 

+ 13 - 1
netbox/extras/forms/scripts.py

@@ -2,7 +2,7 @@ from django import forms
 from django.core.files.storage import storages
 from django.core.files.storage import storages
 from django.utils.translation import gettext_lazy as _
 from django.utils.translation import gettext_lazy as _
 
 
-from core.choices import JobIntervalChoices
+from core.choices import JobIntervalChoices, JobNotificationChoices
 from core.forms import ManagedFileForm
 from core.forms import ManagedFileForm
 from utilities.datetime import local_now
 from utilities.datetime import local_now
 from utilities.forms.widgets import DateTimePicker, NumberWithOptions
 from utilities.forms.widgets import DateTimePicker, NumberWithOptions
@@ -35,6 +35,13 @@ class ScriptForm(forms.Form):
         ),
         ),
         help_text=_("Interval at which this script is re-run (in minutes)")
         help_text=_("Interval at which this script is re-run (in minutes)")
     )
     )
+    _notifications = forms.ChoiceField(
+        required=False,
+        choices=JobNotificationChoices,
+        initial=JobNotificationChoices.NOTIFICATION_ALWAYS,
+        label=_("Notifications"),
+        help_text=_("When to notify the user of job completion")
+    )
 
 
     def __init__(self, *args, scheduling_enabled=True, **kwargs):
     def __init__(self, *args, scheduling_enabled=True, **kwargs):
         super().__init__(*args, **kwargs)
         super().__init__(*args, **kwargs)
@@ -57,6 +64,11 @@ class ScriptForm(forms.Form):
         if self.cleaned_data.get('_interval') and not scheduled_time:
         if self.cleaned_data.get('_interval') and not scheduled_time:
             self.cleaned_data['_schedule_at'] = local_now()
             self.cleaned_data['_schedule_at'] = local_now()
 
 
+        # Fall back to the field's initial value if no notification preference was submitted
+        # (e.g. when running a script via the "Run Script" button on the scripts list view)
+        if not self.cleaned_data.get('_notifications'):
+            self.cleaned_data['_notifications'] = self.fields['_notifications'].initial
+
         return self.cleaned_data
         return self.cleaned_data
 
 
 
 

+ 12 - 1
netbox/extras/scripts.py

@@ -10,6 +10,7 @@ from django.utils import timezone
 from django.utils.functional import classproperty
 from django.utils.functional import classproperty
 from django.utils.translation import gettext as _
 from django.utils.translation import gettext as _
 
 
+from core.choices import JobNotificationChoices
 from extras.choices import LogLevelChoices
 from extras.choices import LogLevelChoices
 from extras.models import ScriptModule
 from extras.models import ScriptModule
 from ipam.formfields import IPAddressFormField, IPNetworkFormField
 from ipam.formfields import IPAddressFormField, IPNetworkFormField
@@ -389,6 +390,10 @@ class BaseScript:
     def scheduling_enabled(self):
     def scheduling_enabled(self):
         return getattr(self.Meta, 'scheduling_enabled', True)
         return getattr(self.Meta, 'scheduling_enabled', True)
 
 
+    @classproperty
+    def notifications_default(self):
+        return getattr(self.Meta, 'notifications_default', JobNotificationChoices.NOTIFICATION_ALWAYS)
+
     @property
     @property
     def filename(self):
     def filename(self):
         return inspect.getfile(self.__class__)
         return inspect.getfile(self.__class__)
@@ -491,7 +496,10 @@ class BaseScript:
             fieldsets.append((_('Script Data'), fields))
             fieldsets.append((_('Script Data'), fields))
 
 
         # Append the default fieldset if defined in the Meta class
         # Append the default fieldset if defined in the Meta class
-        exec_parameters = ('_schedule_at', '_interval', '_commit') if self.scheduling_enabled else ('_commit',)
+        if self.scheduling_enabled:
+            exec_parameters = ('_schedule_at', '_interval', '_commit', '_notifications')
+        else:
+            exec_parameters = ('_commit', '_notifications')
         fieldsets.append((_('Script Execution Parameters'), exec_parameters))
         fieldsets.append((_('Script Execution Parameters'), exec_parameters))
 
 
         return fieldsets
         return fieldsets
@@ -511,6 +519,9 @@ class BaseScript:
         # Set initial "commit" checkbox state based on the script's Meta parameter
         # Set initial "commit" checkbox state based on the script's Meta parameter
         form.fields['_commit'].initial = self.commit_default
         form.fields['_commit'].initial = self.commit_default
 
 
+        # Set initial "notifications" selection based on the script's Meta parameter
+        form.fields['_notifications'].initial = self.notifications_default
+
         # Hide fields if scheduling has been disabled
         # Hide fields if scheduling has been disabled
         if not self.scheduling_enabled:
         if not self.scheduling_enabled:
             form.fields['_schedule_at'].widget = forms.HiddenInput()
             form.fields['_schedule_at'].widget = forms.HiddenInput()

+ 1 - 0
netbox/extras/views.py

@@ -1707,6 +1707,7 @@ class ScriptView(BaseScriptView):
                 user=request.user,
                 user=request.user,
                 schedule_at=form.cleaned_data.pop('_schedule_at'),
                 schedule_at=form.cleaned_data.pop('_schedule_at'),
                 interval=form.cleaned_data.pop('_interval'),
                 interval=form.cleaned_data.pop('_interval'),
+                notifications=form.cleaned_data.pop('_notifications'),
                 data=form.cleaned_data,
                 data=form.cleaned_data,
                 request=copy_safe_request(request),
                 request=copy_safe_request(request),
                 job_timeout=script.python_class.job_timeout,
                 job_timeout=script.python_class.job_timeout,

+ 1 - 0
netbox/netbox/jobs.py

@@ -142,6 +142,7 @@ class JobRunner(ABC):
                     user=job.user,
                     user=job.user,
                     schedule_at=new_scheduled_time,
                     schedule_at=new_scheduled_time,
                     interval=job.interval,
                     interval=job.interval,
+                    notifications=job.notifications,
                     **kwargs,
                     **kwargs,
                 )
                 )