feat: DEV-2147: Calculate only requested fields (HumanSignal#2231)

* feat: DEV-2147: Calculate only requested fields * Review fixes (DEV-2147) * Remove unused function (DEV-2147) * Pretty fixes Co-authored-by: makseq-ubnt <[email protected]>
jotamotk · Jun 8, 2022 · cfa097c · cfa097c
1 parent 5550988
commit cfa097c
Show file tree

Hide file tree

Showing 5 changed files with 136 additions and 35 deletions.
diff --git a/label_studio/projects/api.py b/label_studio/projects/api.py
@@ -25,7 +25,7 @@
     Project, ProjectSummary, ProjectManager
 )
 from projects.serializers import (
-    ProjectSerializer, ProjectLabelConfigSerializer, ProjectSummarySerializer
+    ProjectSerializer, ProjectLabelConfigSerializer, ProjectSummarySerializer, GetFieldsSerializer
 )
 from projects.functions.next_task import get_next_task
 from tasks.models import Task
@@ -129,8 +129,11 @@ class ProjectListAPI(generics.ListCreateAPIView):
     pagination_class = ProjectListPagination
 
     def get_queryset(self):
+        serializer = GetFieldsSerializer(data=self.request.query_params)
+        serializer.is_valid(raise_exception=True)
+        fields = serializer.validated_data.get('include')
         projects = Project.objects.filter(organization=self.request.user.active_organization)
-        return ProjectManager.with_counts_annotate(projects).prefetch_related('members', 'created_by')
+        return ProjectManager.with_counts_annotate(projects, fields=fields).prefetch_related('members', 'created_by')
 
     def get_serializer_context(self):
         context = super(ProjectListAPI, self).get_serializer_context()
@@ -187,7 +190,10 @@ class ProjectAPI(generics.RetrieveUpdateDestroyAPIView):
     redirect_kwarg = 'pk'
 
     def get_queryset(self):
-        return Project.objects.with_counts().filter(organization=self.request.user.active_organization)
+        serializer = GetFieldsSerializer(data=self.request.query_params)
+        serializer.is_valid(raise_exception=True)
+        fields = serializer.validated_data.get('include')
+        return Project.objects.with_counts(fields=fields).filter(organization=self.request.user.active_organization)
 
     def get(self, request, *args, **kwargs):
         return super(ProjectAPI, self).get(request, *args, **kwargs)

diff --git a/label_studio/projects/functions/__init__.py b/label_studio/projects/functions/__init__.py
@@ -0,0 +1,52 @@
+from django.db.models import Count, Q
+
+
+def annotate_task_number(queryset):
+    return queryset.annotate(task_number=Count('tasks', distinct=True))
+
+
+def annotate_finished_task_number(queryset):
+    return queryset.annotate(finished_task_number=Count('tasks', distinct=True, filter=Q(tasks__is_labeled=True)))
+
+
+def annotate_total_predictions_number(queryset):
+    return queryset.annotate(total_predictions_number=Count('tasks__predictions', distinct=True))
+
+
+def annotate_total_annotations_number(queryset):
+    return queryset.annotate(total_annotations_number=Count(
+        'tasks__annotations__id', distinct=True, filter=Q(tasks__annotations__was_cancelled=False)
+    ))
+
+
+def annotate_num_tasks_with_annotations(queryset):
+    return queryset.annotate(num_tasks_with_annotations=Count(
+        'tasks__id',
+        distinct=True,
+        filter=Q(tasks__annotations__isnull=False)
+        & Q(tasks__annotations__ground_truth=False)
+        & Q(tasks__annotations__was_cancelled=False)
+        & Q(tasks__annotations__result__isnull=False),
+    ))
+
+
+def annotate_useful_annotation_number(queryset):
+    return queryset.annotate(useful_annotation_number=Count(
+        'tasks__annotations__id',
+        distinct=True,
+        filter=Q(tasks__annotations__was_cancelled=False)
+        & Q(tasks__annotations__ground_truth=False)
+        & Q(tasks__annotations__result__isnull=False),
+    ))
+
+
+def annotate_ground_truth_number(queryset):
+    return queryset.annotate(ground_truth_number=Count(
+        'tasks__annotations__id', distinct=True, filter=Q(tasks__annotations__ground_truth=True)
+    ))
+
+
+def annotate_skipped_annotations_number(queryset):
+    return queryset.annotate(skipped_annotations_number=Count(
+        'tasks__annotations__id', distinct=True, filter=Q(tasks__annotations__was_cancelled=True)
+    ))
diff --git a/label_studio/projects/models.py b/label_studio/projects/models.py
@@ -28,6 +28,11 @@
 )
 from core.bulk_update_utils import bulk_update
 from label_studio_tools.core.label_config import parse_config
