diff --git a/label_studio/projects/serializers.py b/label_studio/projects/serializers.py index 48273cedc914..742ddb394247 100644 --- a/label_studio/projects/serializers.py +++ b/label_studio/projects/serializers.py @@ -4,7 +4,7 @@ import bleach from constants import SAFE_HTML_ATTRIBUTES, SAFE_HTML_TAGS -from django.db.models import Q +from django.db.models import Exists, OuterRef, Q from drf_spectacular.utils import extend_schema_field, extend_schema_serializer from fsm.serializer_fields import FSMStateField from label_studio_sdk.label_interface import LabelInterface @@ -35,7 +35,7 @@ from rest_flex_fields import FlexFieldsModelSerializer from rest_framework import serializers from rest_framework.serializers import SerializerMethodField -from tasks.models import Task +from tasks.models import Annotation, Task from users.serializers import UserSimpleSerializer @@ -350,10 +350,8 @@ def update(self, instance, validated_data): return super().update(instance, validated_data) def get_queue_total(self, project) -> int: - remain = project.tasks.filter( - Q(is_labeled=False) & ~Q(annotations__completed_by_id=self.user_id) - | Q(annotations__completed_by_id=self.user_id) - ).distinct() + user_annotations = Annotation.objects.filter(task_id=OuterRef('pk'), completed_by_id=self.user_id) + remain = project.tasks.filter(Q(is_labeled=False) | Exists(user_annotations)) return remain.count() def get_queue_done(self, project) -> int: diff --git a/label_studio/projects/tests/test_serializers.py b/label_studio/projects/tests/test_serializers.py index b3ebe8624fb8..f19c771a8087 100644 --- a/label_studio/projects/tests/test_serializers.py +++ b/label_studio/projects/tests/test_serializers.py @@ -1,9 +1,14 @@ -"""Tests for projects.serializers (control_weights validation).""" +"""Tests for project serializers.""" import pytest +from django.db import connection from django.test import TestCase +from django.test.utils import CaptureQueriesContext from projects.serializers import ControlTagWeightSerializer, ProjectSerializer +from projects.tests.factories import ProjectFactory from rest_framework.exceptions import ValidationError +from tasks.tests.factories import AnnotationFactory, TaskFactory +from users.tests.factories import UserFactory class TestControlTagWeightSerializer(TestCase): @@ -84,7 +89,7 @@ def test_rejects_label_with_more_than_three_decimal_places(self): class TestProjectSerializer(TestCase): - """Validates the cross-field validate_control_weights on ProjectSerializer.""" + """Validates project serializer behavior, including weights and queue counts.""" def _validate(self, value): """Run validate_control_weights from a ProjectSerializer instance.""" @@ -124,3 +129,30 @@ def test_accepts_none(self): def test_accepts_empty_dict(self): """Empty dict passes through unchanged.""" assert self._validate({}) == {} + + def test_queue_total_uses_correlated_user_annotation_lookup(self): + project = ProjectFactory() + current_user = project.created_by + other_user = UserFactory() + TaskFactory(project=project, is_labeled=False) + current_user_task = TaskFactory(project=project, is_labeled=True) + other_user_task = TaskFactory(project=project, is_labeled=True) + unlabeled_other_user_task = TaskFactory(project=project, is_labeled=False) + mixed_user_task = TaskFactory(project=project, is_labeled=True) + AnnotationFactory(task=current_user_task, completed_by=current_user) + AnnotationFactory(task=other_user_task, completed_by=other_user) + AnnotationFactory(task=unlabeled_other_user_task, completed_by=other_user) + AnnotationFactory(task=mixed_user_task, completed_by=other_user) + AnnotationFactory(task=mixed_user_task, completed_by=current_user) + project.tasks.filter(pk__in=[current_user_task.pk, other_user_task.pk, mixed_user_task.pk]).update( + is_labeled=True + ) + serializer = ProjectSerializer(context={'user_cache': {current_user.id: current_user}}) + + with CaptureQueriesContext(connection) as captured_queries: + queue_total = serializer.get_queue_total(project) + + assert len(captured_queries) == 1 + assert queue_total == 4, captured_queries[0]['sql'] + assert 'EXISTS' in captured_queries[0]['sql'] + assert 'DISTINCT' not in captured_queries[0]['sql']