Optimize lookup of org membership on prototype and perms APIs

Fixes a major slowdown when working with permissions under organizations with a lot of members

Fixes https://www.pivotaltracker.com/story/show/144076113

data/model/organization.py

@@ -106,15 +106,28 @@ def remove_organization_member(org, user_obj):
       TeamMember.delete().where(TeamMember.id << members).execute()
 
 
-def get_organization_member_set(orgname):
+def get_organization_member_set(org, include_robots=False, users_filter=None):
+  """ Returns the set of all member usernames under the given organization, with optional
+      filtering by robots and/or by a specific set of User objects.
+  """
   Org = User.alias()
   org_users = (User
                .select(User.username)
                .join(TeamMember)
                .join(Team)
-               .join(Org, on=(Org.id == Team.organization))
-               .where(Org.username == orgname)
+               .where(Team.organization == org)
                .distinct())
+
+  if not include_robots:
+    org_users = org_users.where(User.robot == False)
+
+  if users_filter is not None:
+    ids_list = [u.id for u in users_filter if u is not None]
+    if not ids_list:
+      return set()
+
+    org_users = org_users.where(User.id << ids_list)
+
   return {user.username for user in org_users}
 
 

endpoints/api/permission.py

@@ -73,6 +73,9 @@ class RepositoryUserPermissionList(RepositoryParamResource):
       # This repository isn't under an org
       pass
 
+    # Load the permissions.
+    repo_perms = model.user.get_all_repo_users(namespace, repository)
+
     # Determine how to wrap the role(s).
     def wrapped_role_view(repo_perm):
       return wrap_role_view_user(role_view(repo_perm), repo_perm.user)
@@ -80,20 +83,17 @@ class RepositoryUserPermissionList(RepositoryParamResource):
     role_view_func = wrapped_role_view
 
     if org:
-      org_members = model.organization.get_organization_member_set(namespace)
+      users_filter = {perm.user for perm in repo_perms}
+      org_members = model.organization.get_organization_member_set(org, users_filter=users_filter)
       current_func = role_view_func
 
       def wrapped_role_org_view(repo_perm):
-        return wrap_role_view_org(current_func(repo_perm), repo_perm.user,
-                                  org_members)
+        return wrap_role_view_org(current_func(repo_perm), repo_perm.user, org_members)
 
       role_view_func = wrapped_role_org_view
 
-    # Load and return the permissions.
-    repo_perms = model.user.get_all_repo_users(namespace, repository)
     return {
-      'permissions': {perm.user.username: role_view_func(perm)
-                      for perm in repo_perms}
+      'permissions': {perm.user.username: role_view_func(perm) for perm in repo_perms}
     }
 
 
@@ -156,8 +156,8 @@ class RepositoryUserPermission(RepositoryParamResource):
     perm_view = wrap_role_view_user(role_view(perm), perm.user)
 
     try:
-      model.organization.get_organization(namespace)
-      org_members = model.organization.get_organization_member_set(namespace)
+      org = model.organization.get_organization(namespace)
+      org_members = model.organization.get_organization_member_set(org, users_filter={perm.user})
       perm_view = wrap_role_view_org(perm_view, perm.user, org_members)
     except model.InvalidOrganizationException:
       # This repository is not part of an organization
@@ -183,8 +183,8 @@ class RepositoryUserPermission(RepositoryParamResource):
     perm_view = wrap_role_view_user(role_view(perm), perm.user)
 
     try:
-      model.organization.get_organization(namespace)
-      org_members = model.organization.get_organization_member_set(namespace)
+      org = model.organization.get_organization(namespace)
+      org_members = model.organization.get_organization_member_set(org, users_filter={perm.user})
       perm_view = wrap_role_view_org(perm_view, perm.user, org_members)
     except model.InvalidOrganizationException:
       # This repository is not part of an organization

endpoints/api/prototype.py

@@ -133,7 +133,10 @@ class PermissionPrototypeList(ApiResource):
         raise NotFound()
 
       permissions = model.permission.get_prototype_permissions(org)
-      org_members = model.organization.get_organization_member_set(orgname)
+
+      users_filter = ({p.activating_user for p in permissions} |
+                      {p.delegate_user for p in permissions})
+      org_members = model.organization.get_organization_member_set(org, users_filter=users_filter)
       return {'prototypes': [prototype_view(p, org_members)  for p in permissions]}
 
     raise Unauthorized()
@@ -180,7 +183,9 @@ class PermissionPrototypeList(ApiResource):
       prototype = model.permission.add_prototype_permission(org, role_name, activating_user,
                                                             delegate_user, delegate_team)
       log_prototype_action('create_prototype_permission', orgname, prototype)
-      org_members = model.organization.get_organization_member_set(orgname)
+
+      users_filter = {prototype.activating_user, prototype.delegate_user}
+      org_members = model.organization.get_organization_member_set(org, users_filter=users_filter)
       return prototype_view(prototype, org_members)
 
     raise Unauthorized()
@@ -257,7 +262,9 @@ class PermissionPrototype(ApiResource):
 
       log_prototype_action('modify_prototype_permission', orgname, prototype,
                            original_role=existing.role.name)
-      org_members = model.organization.get_organization_member_set(orgname)
+
+      users_filter = {prototype.activating_user, prototype.delegate_user}
+      org_members = model.organization.get_organization_member_set(org, users_filter=users_filter)
       return prototype_view(prototype, org_members)
 
     raise Unauthorized()