mirror of
https://github.com/ansible/awx.git
synced 2026-07-20 21:01:58 -02:30
AAP-82745 — Replace O(N) correlated subqueries in org list with bulk aggregation (#16548)
* AAP-82745 — Replace O(N) correlated subqueries in org list with bulk aggregation The org list endpoint computed per-org user and admin counts using two correlated subqueries against roleuserassignment — one per org per role type. At scale this was the #1 worst query in the customer DB (677ms mean, 336K calls, 3,801 min total DB time). Replace with a single flat conditional-aggregation query that scans roleuserassignment once regardless of org count. Also simplify the detail view to use direct .count() calls and add a query count regression test. * Remove dead code: org_counts is always populated after refactor * Address review feedback on detail view and test - Detail view: combine two .count() calls into single .aggregate() with conditional counts, matching the list-view pattern - Use .update() in else branch for consistency with if branch - Remove redundant setup_managed_roles fixture (transitive via organization_resource_creator) --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)}"
|
||||
|
||||
Reference in New Issue
Block a user