From 9f0f00a59496c57c84aa8cf553150f4fb6f92f37 Mon Sep 17 00:00:00 2001
From: Tom Hacohen <tom@stosb.com>
Date: Thu, 16 Apr 2020 17:36:06 +0300
Subject: [PATCH] Unify the stoken filtering and calculation.

---
 django_etesync/views.py | 45 ++++++++++++++++++-----------------------
 1 file changed, 20 insertions(+), 25 deletions(-)

diff --git a/django_etesync/views.py b/django_etesync/views.py
index 53e6adb..25f2d2a 100644
--- a/django_etesync/views.py
+++ b/django_etesync/views.py
@@ -41,6 +41,7 @@ User = get_user_model()
 class BaseViewSet(viewsets.ModelViewSet):
     authentication_classes = tuple(app_settings.API_AUTHENTICATORS)
     permission_classes = tuple(app_settings.API_PERMISSIONS)
+    stoken_id_field = None
 
     def get_serializer_class(self):
         serializer_class = self.serializer_class
@@ -54,6 +55,22 @@ class BaseViewSet(viewsets.ModelViewSet):
         user = self.request.user
         return queryset.filter(members__user=user)
 
+    def filter_by_stoken_and_limit(self, request, queryset):
+        stoken = request.GET.get('stoken', None)
+        limit = int(request.GET.get('limit', 50))
+
+        stoken_id_field = self.stoken_id_field + '__id'
+
+        if stoken is not None:
+            last_rev = get_object_or_404(CollectionItemRevision.objects.all(), uid=stoken)
+            filter_by = {stoken_id_field + '__gt': last_rev.id}
+            queryset = queryset.filter(**filter_by)
+
+        new_stoken_id = queryset.aggregate(stoken_id=Max(stoken_id_field))['stoken_id']
+        new_stoken = CollectionItemRevision.objects.get(id=new_stoken_id).uid if new_stoken_id is not None else stoken
+
+        return queryset[:limit], new_stoken
+
 
 class CollectionViewSet(BaseViewSet):
     allowed_methods = ['GET', 'POST', 'DELETE']
@@ -61,6 +78,7 @@ class CollectionViewSet(BaseViewSet):
     queryset = Collection.objects.all()
     serializer_class = CollectionSerializer
     lookup_field = 'uid'
+    stoken_id_field = 'items__revisions'
 
     def get_queryset(self, queryset=None):
         if queryset is None:
@@ -91,22 +109,11 @@ class CollectionViewSet(BaseViewSet):
         return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
 
     def list(self, request):
-        stoken = request.GET.get('stoken', None)
-        limit = int(request.GET.get('limit', 50))
-
         queryset = self.get_queryset()
-
-        if stoken is not None:
-            last_rev = get_object_or_404(CollectionItemRevision.objects.all(), uid=stoken)
-            queryset = queryset.filter(items__revisions__id__gt=last_rev.id)
-
-        queryset = queryset[:limit]
+        queryset, new_stoken = self.filter_by_stoken_and_limit(request, queryset)
 
         serializer = self.serializer_class(queryset, context=self.get_serializer_context(), many=True)
 
-        new_stoken_id = queryset.aggregate(stoken_id=Max('items__revisions__id'))['stoken_id']
-        new_stoken = CollectionItemRevision.objects.get(id=new_stoken_id).uid if new_stoken_id is not None else stoken
-
         ret = {
             'data': serializer.data,
         }
@@ -119,6 +126,7 @@ class CollectionItemViewSet(BaseViewSet):
     queryset = CollectionItem.objects.all()
     serializer_class = CollectionItemSerializer
     lookup_field = 'uid'
+    stoken_id_field = 'revisions'
 
     def get_queryset(self):
         collection_uid = self.kwargs['collection_uid']
@@ -201,19 +209,6 @@ class CollectionItemViewSet(BaseViewSet):
         }
         return Response(ret, headers={'X-EteSync-SToken': new_stoken})
 
-    def filter_by_stoken_and_limit(self, request, queryset):
-        stoken = request.GET.get('stoken', None)
-        limit = int(request.GET.get('limit', 50))
-
-        if stoken is not None:
-            last_rev = get_object_or_404(CollectionItemRevision.objects.all(), uid=stoken)
-            queryset = queryset.filter(revisions__id__gt=last_rev.id)
-
-        new_stoken_id = queryset.aggregate(stoken_id=Max('revisions__id'))['stoken_id']
-        new_stoken = CollectionItemRevision.objects.get(id=new_stoken_id).uid if new_stoken_id is not None else stoken
-
-        return queryset[:limit], new_stoken
-
 
 class CollectionItemChunkViewSet(viewsets.ViewSet):
     allowed_methods = ['GET', 'POST']