diff --git a/netbox/extras/models/customfields.py b/netbox/extras/models/customfields.py index 270de3488..dbf086872 100644 --- a/netbox/extras/models/customfields.py +++ b/netbox/extras/models/customfields.py @@ -787,6 +787,8 @@ class CustomFieldChoiceSet(CloningMixin, ExportTemplatesMixin, ChangeLoggedModel 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): @@ -822,33 +824,31 @@ class CustomFieldChoiceSet(CloningMixin, ExportTemplatesMixin, ChangeLoggedModel if not self.base_choices and not self.extra_choices: raise ValidationError(_("Must define base or extra choices.")) - # check if removing any used choices - if self._original_extra_choices: - original_choices = new_choices = [] - original_choices = [obj[0] for obj in self._original_extra_choices] - if self.extra_choices: - new_choices = [obj[0] for obj in self.extra_choices] - - # only care about what fields are being deleted. - if diff_choices := list(set(original_choices) - set(new_choices)): - # Get all CustomFields using this ChoiceSet - for custom_field in self.choices_for.all(): - # Then the models using those custom fields - for object_type in custom_field.object_types.all(): - # unfortunately querying the whole array of diff_choices doesn't work - for choice in diff_choices: - # Need to OR query as can be multiple choice or single choice so - # have to do both contains and equals to catch both - query_args = { - f"custom_field_data__{custom_field.name}__contains": choice, - f"custom_field_data__{custom_field.name}": choice - } - if object_type.model_class().objects.filter(models.Q(**query_args, _connector=models.Q.OR)).exists(): - raise ValidationError( - _("Cannot remove choice {choice} as it is used in records of model {model}").format( - choice=choice, model=object_type - ) - ) + # 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): diff --git a/netbox/extras/tests/test_customfields.py b/netbox/extras/tests/test_customfields.py index c365e9653..2bc9b5acc 100644 --- a/netbox/extras/tests/test_customfields.py +++ b/netbox/extras/tests/test_customfields.py @@ -343,13 +343,16 @@ class CustomFieldTest(TestCase): instance.refresh_from_db() self.assertIsNone(instance.custom_field_data.get(cf.name)) - def test_select_validation(self): + 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'), - ('abcde', 'Option ABCDE'), ) # Create a set of custom field choices @@ -358,7 +361,7 @@ class CustomFieldTest(TestCase): extra_choices=CHOICES ) - # Create a custom field & check that initial value is null + # Create a select custom field cf = CustomField.objects.create( name='select_field', type=CustomFieldTypeChoices.TYPE_SELECT, @@ -367,7 +370,7 @@ class CustomFieldTest(TestCase): ) cf.object_types.set([self.object_type]) - # Create a custom field & check that initial value is null + # Create a multi-select custom field cf_multiselect = CustomField.objects.create( name='multiselect_field', type=CustomFieldTypeChoices.TYPE_MULTISELECT, @@ -376,48 +379,37 @@ class CustomFieldTest(TestCase): ) cf_multiselect.object_types.set([self.object_type]) + # Assign a choice for both custom fields on an object instance = Site.objects.first() - - # Assign a value and check that it is saved instance.custom_field_data[cf.name] = 'a' instance.custom_field_data[cf_multiselect.name] = ['b', 'c'] instance.save() - instance.refresh_from_db() - # check can't delete single choice custom field option + # Attempting to delete a selected choice should fail with self.assertRaises(ValidationError): - CHOICES = ( + choice_set.extra_choices = ( ('b', 'Option B'), ('c', 'Option C'), ('d', 'Option D'), - ('abcde', 'Option ABCDE'), ) - choice_set.extra_choices = CHOICES choice_set.full_clean() - choice_set.save() - # check can't delete multi choice custom field option + # Attempting to delete either of the multi-select choices should fail with self.assertRaises(ValidationError): - CHOICES = ( + choice_set.extra_choices = ( ('a', 'Option A'), ('b', 'Option B'), ('d', 'Option D'), - ('abcde', 'Option ABCDE'), ) - choice_set.extra_choices = CHOICES choice_set.full_clean() - choice_set.save() - # delete non selected option should work fine - CHOICES = ( + # Removing a non-selected choice should succeed + choice_set.extra_choices = ( ('a', 'Option A'), ('b', 'Option B'), ('c', 'Option C'), - ('abcde', 'Option ABCDE'), ) - choice_set.extra_choices = CHOICES choice_set.full_clean() - choice_set.save() def test_object_field(self): value = VLAN.objects.create(name='VLAN 1', vid=1).pk