views.py

from django.contrib.auth import get_user_model
from django.db import transaction
from django_filters.rest_framework import DjangoFilterBackend
from edx_rest_framework_extensions.authentication import JwtAuthentication
from rest_framework import permissions, viewsets
from rest_framework.decorators import list_route
from rest_framework.response import Response

from experiments import filters, serializers
from experiments.models import ExperimentData, ExperimentKeyValue
from experiments.permissions import IsStaffOrOwner, IsStaffOrReadOnly
from openedx.core.lib.api.authentication import SessionAuthenticationAllowInactiveUser

User = get_user_model()  # pylint: disable=invalid-name


class ExperimentDataViewSet(viewsets.ModelViewSet):
    authentication_classes = (JwtAuthentication, SessionAuthenticationAllowInactiveUser,)
    filter_backends = (DjangoFilterBackend,)
    filter_class = filters.ExperimentDataFilter
    permission_classes = (permissions.IsAuthenticated, IsStaffOrOwner,)
    queryset = ExperimentData.objects.all()
    serializer_class = serializers.ExperimentDataSerializer
    _cached_users = {}

    def filter_queryset(self, queryset):
        queryset = queryset.filter(user=self.request.user)
        return super(ExperimentDataViewSet, self).filter_queryset(queryset)

    def get_serializer_class(self):
        if self.action == 'create':
            return serializers.ExperimentDataCreateSerializer
        return serializers.ExperimentDataSerializer

    def create_or_update(self, request, *args, **kwargs):
        # If we have a primary key, treat this as a regular update request
        if self.kwargs.get('pk'):
            return self.update(request, *args, **kwargs)

        # If we only have data, check to see if an instance exists in the database. If so, update it.
        # Otherwise, create a new instance.
        experiment_id = request.data.get('experiment_id')
        key = request.data.get('key')

        if experiment_id and key:
            try:
                obj = self.get_queryset().get(user=self.request.user, experiment_id=experiment_id, key=key)
                self.kwargs['pk'] = obj.pk
                self.request.data['id'] = obj.pk
                return self.update(request, *args, **kwargs)
            except ExperimentData.DoesNotExist:
                pass

        self.action = 'create'
        return self.create(request, *args, **kwargs)

    def _cache_users(self, usernames):
        users = User.objects.filter(username__in=usernames)
        self._cached_users = {user.username: user for user in users}

    def _get_user(self, username):
        user = self._cached_users.get(username)

        if not user:
            user = User.objects.get(username=username)
            self._cached_users[username] = user

        return user

    @list_route(methods=['put'], permission_classes=[permissions.IsAdminUser])
    def bulk_upsert(self, request):
        upserted = []
        self._cache_users([datum['user'] for datum in request.data])

        with transaction.atomic():
            for item in request.data:
                user = self._get_user(username=item['user'])
                datum, __ = ExperimentData.objects.update_or_create(
                    user=user, experiment_id=item['experiment_id'], key=item['key'], defaults={'value': item['value']})
                upserted.append(datum)

            serializer = self.get_serializer(upserted, many=True)
            return Response(serializer.data)


class ExperimentKeyValueViewSet(viewsets.ModelViewSet):
    authentication_classes = (JwtAuthentication, SessionAuthenticationAllowInactiveUser,)
    filter_backends = (DjangoFilterBackend,)
    filter_class = filters.ExperimentKeyValueFilter
    permission_classes = (IsStaffOrReadOnly,)
    queryset = ExperimentKeyValue.objects.all()
    serializer_class = serializers.ExperimentKeyValueSerializer

    @list_route(methods=['put'], permission_classes=[permissions.IsAdminUser])
    def bulk_upsert(self, request):
        upserted = []

        with transaction.atomic():
            for item in request.data:
                datum, __ = ExperimentKeyValue.objects.update_or_create(
                    experiment_id=item['experiment_id'], key=item['key'], defaults={'value': item['value']})
                upserted.append(datum)

            serializer = self.get_serializer(upserted, many=True)
            return Response(serializer.data)