diff --git a/netbox_custom_objects/models.py b/netbox_custom_objects/models.py index fa2f160..3cc0caf 100644 --- a/netbox_custom_objects/models.py +++ b/netbox_custom_objects/models.py @@ -12,7 +12,7 @@ # from django.contrib.contenttypes.management import create_contenttypes from django.contrib.contenttypes.models import ContentType from django.core.validators import RegexValidator, ValidationError -from django.db import connection, models +from django.db import connection, IntegrityError, models, transaction from django.db.models import Q from django.db.models.functions import Lower from django.db.models.signals import pre_delete @@ -52,6 +52,13 @@ from netbox_custom_objects.constants import APP_LABEL, RESERVED_FIELD_NAMES from netbox_custom_objects.field_types import FIELD_TYPE_CLASS + +class UniquenessConstraintTestError(Exception): + """Custom exception used to signal successful uniqueness constraint test.""" + + pass + + USER_TABLE_DATABASE_NAME_PREFIX = "custom_objects_" @@ -872,6 +879,43 @@ def clean(self): {"unique": _("Uniqueness cannot be enforced for boolean fields")} ) + # Check if uniqueness constraint can be applied when changing from non-unique to unique + if ( + self.pk + and self.unique + and not self.original.unique + and not self._state.adding + ): + field_type = FIELD_TYPE_CLASS[self.type]() + model_field = field_type.get_model_field(self) + model = self.custom_object_type.get_model() + model_field.contribute_to_class(model, self.name) + + old_field = field_type.get_model_field(self.original) + old_field.contribute_to_class(model, self._original_name) + + try: + with transaction.atomic(): + with connection.schema_editor() as test_schema_editor: + test_schema_editor.alter_field(model, old_field, model_field) + # If we get here, the constraint was applied successfully + # Now raise a custom exception to rollback the test transaction + raise UniquenessConstraintTestError() + except UniquenessConstraintTestError: + # The constraint can be applied, validation passes + pass + except IntegrityError: + # The constraint cannot be applied due to existing non-unique values + raise ValidationError( + { + "unique": _( + "Custom objects with non-unique values already exist so this action isn't permitted" + ) + } + ) + finally: + self.custom_object_type.clear_model_cache(self.custom_object_type.id) + # Choice set must be set on selection fields, and *only* on selection fields if self.type in ( CustomFieldTypeChoices.TYPE_SELECT,