Switched to a nicer contextmanager implemenation for role hierarchy rebuild batching

#1206

awx/main/fields.py

@@ -23,7 +23,7 @@ from django.db.transaction import TransactionManagementError
 
 
 # AWX
-from awx.main.models.rbac import RolePermission, Role
+from awx.main.models.rbac import RolePermission, Role, batch_role_ancestor_rebuilding
 
 
 __all__ = ['AutoOneToOneField', 'ImplicitRoleField']
@@ -280,12 +280,11 @@ class ImplicitRoleField(models.ForeignKey):
             for parent in parents:
                 new_parents.add(parent)
 
-        Role.pause_role_ancestor_rebuilding()
+        with batch_role_ancestor_rebuilding():
             for role in original_parents - new_parents:
                 this_role.parents.remove(role)
             for role in new_parents - original_parents:
                 this_role.parents.add(role)
-        Role.unpause_role_ancestor_rebuilding()
 
         setattr(self, '__original_parent_roles', new_parents)
 

awx/main/management/commands/generate_dummy_data.py

@@ -104,7 +104,7 @@ class Command(BaseCommand):
         try:
 
             with transaction.atomic():
-                Role.pause_role_ancestor_rebuilding()
+                with batch_role_ancestor_rebuilding():
 
                     print('# Creating %d organizations' % n_organizations)
                     for i in xrange(n_organizations):
@@ -344,10 +344,6 @@ class Command(BaseCommand):
                         if n:
                             print('')
 
-
-
-                Role.unpause_role_ancestor_rebuilding()
-
                 if options['pretend']:
                     raise Rollback()
         except Rollback:

awx/main/models/rbac.py

@@ -3,9 +3,11 @@
 
 # Python
 import logging
+import threading
+import contextlib
 
 # Django
-from django.db import models
+from django.db import models, transaction
 from django.db.models import Q
 from django.db.models.aggregates import Max
 from django.core.urlresolvers import reverse
@@ -19,6 +21,7 @@ from awx.main.models.base import * # noqa
 __all__ = [
     'Role',
     'RolePermission',
+    'batch_role_ancestor_rebuilding',
     'get_user_permissions_on_resource',
     'get_role_permissions_on_resource',
     'ROLE_SINGLETON_SYSTEM_ADMINISTRATOR',
@@ -30,12 +33,43 @@ logger = logging.getLogger('awx.main.models.rbac')
 ROLE_SINGLETON_SYSTEM_ADMINISTRATOR='System Administrator'
 ROLE_SINGLETON_SYSTEM_AUDITOR='System Auditor'
 
-role_rebuilding_paused = False
-roles_needing_rebuilding = set()
-
 ALL_PERMISSIONS = {'create': True, 'read': True, 'update': True, 'delete': True,
                    'write': True, 'scm_update': True, 'use': True, 'execute': True}
 
+
+tls = threading.local() # thread local storage
+
+@contextlib.contextmanager
+def batch_role_ancestor_rebuilding(allow_nesting=False):
+    '''
+    Batches the role ancestor rebuild work necessary whenever role-role
+    relations change. This can result in a big speedup when performing
+    any bulk manipulation.
+
+    WARNING: Calls to anything related to checking access/permissions
+    while within the context of the batch_role_ancestor_rebuilding will
+    likely not work.
+    '''
+
+    batch_role_rebuilding = getattr(tls, 'batch_role_rebuilding', False)
+
+    try:
+        setattr(tls, 'batch_role_rebuilding', True)
+        if not batch_role_rebuilding:
+            setattr(tls, 'roles_needing_rebuilding', set())
+        yield
+
+    finally:
+        setattr(tls, 'batch_role_rebuilding', batch_role_rebuilding)
+        if not batch_role_rebuilding:
+            rebuild_set = getattr(tls, 'roles_needing_rebuilding')
+            with transaction.atomic():
+                for role in Role.objects.filter(id__in=list(rebuild_set)).all():
+                    # TODO: We can reduce this to one rebuild call with our new upcoming rebuild method.. do this
+                    role.rebuild_role_ancestor_list()
+            delattr(tls, 'roles_needing_rebuilding')
+
+
 class Role(CommonModelNameNotUnique):
     '''
     Role model
@@ -61,35 +95,6 @@ class Role(CommonModelNameNotUnique):
     def get_absolute_url(self):
         return reverse('api:role_detail', args=(self.pk,))
 
-    @staticmethod
-    def pause_role_ancestor_rebuilding():
-        '''
-        Pauses role ancestor list updating. This is useful when you're making
-        many changes to the same roles, for example doing bulk inserts or
-        making many changes to the same object in succession.
-
-        Note that the unpause_role_ancestor_rebuilding MUST be called within
-        the same execution context (preferably within the same transaction),
-        otherwise the RBAC role ancestor hierarchy will not be properly
-        updated.
-        '''
-
-        global role_rebuilding_paused
-        role_rebuilding_paused = True
-
-    @staticmethod
-    def unpause_role_ancestor_rebuilding():
-        '''
-        Unpauses the role ancestor list updating. This will will rebuild all
-        roles that need updating since the last call to
-        pause_role_ancestor_rebuilding and bring everything back into sync.
-        '''
-        global role_rebuilding_paused
-        global roles_needing_rebuilding
-        role_rebuilding_paused = False
-        for role in Role.objects.filter(id__in=list(roles_needing_rebuilding)).all():
-            role.rebuild_role_ancestor_list()
-        roles_needing_rebuilding = set()
 
     def rebuild_role_ancestor_list(self):
         '''
@@ -100,9 +105,11 @@ class Role(CommonModelNameNotUnique):
 
         Note that this method relies on any parents' ancestor list being correct.
         '''
-        global role_rebuilding_paused, roles_needing_rebuilding
+        global tls
+        batch_role_rebuilding = getattr(tls, 'batch_role_rebuilding', False)
 
-        if role_rebuilding_paused:
+        if batch_role_rebuilding:
+            roles_needing_rebuilding = getattr(tls, 'roles_needing_rebuilding')
             roles_needing_rebuilding.add(self.id)
             return