oncall-engine/engine/common/api_helpers/custom_fields.py

import time

from django.core.exceptions import ObjectDoesNotExist
from rest_framework import fields, serializers
from rest_framework.exceptions import ValidationError
from rest_framework.relations import RelatedField

from apps.alerts.models import ChannelFilter
from apps.user_management.models import User
from common.api_helpers.exceptions import BadRequest


class OrganizationFilteredPrimaryKeyRelatedField(RelatedField):
    """
    This field is used to filter entities by organization
    """

    def __init__(self, **kwargs):
        self.filter_field = kwargs.pop("filter_field", "organization")
        self.display_func = kwargs.pop("display_func", lambda instance: str(instance))
        super().__init__(**kwargs)

    def to_representation(self, value):
        return value.public_primary_key

    def to_internal_value(self, data):
        try:
            return self.get_queryset().get(public_primary_key=data)
        except ObjectDoesNotExist:
            raise ValidationError("Object does not exist")
        except (TypeError, ValueError):
            raise ValidationError("Invalid values")

    def get_queryset(self):
        request = self.context.get("request", None)
        queryset = self.queryset
        if not request or not queryset:
            return None
        filter_kwargs = {self.filter_field: request.auth.organization}
        return queryset.filter(**filter_kwargs).distinct()

    def display_value(self, instance):
        return self.display_func(instance)


class TeamPrimaryKeyRelatedField(RelatedField):
    """
    This field is used to get user teams
    """

    def __init__(self, **kwargs):
        self.display_func = kwargs.pop("display_func", lambda instance: str(instance))
        super().__init__(**kwargs)

    def to_representation(self, value):
        return value.public_primary_key

    def to_internal_value(self, data):
        try:
            return self.get_queryset().get(public_primary_key=data)
        except ObjectDoesNotExist:
            raise ValidationError("Object does not exist")
        except (TypeError, ValueError):
            raise ValidationError("Invalid values")

    def get_queryset(self):
        request = self.context.get("request", None)
        if not request:
            return None
        return request.user.teams.all()

    def display_value(self, instance):
        return self.display_func(instance)


class UsersFilteredByOrganizationField(serializers.Field):
    """
    This field reduces queries count when accessing User many related field (ex: notify_to_users_queue).
    Check if you can use OrganizationFilteredPrimaryKeyRelatedField before using this one.
    """

    def __init__(self, **kwargs):
        self.queryset = kwargs.pop("queryset", None)
        super().__init__(**kwargs)

    def to_representation(self, value):
        return list(map(lambda v: v.public_primary_key, value.all()))

    def to_internal_value(self, data):
        queryset = self.queryset
        request = self.context.get("request", None)

        if not request or not queryset:
            return None

        return queryset.filter(organization=request.user.organization, public_primary_key__in=data).distinct()


class WritableSerializerMethodField(serializers.SerializerMethodField):
    """
    Please, NEVER use this field.
    It was a mistake to create this one due to necessity to dig deep in drf to fix bugs there.
    This field is a workaround to allow to write into SerializerMethodField.
    """

    def __init__(self, method_name=None, **kwargs):
        self.method_name = method_name
        self.setter_method_name = kwargs.pop("setter_method_name", None)
        self.deserializer_field = kwargs.pop("deserializer_field")

        kwargs["source"] = "*"
        super(serializers.SerializerMethodField, self).__init__(**kwargs)

    def bind(self, field_name, parent):
        retval = super().bind(field_name, parent)
        if not self.setter_method_name:
            self.setter_method_name = f"set_{field_name}"

        return retval

    def to_internal_value(self, data):
        value = self.deserializer_field.to_internal_value(data)
        method = getattr(self.parent, self.setter_method_name)
        method(value)
        return {self.method_name: value}


class CustomTimeField(fields.TimeField):
    def to_representation(self, value):
        result = super().to_representation(value)
        if result[-1] != "Z":
            result += "Z"
        return result

    def to_internal_value(self, data):
        TIME_FORMAT_LEN = len("00:00:00Z")
        if len(data) == TIME_FORMAT_LEN:
            try:
                time.strptime(data, "%H:%M:%SZ")
            except ValueError:
                raise BadRequest(detail="Invalid time format, should be '00:00:00Z'")
        else:
            raise BadRequest(detail="Invalid time format, should be '00:00:00Z'")
        return data


class RouteIdField(fields.CharField):
    def to_internal_value(self, data):
        try:
            channel_filter = ChannelFilter.objects.get(public_primary_key=data)
        except ChannelFilter.DoesNotExist:
            raise BadRequest(detail="Route does not exist")
        return channel_filter

    def to_representation(self, value):
        if value is not None:
            return value.public_primary_key
        return value


class UserIdField(fields.CharField):
    def to_internal_value(self, data):
        request = self.context.get("request", None)

        user = User.objects.filter(organization=request.auth.organization, public_primary_key=data).first()
        if user is None:
            raise BadRequest(detail="User does not exist")
        return user

    def to_representation(self, value):
        if value is not None:
            return value.public_primary_key
        return value