diff --git a/awx/api/views/mixin.py b/awx/api/views/mixin.py index 24b5f859ab..1f114cb675 100644 --- a/awx/api/views/mixin.py +++ b/awx/api/views/mixin.py @@ -4,8 +4,8 @@ import dateutil import logging -from django.db.models import Count, OuterRef, Subquery, TextField -from django.db.models.functions import Cast, Coalesce +from django.db.models import Count, Q, TextField +from django.db.models.functions import Cast from django.db import transaction from django.shortcuts import get_object_or_404 from django.utils.timezone import now @@ -179,48 +179,37 @@ class OrganizationCountsMixin(object): db_results['projects'] = project_qs.values('organization').annotate(Count('organization')).order_by('organization') - member_rd = RoleDefinition.objects.filter(name='Organization Member').first() - admin_rd = RoleDefinition.objects.filter(name='Organization Admin').first() - - if member_rd and admin_rd: - - def assignment_count(rd): - return Coalesce( - Subquery( - RoleUserAssignment.objects.filter( - object_id=Cast(OuterRef('pk'), output_field=TextField()), - role_definition=rd, - ) - .values('role_definition') - .annotate(c=Count('pk')) - .values('c') - ), - 0, - ) - - db_results['users'] = org_qs.annotate( - users=assignment_count(member_rd), - admins=assignment_count(admin_rd), - ).values('id', 'users', 'admins') - count_context = {} for org in org_id_list: org_id = org['id'] count_context[org_id] = {'inventories': 0, 'teams': 0, 'users': 0, 'job_templates': 0, 'admins': 0, 'projects': 0} for res, count_qs in db_results.items(): - if res == 'users': - org_reference = 'id' - else: - org_reference = 'organization' for entry in count_qs: - org_id = entry[org_reference] + org_id = entry['organization'] if org_id in count_context: - if res == 'users': - count_context[org_id]['admins'] = entry['admins'] - count_context[org_id]['users'] = entry['users'] - continue - count_context[org_id][res] = entry['%s__count' % org_reference] + count_context[org_id][res] = entry['organization__count'] + + member_rd = RoleDefinition.objects.filter(name='Organization Member').first() + admin_rd = RoleDefinition.objects.filter(name='Organization Admin').first() + + if member_rd and admin_rd: + user_admin_counts = ( + RoleUserAssignment.objects.filter( + role_definition__in=[member_rd, admin_rd], + object_id__in=org_qs.annotate(text_pk=Cast('pk', TextField())).values('text_pk'), + ) + .values('object_id') + .annotate( + users=Count('pk', filter=Q(role_definition=member_rd)), + admins=Count('pk', filter=Q(role_definition=admin_rd)), + ) + ) + for entry in user_admin_counts: + org_id = int(entry['object_id']) + if org_id in count_context: + count_context[org_id]['users'] = entry['users'] + count_context[org_id]['admins'] = entry['admins'] full_context['related_field_counts'] = count_context diff --git a/awx/api/views/organization.py b/awx/api/views/organization.py index 24054383f1..2e4ef01dd7 100644 --- a/awx/api/views/organization.py +++ b/awx/api/views/organization.py @@ -5,8 +5,7 @@ import logging # Django -from django.db.models import Count, OuterRef, Subquery, TextField -from django.db.models.functions import Cast, Coalesce +from django.db.models import Count, Q from django.contrib.contenttypes.models import ContentType from django.utils.translation import gettext_lazy as _ @@ -83,37 +82,17 @@ class OrganizationDetail(RelatedJobsPreventDeleteMixin, RetrieveUpdateDestroyAPI admin_rd = RoleDefinition.objects.filter(name='Organization Admin').first() if member_rd and admin_rd: - - def assignment_count(rd): - return Coalesce( - Subquery( - RoleUserAssignment.objects.filter( - object_id=Cast(OuterRef('pk'), output_field=TextField()), - role_definition=rd, - ) - .values('role_definition') - .annotate(c=Count('pk')) - .values('c') - ), - 0, - ) - - direct_counts = ( - Organization.objects.filter(id=org_id) - .annotate( - users=assignment_count(member_rd), - admins=assignment_count(admin_rd), - ) - .values('users', 'admins') + counts = RoleUserAssignment.objects.filter( + role_definition__in=[member_rd, admin_rd], + object_id=str(org_id), + ).aggregate( + users=Count('pk', filter=Q(role_definition=member_rd)), + admins=Count('pk', filter=Q(role_definition=admin_rd)), ) - - if direct_counts: - org_counts = direct_counts[0] + org_counts.update(counts) else: - org_counts = {'users': 0, 'admins': 0} + org_counts.update({'users': 0, 'admins': 0}) - if not org_counts: - return full_context org_counts['inventories'] = Inventory.accessible_objects(**access_kwargs).filter(organization__id=org_id).count() org_counts['teams'] = Team.accessible_objects(**access_kwargs).filter(organization__id=org_id).count() org_counts['projects'] = Project.accessible_objects(**access_kwargs).filter(organization__id=org_id).count() diff --git a/awx/main/tests/functional/api/test_organization_counts.py b/awx/main/tests/functional/api/test_organization_counts.py index f647cf7d53..1c20f2fc71 100644 --- a/awx/main/tests/functional/api/test_organization_counts.py +++ b/awx/main/tests/functional/api/test_organization_counts.py @@ -1,5 +1,8 @@ import pytest +from django.test.utils import CaptureQueriesContext +from django.db import connection + from ansible_base.rbac.models import RoleDefinition from awx.api.versioning import reverse @@ -213,3 +216,24 @@ def test_JT_not_double_counted(resourced_organization, user, get): assert 'hosts' in counts counts.pop('hosts') assert counts == counts_dict + + +@pytest.mark.django_db +def test_org_list_user_admin_query_count(organization_resource_creator, organizations, user, get): + """User/admin counts use O(1) queries against roleuserassignment, not O(N) correlated subqueries.""" + admin_user = user('admin', True) + extra_orgs = organizations(4) + member_rd = RoleDefinition.objects.get(name='Organization Member') + admin_rd = RoleDefinition.objects.get(name='Organization Admin') + for org in extra_orgs: + for i in range(3): + member_rd.give_permission(user(f'member-{org.pk}-{i}'), org) + admin_rd.give_permission(user(f'admin-{org.pk}'), org) + + with CaptureQueriesContext(connection) as ctx: + response = get(reverse('api:organization_list'), admin_user) + + assert response.status_code == 200 + + rua_queries = [q for q in ctx.captured_queries if 'dab_rbac_roleuserassignment' in q['sql']] + assert len(rua_queries) <= 2, f"Expected at most 2 roleuserassignment queries, got {len(rua_queries)}"