Introduce a bulk re-computation context manager for RBAC

AlanCoding · AlanCoding · commit 641b303995d2 · 2025-09-16T16:49:04.000-04:00
diff --git a/ansible_base/rbac/__init__.py b/ansible_base/rbac/__init__.py
@@ -2,4 +2,13 @@
 
 __all__ = [
     'permission_registry',
+    'bulk_rbac_caching',
 ]
+
+
+def __getattr__(name):
+    if name == 'bulk_rbac_caching':
+        from ansible_base.rbac.triggers import bulk_rbac_caching
+
+        return bulk_rbac_caching
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
diff --git a/ansible_base/rbac/models/role.py b/ansible_base/rbac/models/role.py
@@ -281,7 +281,7 @@ def give_or_remove_permission(self, actor, content_object, giving=True, sync_act
 
         from ansible_base.rbac.triggers import needed_updates_on_assignment, update_after_assignment
 
-        update_teams, to_update = needed_updates_on_assignment(self, actor, object_role, created=created, giving=True)
+        update_teams, to_update = needed_updates_on_assignment(self, actor, object_role, created=created, giving=giving)
 
         assignment = None
         if actor._meta.model_name == 'user':
diff --git a/ansible_base/rbac/triggers.py b/ansible_base/rbac/triggers.py
@@ -1,5 +1,6 @@
 import logging
 from contextlib import contextmanager
+from threading import local
 from typing import Optional, Union
 from uuid import UUID
 
@@ -26,6 +27,68 @@
 dab_post_migrate = Signal()
 
 
+# Thread-local storage for bulk caching state
+_bulk_cache_state = local()
+
+
+def _get_bulk_cache_state():
+    """Get or initialize the bulk cache state for the current thread"""
+    if not hasattr(_bulk_cache_state, 'active'):
+        _bulk_cache_state.active = False
+        _bulk_cache_state.needs_team_update = False
+        _bulk_cache_state.object_roles_to_update = set()
+    return _bulk_cache_state
+
+
+@contextmanager
+def bulk_rbac_caching():
+    """
+    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.
+
+    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
+    """
+    state = _get_bulk_cache_state()
+
+    if state.active:
+        # Already in a bulk context, just yield (nested calls)
+        yield
+        return
+
+    # Enter bulk mode
+    state.active = True
+    state.needs_team_update = False
+    state.object_roles_to_update = set()
+
+    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()
+
+        # Reset state
+        state.active = False
+        state.needs_team_update = False
+        state.object_roles_to_update = set()
+
+        # 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)
+
+
 def team_ancestor_roles(team):
     """
     Return a queryset of all roles that directly or indirectly grant any form of permission to a team.
@@ -85,6 +148,17 @@ def needed_updates_on_assignment(role_definition, actor, object_role, created=Fa
 
 def update_after_assignment(update_teams, to_update):
     "Call this with the output of needed_updates_on_assignment"
+    state = _get_bulk_cache_state()
+
+    # If we're in bulk mode, defer the updates
+    if state.active:
+        if update_teams:
+            state.needs_team_update = True
+        if to_update:
+            state.object_roles_to_update.update(to_update)
+        return
+
+    # Normal mode - perform updates immediately
     if update_teams:
         compute_team_member_roles()
 
diff --git a/test_app/tests/rbac/test_triggers.py b/test_app/tests/rbac/test_triggers.py
@@ -1,12 +1,12 @@
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 from django.apps import apps
 from django.test.utils import override_settings
 
 from ansible_base.rbac.models import ObjectRole, RoleEvaluation, RoleTeamAssignment, RoleUserAssignment
 from ansible_base.rbac.permission_registry import permission_registry
-from ansible_base.rbac.triggers import dab_post_migrate, post_migration_rbac_setup
+from ansible_base.rbac.triggers import bulk_rbac_caching, dab_post_migrate, post_migration_rbac_setup
 from test_app.models import Inventory, Organization
 
 
@@ -172,3 +172,177 @@ def test_delete_signals_team_organization(organization, inventory, team, org_inv
         assert not RoleEvaluation.objects.filter(**org_gfk).exists()
 
     assert not RoleEvaluation.objects.filter(**inv_gfk).exists()
+
+
+@pytest.mark.django_db
+class TestBulkRBACCaching:
+    """Tests for the bulk_rbac_caching context manager"""
+
+    def test_bulk_caching_defers_updates(self, rando, inv_rd, inventory):
+        """Test that updates are deferred during bulk operations"""
+        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():
+                # Multiple permission assignments
+                inv_rd.give_permission(rando, inventory)
+                # During bulk mode, the expensive cache functions should not be called
+                mock_team_update.assert_not_called()
+                mock_obj_update.assert_not_called()
+
+            # After exiting context, they should be called once
+            mock_obj_update.assert_called_once()
+
+    def test_bulk_caching_collects_object_roles(self, rando, team, inv_rd, org_inv_rd, inventory, organization):
+        """Test that object roles are properly collected and updated in bulk"""
+        with patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update:
+
+            with bulk_rbac_caching():
+                # Multiple assignments that affect different object roles
+                assignment1 = inv_rd.give_permission(rando, inventory)
+                assignment2 = 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 collected object roles
+            mock_obj_update.assert_called_once()
+            call_args = mock_obj_update.call_args
+            object_roles = call_args.kwargs['object_roles']
+
+            # Should contain both object roles
+            assert assignment1.object_role in object_roles
+            assert assignment2.object_role in object_roles
+
+    def test_bulk_caching_handles_team_updates(self, rando, team, member_rd):
+        """Test that team updates are properly deferred and executed"""
+        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():
+                # 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 after exiting context
+            mock_team_update.assert_called_once()
+            mock_obj_update.assert_called_once()
+
+    def test_bulk_caching_nested_contexts(self, rando, inv_rd, inventory):
+        """Test that nested bulk contexts work correctly"""
+        with patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update:
+
+            with bulk_rbac_caching():
+                inv_rd.give_permission(rando, inventory)
+
+                # Nested context
+                with bulk_rbac_caching():
+                    # 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
+            mock_obj_update.assert_called_once()
+
+    def test_bulk_caching_with_multiple_assignments(self, rando, team, inv_rd, org_inv_rd, inventory, organization):
+        """Test bulk caching works with multiple assignments that require updates"""
+        with patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update:
+
+            with bulk_rbac_caching():
+                # Multiple operations that create new object roles
+                inv_rd.give_permission(rando, inventory)
+                org_inv_rd.give_permission(team, organization)
+                mock_obj_update.assert_not_called()
+
+            # Should be called once after all operations
+            mock_obj_update.assert_called_once()
+
+    def test_bulk_caching_with_object_role_deletion(self, rando, inv_rd, inventory):
+        """Test bulk caching 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():
+                # 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 be called since object role was deleted (nothing to update)
+            mock_obj_update.assert_not_called()
+
+    def test_bulk_caching_performance_benefit(self, rando, team, inv_rd, inventory):
+        """Test that bulk operations actually reduce cache update calls"""
+        with patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update:
+
+            # Without bulk caching - each assignment triggers an update
+            inv_rd.give_permission(rando, inventory)
+            inv_rd.remove_permission(rando, inventory)
+
+            # Should have been called multiple times
+            assert mock_obj_update.call_count >= 2
+
+            mock_obj_update.reset_mock()
+
+            # With bulk caching - should only be called once
+            with bulk_rbac_caching():
+                inv_rd.give_permission(rando, inventory)
+                inv_rd.give_permission(team, inventory)
+                inv_rd.remove_permission(rando, inventory)
+
+            # Should only be called once
+            mock_obj_update.assert_called_once()
+
+    def test_bulk_caching_exception_handling(self, rando, inv_rd, inventory):
+        """Test that bulk caching state is reset even if exception occurs"""
+        with patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update:
+
+            try:
+                with bulk_rbac_caching():
+                    inv_rd.give_permission(rando, inventory)
+                    raise ValueError("Test exception")
+            except ValueError:
+                pass
+
+            # Should still be called even though exception occurred
+            mock_obj_update.assert_called_once()
+
+            mock_obj_update.reset_mock()
+
+            # State should be reset, next operation should work normally
+            inv_rd.give_permission(rando, inventory)
+            mock_obj_update.assert_called()
+
+    def test_no_bulk_caching_without_context(self, rando, inv_rd, inventory):
+        """Test that normal operations still work when not in bulk context"""
+        with patch('ansible_base.rbac.triggers.compute_object_role_permissions') as mock_obj_update:
+
+            # Normal assignment should trigger immediate update
+            inv_rd.give_permission(rando, inventory)
+            mock_obj_update.assert_called()
+
+    def test_bulk_caching_empty_context(self):
+        """Test that empty bulk 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():
+                # No operations performed
+                pass
+
+            # Should not call expensive functions if no changes were made
+            mock_team_update.assert_not_called()
+            mock_obj_update.assert_not_called()
diff --git a/test_app/views.py b/test_app/views.py
@@ -178,6 +178,18 @@ def api_root(request, format=None):
     list_endpoints['role-metadata'] = get_fully_qualified_url('role-metadata')
     list_endpoints['timeout-view'] = get_fully_qualified_url('test-timeout-view')
 
+    from uuid import uuid4
+
+    from ansible_base.rbac.models import DABContentType, RoleDefinition
+    from test_app.models import Organization, User
+
+    DABContentType.objects.warm_cache()
+
+    rd = RoleDefinition.objects.get(name='Organization Admin')
+    user = User.objects.create(username=str(uuid4()))
+    org, _ = Organization.objects.get_or_create(name='Default')
+    rd.give_permission(user, org)
+
     return Response(list_endpoints)
 
 