+from projects.functions import (
+    annotate_task_number, annotate_finished_task_number, annotate_total_predictions_number,
+    annotate_total_annotations_number, annotate_num_tasks_with_annotations,
+    annotate_useful_annotation_number, annotate_ground_truth_number, annotate_skipped_annotations_number
+)
 from labels_manager.models import Label
 
 
@@ -49,40 +54,30 @@ def for_user(self, user):
         'skipped_annotations_number',
     ]
 
-    def with_counts(self):
-        return self.with_counts_annotate(self)
+    def with_counts(self, fields=None):
+        return self.with_counts_annotate(self, fields=fields)
 
     @staticmethod
-    def with_counts_annotate(queryset):
-        return queryset.annotate(
-            task_number=Count('tasks', distinct=True),
-            finished_task_number=Count('tasks', distinct=True, filter=Q(tasks__is_labeled=True)),
-            total_predictions_number=Count('tasks__predictions', distinct=True),
-            total_annotations_number=Count(
-                'tasks__annotations__id', distinct=True, filter=Q(tasks__annotations__was_cancelled=False)
-            ),
-            num_tasks_with_annotations=Count(
-                'tasks__id',
-                distinct=True,
-                filter=Q(tasks__annotations__isnull=False)
-                & Q(tasks__annotations__ground_truth=False)
-                & Q(tasks__annotations__was_cancelled=False)
-                & Q(tasks__annotations__result__isnull=False),
-            ),
-            useful_annotation_number=Count(
-                'tasks__annotations__id',
-                distinct=True,
-                filter=Q(tasks__annotations__was_cancelled=False)
-                & Q(tasks__annotations__ground_truth=False)
-                & Q(tasks__annotations__result__isnull=False),
-            ),
-            ground_truth_number=Count(
-                'tasks__annotations__id', distinct=True, filter=Q(tasks__annotations__ground_truth=True)
-            ),
-            skipped_annotations_number=Count(
-                'tasks__annotations__id', distinct=True, filter=Q(tasks__annotations__was_cancelled=True)
-            ),
-        )
+    def with_counts_annotate(queryset, fields=None):
+        available_fields = {
+            'task_number': annotate_task_number,
+            'finished_task_number': annotate_finished_task_number,
+            'total_predictions_number': annotate_total_predictions_number,
+            'total_annotations_number': annotate_total_annotations_number,
+            'num_tasks_with_annotations': annotate_num_tasks_with_annotations,
+            'useful_annotation_number': annotate_useful_annotation_number,
+            'ground_truth_number': annotate_ground_truth_number,
+            'skipped_annotations_number': annotate_skipped_annotations_number,
+        }
+        if fields is None:
+            to_annotate = available_fields
+        else:
+            to_annotate = {field: available_fields[field] for field in fields if field in available_fields}
+
+        for _, annotate_func in to_annotate.items():
+            queryset = annotate_func(queryset)
+
+        return queryset
 
 
 ProjectMixin = load_func(settings.PROJECT_MIXIN)

diff --git a/label_studio/projects/serializers.py b/label_studio/projects/serializers.py
@@ -113,3 +113,12 @@ class ProjectSummarySerializer(serializers.ModelSerializer):
     class Meta:
         model = ProjectSummary
         fields = '__all__'
+
+
+class GetFieldsSerializer(serializers.Serializer):
+    include = serializers.CharField(required=False)
+
+    def validate_include(self, value):
+        if value is not None:
+            value = value.split(',')
+        return value
diff --git a/label_studio/tests/test_projects.tavern.yml b/label_studio/tests/test_projects.tavern.yml
@@ -0,0 +1,39 @@
+---
+
+test_name: test_project_id
+
+strict: True
+
+marks:
+  - usefixtures:
+       - django_live_url
+
+stages:
+  - id: signup
+    type: ref
+  - id: get_user_token
+    type: ref
+  - id: create_project
+    name: Create project
+    request:
+      url: "{django_live_url}/api/projects"
+      json:
+        title: create_batch_tasks_assignments
+        label_config: <View><Text name="text" value="$text"/><Choices name="label" toName="text"><Choice value="pos"/><Choice value="neg"/></Choices></View>
+        is_published: true
+      method: POST
+      headers:
+        content-type: application/json
+    response:
+      status_code: 201
+      save:
+        json:
+          project_pk: id
+  - name: get_only_id
+    request:
+      method: GET
+      url: '{django_live_url}/api/projects/{project_pk}?include=id'
+    response:
+      status_code: 200
+      json:
+        id: !int "{project_pk}"