Add memory-safe option

Author: AlanCoding

ansible_base/rbac/triggers.py

@@ -37,24 +37,36 @@ def _get_bulk_cache_state():
         _bulk_cache_state.active = False
         _bulk_cache_state.needs_team_update = False
         _bulk_cache_state.object_roles_to_update = set()
+        _bulk_cache_state.memory_safe = False
+        _bulk_cache_state.needs_object_role_update = False
     return _bulk_cache_state
 
 
 @contextmanager
-def bulk_rbac_caching():
+def bulk_rbac_caching(memory_safe=False):
     """
     Context manager that defers expensive RBAC cache updates during bulk operations.
 
     Instead of calling update_after_assignment for each permission change,
     this collects all the updates and performs them once when exiting the context.
 
+    Args:
+        memory_safe (bool): If True, doesn't store individual object roles to save memory.
+                           Instead performs a full re-computation of all object roles
+                           when exiting the context, similar to post_migrate behavior.
+
     Usage:
         with bulk_rbac_caching():
             # Multiple permission assignments/removals
             role_def.give_permission(user1, obj1)
             role_def.give_permission(user2, obj2)
             role_def.remove_permission(user3, obj3)
             # Cache updates happen here when exiting the context
+
+        with bulk_rbac_caching(memory_safe=True):
+            # For very large bulk operations where memory usage is a concern
+            # This will recompute all object roles instead of tracking specific ones
+            # ... many operations ...
     """
     state = _get_bulk_cache_state()
 
@@ -67,26 +79,39 @@ def bulk_rbac_caching():
     state.active = True
     state.needs_team_update = False
     state.object_roles_to_update = set()
+    state.memory_safe = memory_safe
+    state.needs_object_role_update = False
 
     try:
         yield
     finally:
         # Exit bulk mode and perform deferred updates
         needs_team_update = state.needs_team_update
         object_roles_to_update = state.object_roles_to_update.copy()
+        is_memory_safe = state.memory_safe
+        needs_object_role_update = state.needs_object_role_update
 
         # Reset state
         state.active = False
         state.needs_team_update = False
         state.object_roles_to_update = set()
+        state.memory_safe = False
+        state.needs_object_role_update = False
 
         # Perform the global update
-        if needs_team_update or object_roles_to_update:
-            logger.info(f'Performing bulk RBAC cache update: teams={needs_team_update}, object_roles={len(object_roles_to_update)}')
-            if needs_team_update:
-                compute_team_member_roles()
-            if object_roles_to_update:
-                compute_object_role_permissions(object_roles=object_roles_to_update)
+        if needs_team_update or object_roles_to_update or (is_memory_safe and needs_object_role_update):
+            if is_memory_safe and needs_object_role_update:
+                logger.info('Performing bulk RBAC cache update: memory_safe=True, recomputing all object roles')
+                if needs_team_update:
+                    compute_team_member_roles()
+                # In memory-safe mode, always recompute all object roles
+                compute_object_role_permissions()
+            elif not is_memory_safe:
+                logger.info(f'Performing bulk RBAC cache update: teams={needs_team_update}, object_roles={len(object_roles_to_update)}')
+                if needs_team_update:
+                    compute_team_member_roles()
+                if object_roles_to_update:
+                    compute_object_role_permissions(object_roles=object_roles_to_update)
 
 
 def team_ancestor_roles(team):
@@ -155,7 +180,12 @@ def update_after_assignment(update_teams, to_update):
         if update_teams:
             state.needs_team_update = True
         if to_update:
-            state.object_roles_to_update.update(to_update)
+            if state.memory_safe:
+                # In memory-safe mode, don't store individual object roles to save memory,
+                # but track that we need to do a full recomputation
+                state.needs_object_role_update = True
+            else:
+                state.object_roles_to_update.update(to_update)
         return
 
     # Normal mode - perform updates immediately

test_app/tests/rbac/test_triggers.py

@@ -346,3 +346,131 @@ def test_bulk_caching_empty_context(self):
             # Should not call expensive functions if no changes were made
             mock_team_update.assert_not_called()
             mock_obj_update.assert_not_called()
