From 02e564ae44f582d6d39c9bf9c289f2c7d7578309 Mon Sep 17 00:00:00 2001
From: Sergei Ivashchenko <triklozoid@gmail.com>
Date: Tue, 5 Sep 2023 15:52:32 +0100
Subject: [PATCH 1/2] fix: LEAP-24: /api/tasks performance improvement

---
 label_studio/data_manager/api.py      | 29 +++++++++++++++++++++++++--
 label_studio/data_manager/managers.py |  9 +++++----
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/label_studio/data_manager/api.py b/label_studio/data_manager/api.py
index fa390badfd9d..b24906b07469 100644
--- a/label_studio/data_manager/api.py
+++ b/label_studio/data_manager/api.py
@@ -12,10 +12,12 @@
 from drf_yasg.utils import swagger_auto_schema
 from drf_yasg import openapi
 from django.conf import settings
+from asgiref.sync import sync_to_async, async_to_sync
 
 from core.utils.common import int_from_request, load_func
 from core.utils.params import bool_from_request
 from core.permissions import all_permissions, ViewClassPermission
+from core.feature_flags import flag_set
 from projects.models import Project
 from projects.serializers import ProjectSerializer
 from tasks.models import Task, Annotation, Prediction
@@ -130,11 +132,26 @@ class TaskPagination(PageNumberPagination):
     total_predictions = 0
     max_page_size = settings.TASK_API_PAGE_SIZE_MAX
 
-    def paginate_queryset(self, queryset, request, view=None):
+    @async_to_sync
+    async def async_paginate_queryset(self, queryset, request, view=None):
+        predictions_count_qs = Prediction.objects.filter(task_id__in=queryset)
+        self.total_predictions = await sync_to_async(predictions_count_qs.count, thread_sensitive=True)()
+
+        annotations_count_qs = Annotation.objects.filter(task_id__in=queryset, was_cancelled=False)
+        self.total_annotations = await sync_to_async(annotations_count_qs.count, thread_sensitive=True)()
+        return await sync_to_async(super().paginate_queryset, thread_sensitive=True)(queryset, request, view)
+
+    def sync_paginate_queryset(self, queryset, request, view=None):
         self.total_predictions = Prediction.objects.filter(task_id__in=queryset).count()
         self.total_annotations = Annotation.objects.filter(task_id__in=queryset, was_cancelled=False).count()
         return super().paginate_queryset(queryset, request, view)
 
+    def paginate_queryset(self, queryset, request, view=None):
+        if flag_set('fflag_fix_back_leap_24_tasks_api_optimization_05092023_short'):
+            return self.async_paginate_queryset(queryset, request, view)
+        else:
+            return self.sync_paginate_queryset(queryset, request, view)
+
     def get_paginated_response(self, data):
         return Response(
             {
@@ -236,7 +253,15 @@ def get(self, request):
                 evaluate_predictions(tasks_for_predictions)
                 [tasks_by_ids[_id].refresh_from_db() for _id in ids]
 
-            serializer = self.task_serializer_class(page, many=True, context=context)
+            if flag_set('fflag_fix_back_leap_24_tasks_api_optimization_05092023_short'):
+                serializer = self.task_serializer_class(
+                    page,
+                    many=True,
+                    context=context,
+                    include=get_fields_for_evaluation(prepare_params, request.user, skip_regular=False)
+                )
+            else:
+                serializer = self.task_serializer_class(page, many=True, context=context)
             return self.get_paginated_response(serializer.data)
         # all tasks
         if project.evaluate_predictions_automatically:
diff --git a/label_studio/data_manager/managers.py b/label_studio/data_manager/managers.py
index aedddf9377a3..b5a6e7e39700 100644
--- a/label_studio/data_manager/managers.py
+++ b/label_studio/data_manager/managers.py
@@ -81,7 +81,7 @@ def get_fields_for_filter_ordering(prepare_params):
     return result
 
 
-def get_fields_for_evaluation(prepare_params, user):
+def get_fields_for_evaluation(prepare_params, user, skip_regular=True):
     """ Collecting field names to annotate them
 
     :param prepare_params: structure with filters and ordering
@@ -111,9 +111,10 @@ def get_fields_for_evaluation(prepare_params, user):
     result = set(result)
 
     # we don't need to annotate regular model fields, so we skip them
-    skipped_fields = [field.attname for field in Task._meta.fields]
-    skipped_fields.append("id")
-    result = [f for f in result if f not in skipped_fields]
+    if skip_regular:
+        skipped_fields = [field.attname for field in Task._meta.fields]
+        skipped_fields.append("id")
+        result = [f for f in result if f not in skipped_fields]
     result = [f for f in result if not f.startswith("data.")]
 
     return result

From eb052e63d4b8bb1b14d4cb2a4545ac9857d56b99 Mon Sep 17 00:00:00 2001
From: Sergei Ivashchenko <triklozoid@gmail.com>
Date: Thu, 21 Sep 2023 15:30:53 +0100
Subject: [PATCH 2/2] Fix skip_regular

---
 label_studio/data_manager/managers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/label_studio/data_manager/managers.py b/label_studio/data_manager/managers.py
index 8aaacd3de923..a17376e5415a 100644
--- a/label_studio/data_manager/managers.py
+++ b/label_studio/data_manager/managers.py
@@ -114,7 +114,7 @@ def get_fields_for_evaluation(prepare_params, user, skip_regular=True):
         skipped_fields = [field.attname for field in Task._meta.fields]
         skipped_fields.append('id')
         result = [f for f in result if f not in skipped_fields]
-    result = [f for f in result if not f.startswith('data.')]
+        result = [f for f in result if not f.startswith('data.')]
 
     return result