diff --git a/netbox/extras/models/customfields.py b/netbox/extras/models/customfields.py index aff41521b..91de7df0d 100644 --- a/netbox/extras/models/customfields.py +++ b/netbox/extras/models/customfields.py @@ -785,6 +785,12 @@ class CustomFieldChoiceSet(CloningMixin, ExportTemplatesMixin, ChangeLoggedModel def __str__(self): return self.name + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # Cache the initial set of choices for comparison under clean() + self._original_extra_choices = self.__dict__.get('extra_choices') + def get_absolute_url(self): return reverse('extras:customfieldchoiceset', args=[self.pk]) @@ -818,6 +824,32 @@ class CustomFieldChoiceSet(CloningMixin, ExportTemplatesMixin, ChangeLoggedModel if not self.base_choices and not self.extra_choices: raise ValidationError(_("Must define base or extra choices.")) + # Check whether any choices have been removed. If so, check whether any of the removed + # choices are still set in custom field data for any object. + original_choices = set([ + c[0] for c in self._original_extra_choices + ]) if self._original_extra_choices else set() + current_choices = set([ + c[0] for c in self.extra_choices + ]) if self.extra_choices else set() + if removed_choices := original_choices - current_choices: + for custom_field in self.choices_for.all(): + for object_type in custom_field.object_types.all(): + model = object_type.model_class() + for choice in removed_choices: + # Form the query based on the type of custom field + if custom_field.type == CustomFieldTypeChoices.TYPE_MULTISELECT: + query_args = {f"custom_field_data__{custom_field.name}__contains": choice} + else: + query_args = {f"custom_field_data__{custom_field.name}": choice} + # Raise a ValidationError if there are any objects which still reference the removed choice + if model.objects.filter(models.Q(**query_args)).exists(): + raise ValidationError( + _( + "Cannot remove choice {choice} as there are {model} objects which reference it." + ).format(choice=choice, model=object_type) + ) + def save(self, *args, **kwargs): # Sort choices if alphabetical ordering is enforced diff --git a/netbox/extras/tests/test_customfields.py b/netbox/extras/tests/test_customfields.py index 697b756ec..2bc9b5acc 100644 --- a/netbox/extras/tests/test_customfields.py +++ b/netbox/extras/tests/test_customfields.py @@ -343,6 +343,74 @@ class CustomFieldTest(TestCase): instance.refresh_from_db() self.assertIsNone(instance.custom_field_data.get(cf.name)) + def test_remove_selected_choice(self): + """ + Removing a ChoiceSet choice that is referenced by an object should raise + a ValidationError exception. + """ + CHOICES = ( + ('a', 'Option A'), + ('b', 'Option B'), + ('c', 'Option C'), + ('d', 'Option D'), + ) + + # Create a set of custom field choices + choice_set = CustomFieldChoiceSet.objects.create( + name='Custom Field Choice Set 1', + extra_choices=CHOICES + ) + + # Create a select custom field + cf = CustomField.objects.create( + name='select_field', + type=CustomFieldTypeChoices.TYPE_SELECT, + required=False, + choice_set=choice_set + ) + cf.object_types.set([self.object_type]) + + # Create a multi-select custom field + cf_multiselect = CustomField.objects.create( + name='multiselect_field', + type=CustomFieldTypeChoices.TYPE_MULTISELECT, + required=False, + choice_set=choice_set + ) + cf_multiselect.object_types.set([self.object_type]) + + # Assign a choice for both custom fields on an object + instance = Site.objects.first() + instance.custom_field_data[cf.name] = 'a' + instance.custom_field_data[cf_multiselect.name] = ['b', 'c'] + instance.save() + + # Attempting to delete a selected choice should fail + with self.assertRaises(ValidationError): + choice_set.extra_choices = ( + ('b', 'Option B'), + ('c', 'Option C'), + ('d', 'Option D'), + ) + choice_set.full_clean() + + # Attempting to delete either of the multi-select choices should fail + with self.assertRaises(ValidationError): + choice_set.extra_choices = ( + ('a', 'Option A'), + ('b', 'Option B'), + ('d', 'Option D'), + ) + choice_set.full_clean() + + # Removing a non-selected choice should succeed + choice_set.extra_choices = ( + ('a', 'Option A'), + ('b', 'Option B'), + ('c', 'Option C'), + ) + choice_set.full_clean() + def test_object_field(self): value = VLAN.objects.create(name='VLAN 1', vid=1).pk