From 507ba6a7787bcf584cced384ffac0dcdc75816c1 Mon Sep 17 00:00:00 2001
From: AlanCoding <arominge@redhat.com>
Date: Thu, 8 Sep 2016 15:03:55 -0400
Subject: [PATCH] add new case in prefetch method for foreign-key roles

---
 awx/api/views.py                              |  1 +
 .../functional/api/test_rbac_displays.py      | 22 +++++++++++++---
 awx/main/utils.py                             | 26 ++++++++++++++++---
 3 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/awx/api/views.py b/awx/api/views.py
index 579566766c..300ae3e209 100644
--- a/awx/api/views.py
+++ b/awx/api/views.py
@@ -1749,6 +1749,7 @@ class GroupList(ListCreateAPIView):
 
     model = Group
     serializer_class = GroupSerializer
+    capabilities_prefetch = ['inventory.admin', 'inventory.adhoc', 'inventory.update']
 
 class GroupChildrenList(SubListCreateAttachDetachAPIView):
 
diff --git a/awx/main/tests/functional/api/test_rbac_displays.py b/awx/main/tests/functional/api/test_rbac_displays.py
index c203e45b4b..5424d06616 100644
--- a/awx/main/tests/functional/api/test_rbac_displays.py
+++ b/awx/main/tests/functional/api/test_rbac_displays.py
@@ -4,10 +4,10 @@ from django.core.urlresolvers import reverse
 from django.test.client import RequestFactory
 
 from awx.main.models.jobs import JobTemplate
-from awx.main.models import Role
-from awx.api.serializers import JobTemplateSerializer
+from awx.main.models import Role, Group
 from awx.main.access import access_registry
-
+from awx.main.utils import cache_list_capabilities
+from awx.api.serializers import JobTemplateSerializer
 
 # This file covers special-cases of displays of user_capabilities
 # general functionality should be covered fully by unit tests, see:
@@ -253,10 +253,24 @@ def test_team_roles_unattach_functional(team, team_member, inventory, get):
 
 @pytest.mark.django_db
 def test_user_roles_unattach_functional(organization, alice, bob, get):
-    # Add to same organization so that alice and bob can see each other
     organization.member_role.members.add(alice)
     organization.member_role.members.add(bob)
     response = get(reverse('api:user_roles_list', args=(alice.id,)), bob)
     # Org members can not revoke the membership of other members
     assert response.data['results'][0]['summary_fields']['user_capabilities']['unattach'] == False
 
+
+@pytest.mark.django_db
+def test_prefetch_jt_capabilities(job_template, rando):
+    job_template.execute_role.members.add(rando)
+    qs = JobTemplate.objects.all()
+    cache_list_capabilities(qs, ['admin', 'execute'], JobTemplate, rando)
+    assert qs[0].capabilities_cache == {'edit': False, 'start': True}
+
+@pytest.mark.django_db
+def test_prefetch_group_capabilities(group, rando):
+    group.inventory.adhoc_role.members.add(rando)
+    qs = Group.objects.all()
+    cache_list_capabilities(qs, ['inventory.admin', 'inventory.adhoc'], Group, rando)
+    assert qs[0].capabilities_cache == {'edit': False, 'adhoc': True}
+
diff --git a/awx/main/utils.py b/awx/main/utils.py
index 06227c59e3..84cfa41a54 100644
--- a/awx/main/utils.py
+++ b/awx/main/utils.py
@@ -413,12 +413,27 @@ def cache_list_capabilities(page, role_types, model, user):
     '''
     Given a `page` list of objects, the specified roles for the specified user
     are save on each object in the list, using 1 query for each role type
+
+    Examples:
+    capabilities_prefetch = ['admin', 'execute']
+      --> prefetch the admin (edit) and execute (start) permissions for
+          items in list for current user
+    capabilities_prefetch = ['inventory.admin_role']
+      --> prefetch the related inventory FK permissions for current user,
+          and put it into the object's cache
     '''
     page_ids = [obj.id for obj in page]
     for obj in page:
         obj.capabilities_cache = {}
 
-    for role_type in role_types:
+    for role_path in role_types:
+        if '.' in role_path:
+            path = '__'.join(role_path.split('.')[:-1])
+            role_type = role_path.split('.')[-1]
+        else:
+            path = None
+            role_type = role_path
+
         # Role name translation to UI names for methods
         display_method = role_type
         if role_type == 'admin':
@@ -427,8 +442,13 @@ def cache_list_capabilities(page, role_types, model, user):
             display_method = 'start'
 
         # Query for union of page objects & role accessible_objects
-        ids_with_role = set(model.accessible_objects(
-            user, '%s_role' % role_type).filter(pk__in=page_ids).values_list('pk', flat=True))
+        if path:
+            parent_model = model._meta.get_field(path).related_model
+            kwargs = {'%s__in' % path: parent_model.accessible_objects(user, '%s_role' % role_type)}
+            qs_obj = model.objects.filter(**kwargs)
+        else:
+            qs_obj = model.accessible_objects(user, '%s_role' % role_type)
+        ids_with_role = set(qs_obj.filter(pk__in=page_ids).values_list('pk', flat=True))
 
         # Save data item-by-item
         for obj in page: