Added ExperimentData API endpoint for bulk upserts

This endpoint allows the Rapid Experiments Team to quickly update data for multiple learners.

lms/djangoapps/experiments/tests/test_views.py

@@ -149,3 +149,62 @@ class ExperimentDataViewSetTests(APITestCase):
         self.client.login(username=other_user.username, password=UserFactory._DEFAULT_PASSWORD)
         response = self.client.patch(url, data)
         self.assertEqual(response.status_code, 404)
+
+    def test_bulk_upsert_permissions(self):
+        """ Only staff users can access the bulk upsert endpoint. """
+        url = reverse('api_experiments:v0:data-bulk-upsert')
+        data = []
+
+        # Authentication is required
+        response = self.client.put(url, data, format='json')
+        self.assertEqual(response.status_code, 401)
+
+        user = UserFactory()
+        self.client.login(username=user.username, password=UserFactory._DEFAULT_PASSWORD)
+
+        # No access to non-staff users
+        response = self.client.put(url, data, format='json')
+        self.assertEqual(response.status_code, 403)
+
+        user.is_staff = True
+        user.save()
+        response = self.client.put(url, data, format='json')
+        self.assertEqual(response.status_code, 200)
+
+    def test_bulk_upsert(self):
+        """ The endpoint should support creating/updating multiple ExperimentData objects with a single call. """
+        url = reverse('api_experiments:v0:data-bulk-upsert')
+        experiment_id = 1
+        user = UserFactory(is_staff=True)
+        other_user = UserFactory()
+        self.client.login(username=user.username, password=UserFactory._DEFAULT_PASSWORD)
+
+        data = [
+            {
+                'experiment_id': experiment_id,
+                'key': 'foo',
+                'value': 'bar',
+                'user': user.username,
+            },
+            {
+                'experiment_id': experiment_id,
+                'key': 'foo',
+                'value': 'bar',
+                'user': other_user.username,
+            },
+        ]
+        response = self.client.put(url, data, format='json')
+        self.assertEqual(response.status_code, 200)
+        kwargs = {
+            'experiment_id': experiment_id,
+            'key': 'foo',
+            'value': 'bar',
+        }
+        ExperimentData.objects.get(user=user, **kwargs)
+        ExperimentData.objects.get(user=other_user, **kwargs)
+
+        # Subsequent calls should update the existing data rather than create more
+        response = self.client.put(url, data, format='json')
+        self.assertEqual(response.status_code, 200)
+        ExperimentData.objects.get(user=user, **kwargs)
+        ExperimentData.objects.get(user=other_user, **kwargs)

lms/djangoapps/experiments/views.py

+from django.contrib.auth import get_user_model
+from django.db import transaction
 from edx_rest_framework_extensions.authentication import JwtAuthentication
 from rest_framework import permissions, viewsets
+from rest_framework.decorators import list_route
 from rest_framework.filters import DjangoFilterBackend
+from rest_framework.response import Response
 
 from experiments import filters
 from experiments.models import ExperimentData
@@ -8,6 +12,8 @@ from experiments.permissions import IsStaffOrOwner
 from experiments.serializers import ExperimentDataCreateSerializer, ExperimentDataSerializer
 from openedx.core.lib.api.authentication import SessionAuthenticationAllowInactiveUser
 
+User = get_user_model()  # pylint: disable=invalid-name
+
 
 class ExperimentDataViewSet(viewsets.ModelViewSet):
     authentication_classes = (JwtAuthentication, SessionAuthenticationAllowInactiveUser,)
@@ -16,6 +22,7 @@ class ExperimentDataViewSet(viewsets.ModelViewSet):
     permission_classes = (permissions.IsAuthenticated, IsStaffOrOwner,)
     queryset = ExperimentData.objects.all()
     serializer_class = ExperimentDataSerializer
+    _cached_users = {}
 
     def filter_queryset(self, queryset):
         queryset = queryset.filter(user=self.request.user)
@@ -47,3 +54,31 @@ class ExperimentDataViewSet(viewsets.ModelViewSet):
 
         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)