netbox/netbox/extras/api/customfields.py

from functools import cached_property

from django.core.exceptions import ValidationError as DjangoValidationError
from django.db.models import prefetch_related_objects
from django.utils.translation import gettext as _
from drf_spectacular.types import OpenApiTypes
from drf_spectacular.utils import extend_schema_field
from rest_framework.fields import Field
from rest_framework.serializers import ListSerializer, ValidationError

from extras.choices import CustomFieldTypeChoices
from extras.constants import CUSTOMFIELD_EMPTY_VALUES
from extras.models import CustomField
from utilities.api import get_prefetches_for_serializer, get_serializer_for_model
from utilities.forms.fields import LaxURLField

#
# Custom fields
#


class CustomFieldDefaultValues:
    """
    Return a dictionary of all CustomFields assigned to the parent model and their default values.
    """
    requires_context = True

    def __call__(self, serializer_field):
        self.model = serializer_field.parent.Meta.model

        # Populate the default value for each CustomField on the model
        value = {}
        for field in CustomField.objects.get_for_model(self.model):
            if field.default is not None:
                value[field.name] = field.default
            else:
                value[field.name] = None

        return value


@extend_schema_field(OpenApiTypes.OBJECT)
class CustomFieldsDataField(Field):

    def _get_custom_fields(self):
        """
        Cache CustomFields assigned to this model to avoid redundant database queries
        """
        if not hasattr(self, '_custom_fields'):
            self._custom_fields = CustomField.objects.get_for_model(self.parent.Meta.model)
        return self._custom_fields

    @cached_property
    def _nested_serializers(self):
        return {}

    def _get_nested_serializer(self, cf, many):
        """
        Return the nested serializer for the objects a custom field references, built once per field instance.
        """
        serializer_class = get_serializer_for_model(cf.related_object_type.model_class())
        if (serializer_class, many) not in self._nested_serializers:
            self._nested_serializers[(serializer_class, many)] = serializer_class(
                nested=True, many=many, context=self.parent.context
            )
        return self._nested_serializers[(serializer_class, many)]

    def to_representation(self, obj):
        data = {}
        cache = self.parent.context.get('cf_object_cache')

        for cf in self._get_custom_fields():
            if cache is not None and cf.type in (
                CustomFieldTypeChoices.TYPE_OBJECT,
                CustomFieldTypeChoices.TYPE_MULTIOBJECT,
            ):
                raw = obj.get(cf.name)
                if raw is None:
                    value = None
                elif cf.type == CustomFieldTypeChoices.TYPE_OBJECT:
                    model = cf.related_object_type.model_class()
                    value = cache.get((model, raw))
                else:
                    model = cf.related_object_type.model_class()
                    value = [cache[(model, pk)] for pk in raw if (model, pk) in cache] or None
            else:
                value = cf.deserialize(obj.get(cf.name))

            if value is not None and cf.type == CustomFieldTypeChoices.TYPE_OBJECT:
                value = self._get_nested_serializer(cf, many=False).to_representation(value)
            elif value is not None and cf.type == CustomFieldTypeChoices.TYPE_MULTIOBJECT:
                value = self._get_nested_serializer(cf, many=True).to_representation(value)
            elif cf.type in (CustomFieldTypeChoices.TYPE_SELECT, CustomFieldTypeChoices.TYPE_MULTISELECT):
                value = cf.resolve_selection_value(value)
            data[cf.name] = value

        return data

    def to_internal_value(self, data):
        if type(data) is not dict:
            raise ValidationError(
                "Invalid data format. Custom field data must be passed as a dictionary mapping field names to their "
                "values."
            )

        custom_fields = {cf.name: cf for cf in self._get_custom_fields()}

        # Reject any unknown custom field names
        invalid_fields = set(data) - set(custom_fields)
        if invalid_fields:
            raise ValidationError({
                field: _("Custom field '{name}' does not exist for this object type.").format(name=field)
                for field in sorted(invalid_fields)
            })

        # Serialize object and multi-object values
        for cf in custom_fields.values():
            if cf.name in data and data[cf.name] not in CUSTOMFIELD_EMPTY_VALUES and cf.type in (
                    CustomFieldTypeChoices.TYPE_OBJECT,
                    CustomFieldTypeChoices.TYPE_MULTIOBJECT
            ):
                serializer_class = get_serializer_for_model(cf.related_object_type.model_class())
                many = cf.type == CustomFieldTypeChoices.TYPE_MULTIOBJECT
                serializer = serializer_class(data=data[cf.name], nested=True, many=many, context=self.parent.context)
                if serializer.is_valid():
                    data[cf.name] = [obj['id'] for obj in serializer.data] if many else serializer.data['id']
                else:
                    raise ValidationError(_("Unknown related object(s): {name}").format(name=data[cf.name]))

            # Normalize URL values the same way the UI does (LaxURLField with assume_scheme='https'), so a
            # schemeless value (e.g. "example.com") is stored as an absolute URL ("https://example.com").
            # Malformed values are left untouched for CustomField.validate() to report.
            elif cf.type == CustomFieldTypeChoices.TYPE_URL and isinstance(data.get(cf.name), str) and data[cf.name]:
                try:
                    data[cf.name] = LaxURLField(assume_scheme='https').to_python(data[cf.name])
                except DjangoValidationError:
                    pass

        # If updating an existing instance, start with existing custom_field_data
        if self.parent.instance:
            data = {**self.parent.instance.custom_field_data, **data}

        return data


class CustomFieldListSerializer(ListSerializer):
    """
    ListSerializer that pre-fetches all OBJECT/MULTIOBJECT custom field related objects
    in bulk before per-item serialization.
    """
    def to_representation(self, data):
        cf_field = self.child.fields.get('custom_fields')
        if isinstance(cf_field, CustomFieldsDataField):
            object_type_cfs = [
                cf for cf in cf_field._get_custom_fields()
                if cf.type in (CustomFieldTypeChoices.TYPE_OBJECT, CustomFieldTypeChoices.TYPE_MULTIOBJECT)
            ]
            pks_by_model = {}
            for cf in object_type_cfs:
                pks = pks_by_model.setdefault(cf.related_object_type.model_class(), set())
                for item in data:
                    raw = item.custom_field_data.get(cf.name)
                    if raw is not None:
                        if cf.type == CustomFieldTypeChoices.TYPE_MULTIOBJECT:
                            pks.update(raw)
                        else:
                            pks.add(raw)
            cache = {}
            for model, pks in pks_by_model.items():
                objects = list(model.objects.filter(pk__in=pks))
                # A target model may lack a REST serializer, so plan only for found targets
                if objects:
                    serializer_class = get_serializer_for_model(model)
                    # Targets render with their nested (brief) representation
                    fields = getattr(serializer_class.Meta, 'brief_fields', None)
                    prefetch_related_objects(objects, *get_prefetches_for_serializer(serializer_class, fields=fields))
                for obj in objects:
                    cache[(model, obj.pk)] = obj
            self.child.context['cf_object_cache'] = cache
        return super().to_representation(data)