diff --git a/ansible_base/rbac/__init__.py b/ansible_base/rbac/__init__.py index fd0666930..c7349ac3e 100644 --- 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 index ee90a1446..99c69ca41 100644 --- 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 index 12e806bc8..554a8297a 100644 --- 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,97 @@ 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() + _bulk_cache_state.memory_safe = False + _bulk_cache_state.needs_object_role_update = False + return _bulk_cache_state + + +@contextmanager +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() + + 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() + 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 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() + # When team memberships change, always recompute object role permissions + # to ensure the permission cache reflects new team relationships + compute_object_role_permissions() + elif object_roles_to_update: + # Only object roles changed, no team updates needed + 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 +177,22 @@ 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: + 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 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 index 6ec94246b..0d9f55036 100644 --- a/test_app/tests/rbac/test_triggers.py +++ b/test_app/tests/rbac/test_triggers.py @@ -1,4 +1,4 @@ -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest from django.apps import apps @@ -6,7 +6,7 @@ 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,343 @@ 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() + + 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() + + def test_bulk_caching_with_removal(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_team_only_updates_fix(self, rando, team, member_rd): + """Test that team-only updates properly call compute_object_role_permissions""" + 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(): + # This assignment affects team membership but doesn't create a new object role + # that would be added to object_roles_to_update + 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() + + # CRITICAL: Both should be called when team updates happen + # This ensures the permission cache reflects new team relationships + mock_team_update.assert_called_once() + # This is the bug fix - compute_object_role_permissions should be called + # even when only team updates occurred (no object_roles_to_update) + mock_obj_update.assert_called_once_with()