+
+    def test_bulk_caching_memory_safe_mode(self, rando, team, inv_rd, org_inv_rd, inventory, organization):
+        """Test that memory_safe=True mode doesn't store object roles but still recomputes"""
+        with patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update:
+
+            with bulk_rbac_caching(memory_safe=True):
+                # Multiple assignments that would normally collect object roles
+                inv_rd.give_permission(rando, inventory)
+                org_inv_rd.give_permission(team, organization)
+
+                # Should not be called during bulk mode
+                mock_obj_update.assert_not_called()
+
+            # Should be called once with no specific object roles (full recomputation)
+            mock_obj_update.assert_called_once_with()
+
+    def test_bulk_caching_memory_safe_vs_normal_mode(self, rando, inv_rd, inventory):
+        """Test that memory_safe mode calls recomputation while normal mode passes specific roles"""
+        # Test normal mode
+        with patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update:
+            with bulk_rbac_caching(memory_safe=False):
+                assignment = inv_rd.give_permission(rando, inventory)
+
+            # Should be called with specific object roles
+            mock_obj_update.assert_called_once()
+            call_args = mock_obj_update.call_args
+            object_roles = call_args.kwargs['object_roles']
+            assert assignment.object_role in object_roles
+
+        # Clean up the assignment
+        inv_rd.remove_permission(rando, inventory)
+
+        # Test memory-safe mode
+        with patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update:
+            with bulk_rbac_caching(memory_safe=True):
+                inv_rd.give_permission(rando, inventory)
+
+            # Should be called with no specific object roles (full recomputation)
+            mock_obj_update.assert_called_once_with()
+
+    def test_bulk_caching_memory_safe_with_team_updates(self, rando, team, member_rd):
+        """Test that memory_safe mode works correctly with team updates"""
+        with (
+            patch('ansible_base.rbac.triggers.compute_team_member_roles') as mock_team_update,
+            patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update,
+        ):
+
+            with bulk_rbac_caching(memory_safe=True):
+                # Assignment that affects team membership
+                member_rd.give_permission(rando, team)
+
+                # Should not be called during bulk mode
+                mock_team_update.assert_not_called()
+                mock_obj_update.assert_not_called()
+
+            # Both should be called, but object role update should be full recomputation
+            mock_team_update.assert_called_once()
+            mock_obj_update.assert_called_once_with()
+
+    def test_bulk_caching_memory_safe_empty_context(self):
+        """Test that empty memory_safe context doesn't call cache functions"""
+        with (
+            patch('ansible_base.rbac.triggers.compute_team_member_roles') as mock_team_update,
+            patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update,
+        ):
+
+            with bulk_rbac_caching(memory_safe=True):
+                # No operations performed
+                pass
+
+            # Should not call expensive functions if no changes were made, even in memory_safe mode
+            mock_team_update.assert_not_called()
+            mock_obj_update.assert_not_called()
+
+    def test_bulk_caching_memory_safe_nested_contexts(self, rando, inv_rd, inventory):
+        """Test that nested contexts with memory_safe work correctly"""
+        with patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update:
+
+            with bulk_rbac_caching(memory_safe=True):
+                inv_rd.give_permission(rando, inventory)
+
+                # Nested context (memory_safe is ignored in nested calls)
+                with bulk_rbac_caching(memory_safe=False):
+                    # Another assignment in nested context
+                    # Should still not trigger updates
+                    pass
+
+                # Still in outer context, should not be called yet
+                mock_obj_update.assert_not_called()
+
+            # Only called once when exiting outermost context, with full recomputation
+            mock_obj_update.assert_called_once_with()
+
+    def test_bulk_caching_memory_safe_with_removal(self, rando, inv_rd, inventory):
+        """Test memory_safe mode when object role gets deleted during removal"""
+        # First give permission normally
+        inv_rd.give_permission(rando, inventory)
+
+        with patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update:
+
+            with bulk_rbac_caching(memory_safe=True):
+                # Remove permission in bulk mode - this will delete the object role
+                inv_rd.remove_permission(rando, inventory)
+                mock_obj_update.assert_not_called()
+
+            # Should not call recomputation since object role was deleted (nothing to update)
+            # This matches the existing behavior in regular bulk mode
+            mock_obj_update.assert_not_called()
+
+    def test_bulk_caching_memory_safe_mixed_operations(self, rando, team, inv_rd, org_inv_rd, inventory, organization):
+        """Test memory_safe mode with mixed add/remove operations that result in net changes"""
+        with patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update:
+
+            with bulk_rbac_caching(memory_safe=True):
+                # Add permissions (creates object roles)
+                inv_rd.give_permission(rando, inventory)
+                org_inv_rd.give_permission(team, organization)
+
+                # Remove one permission (may or may not delete object role)
+                inv_rd.remove_permission(rando, inventory)
+
+                # Add it back
+                inv_rd.give_permission(rando, inventory)
+
+                mock_obj_update.assert_not_called()
+
+            # Should call full recomputation since we had net updates
+            mock_obj_update.assert_called_once_with()