Merge branch 'devel' into AAP-77251

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
 

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()

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)}"