From cf9a32341ea7dca8f914ee6b69bb94847ae2f66a Mon Sep 17 00:00:00 2001 From: Oliver Allweyer Date: Tue, 26 May 2026 14:36:03 +0200 Subject: [PATCH] fix: add select_related/prefetch_related to UserAPI to fix N+1 queries Co-Authored-By: Claude Opus 4.6 --- label_studio/users/api.py | 6 +- label_studio/users/tests/test_users_api.py | 66 ++++++++++++++++++++++ 2 files changed, 71 insertions(+), 1 deletion(-) create mode 100644 label_studio/users/tests/test_users_api.py diff --git a/label_studio/users/api.py b/label_studio/users/api.py index e64f4452102e..503a9b3be20d 100644 --- a/label_studio/users/api.py +++ b/label_studio/users/api.py @@ -179,7 +179,11 @@ class UserAPI(viewsets.ModelViewSet): http_method_names = ['get', 'post', 'head', 'patch', 'delete'] def get_queryset(self): - return User.objects.filter(organizations=self.request.user.active_organization) + return ( + User.objects.filter(organizations=self.request.user.active_organization) + .select_related('active_organization', 'active_organization__created_by') + .prefetch_related('om_through') + ) @extend_schema(exclude=True) @action(detail=True, methods=['delete', 'post'], permission_required=all_permissions.avatar_any) diff --git a/label_studio/users/tests/test_users_api.py b/label_studio/users/tests/test_users_api.py new file mode 100644 index 000000000000..872d2fb46d15 --- /dev/null +++ b/label_studio/users/tests/test_users_api.py @@ -0,0 +1,66 @@ +from django.test import TestCase, override_settings +from django.urls import reverse +from organizations.models import Organization, OrganizationMember +from rest_framework import status +from rest_framework.test import APIClient +from users.tests.factories import UserFactory + + +class UserListQueryOptimizationTest(TestCase): + """Test that GET /api/users uses select_related/prefetch_related to avoid N+1 queries.""" + + def setUp(self): + self.client = APIClient() + self.owner = UserFactory() + self.org = Organization.create_organization(created_by=self.owner, title='Test Org') + self.owner.active_organization = self.org + self.owner.save(update_fields=['active_organization']) + OrganizationMember.objects.get_or_create(user=self.owner, organization=self.org) + self.client.force_authenticate(user=self.owner) + + def _create_users(self, count: int) -> list: + return [UserFactory(active_organization=self.org) for _ in range(count)] + + def test_list_query_count_does_not_scale_with_user_count(self): + """The number of queries should be constant regardless of how many users are in the org.""" + self._create_users(5) + with self.assertNumQueries(3) as ctx_small: + response_small = self.client.get(reverse('user-list')) + self.assertEqual(response_small.status_code, status.HTTP_200_OK) + baseline = len(ctx_small.captured_queries) + + self._create_users(20) + with self.assertNumQueries(baseline): + response_large = self.client.get(reverse('user-list')) + self.assertEqual(response_large.status_code, status.HTTP_200_OK) + self.assertEqual(len(response_large.data), 26) + + def test_list_returns_active_organization_meta(self): + """Verify that active_organization_meta is populated correctly with prefetch.""" + self._create_users(3) + response = self.client.get(reverse('user-list')) + self.assertEqual(response.status_code, status.HTTP_200_OK) + for user_data in response.data: + self.assertIn('active_organization_meta', user_data) + meta = user_data['active_organization_meta'] + self.assertIn('title', meta) + self.assertIn('email', meta) + + def test_list_deleted_member_shown_as_deleted(self): + """Verify that soft-deleted org members are shown as 'Deleted User' with prefetch.""" + users = self._create_users(2) + om = OrganizationMember.objects.get(user=users[0], organization=self.org) + om.soft_delete() + + response = self.client.get(reverse('user-list')) + self.assertEqual(response.status_code, status.HTTP_200_OK) + deleted_user_data = next(u for u in response.data if u['id'] == users[0].id) + self.assertEqual(deleted_user_data['first_name'], 'Deleted') + self.assertEqual(deleted_user_data['last_name'], 'User') + + def test_retrieve_still_works(self): + """Verify that GET /api/users/ still works with the optimized queryset.""" + users = self._create_users(1) + response = self.client.get(reverse('user-detail', args=[users[0].id])) + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.data['id'], users[0].id)