AAP-53278 General API optimizations for RBAC

ansible_base/activitystream/views.py

@@ -37,7 +37,7 @@ class EntryReadOnlyViewSet(ReadOnlyModelViewSet, AnsibleBaseDjangoAppApiView):
     API endpoint that allows for read-only access to activity stream entries.
     """
 
-    queryset = Entry.objects.all()
+    queryset = Entry.objects.prefetch_related('created_by', 'content_type').all()
     serializer_class = EntrySerializer
     filter_backends = calculate_filter_backends()
     ordering = ['-id']

ansible_base/rbac/api/views.py

@@ -63,11 +63,20 @@ class RoleMetadataView(AnsibleBaseDjangoAppApiView, GenericAPIView):
     permission_classes = try_add_oauth2_scope_permission([permissions.IsAuthenticated])
     serializer_class = RoleMetadataSerializer
 
+    def __init__(self, *args, **kwargs):
+        self.permission_cache = {}
+
     def dispatch(self, request, *args, **kwargs):
         # Warm cache to avoid hits to basically all types from serializer
         DABContentType.objects.get_for_models(*permission_registry.all_registered_models)
         return super().dispatch(request, *args, **kwargs)
 
+    def get_for_codename(self, codename):
+        if codename not in self.permission_cache:
+            for permission in permission_registry.permission_qs.all():
+                self.permission_cache[permission.codename] = permission
+        return self.permission_cache[codename]
+
     def get(self, request, format=None):
         data = OrderedDict()
         allowed_permissions = OrderedDict()
@@ -84,7 +93,7 @@ def get(self, request, format=None):
                 cls_repr = f"{get_resource_prefix(cls)}.{cls._meta.model_name}"
             allowed_permissions[cls_repr] = []
             for codename in list_combine_values(permissions_allowed_for_role(cls)):
-                perm = permission_registry.permission_qs.get(codename=codename)
+                perm = self.get_for_codename(codename)
                 ct = permission_registry.content_type_model.objects.get_for_id(perm.content_type_id)
                 perm_repr = f"{get_resource_prefix(ct.model_class())}.{codename}"
                 allowed_permissions[cls_repr].append(perm_repr)
@@ -107,7 +116,7 @@ class RoleDefinitionViewSet(AnsibleBaseDjangoAppApiView, ModelViewSet):
     but can be assigned to users.
     """
 
-    queryset = RoleDefinition.objects.prefetch_related('created_by', 'modified_by', 'content_type', 'permissions')
+    queryset = RoleDefinition.objects.prefetch_related('created_by', 'modified_by', 'content_type', 'permissions', 'resource')
     serializer_class = RoleDefinitionSerializer
     permission_classes = try_add_oauth2_scope_permission([RoleDefinitionPermissions])
 
@@ -216,7 +225,7 @@ class RoleTeamAssignmentViewSet(BaseAssignmentViewSet):
     """
 
     serializer_class = RoleTeamAssignmentSerializer
-    prefetch_related = ('team',)
+    prefetch_related = ('team__resource',)
     filter_backends = BaseAssignmentViewSet.filter_backends + [
         ansible_id_backend.TeamAnsibleIdAliasFilterBackend,
         ansible_id_backend.RoleAssignmentFilterBackend,
@@ -236,7 +245,7 @@ class RoleUserAssignmentViewSet(BaseAssignmentViewSet):
     """
 
     serializer_class = RoleUserAssignmentSerializer
-    prefetch_related = ('user',)
+    prefetch_related = ('user__resource',)
     filter_backends = BaseAssignmentViewSet.filter_backends + [
         ansible_id_backend.UserAnsibleIdAliasFilterBackend,
         ansible_id_backend.RoleAssignmentFilterBackend,

ansible_base/rbac/service_api/serializers.py

@@ -1,5 +1,4 @@
 from crum import impersonate
-from django.apps import apps
 from django.db import transaction
 from django.utils.translation import gettext_lazy as _
 from rest_framework import serializers
@@ -9,6 +8,49 @@
 from ..remote import RemoteObject
 
 
+class ObjectAnsibleIdField(serializers.Field):
+    """
+    Field for object_ansible_id that supports both annotation optimization and fallback.
+
+    For read operations: Uses annotation when available, falls back to manual lookup.
+    For write operations: Converts ansible_id to object_id for internal use.
+    """
+
+    def to_representation(self, obj):
+        """Get object_ansible_id, using annotation when available, falling back to manual lookup"""
+        # First try to use the annotation from the queryset (for optimized list operations)
+        if hasattr(obj, '_object_ansible_id_annotation') and obj._object_ansible_id_annotation:
+            return str(obj._object_ansible_id_annotation)
+
+        # Fallback for cases where annotation is not available (creation, etc.)
+        if not obj.content_type_id or not obj.object_id:
+            return None
+
+        content_object = obj.content_object
+        if isinstance(content_object, RemoteObject):
+            return None
+        if hasattr(content_object, 'resource'):
+            return str(content_object.resource.ansible_id)
+        return None
+
+    def get_attribute(self, instance):
+        """Override to return the full instance instead of a specific attribute"""
+        return instance
+
+    def to_internal_value(self, value):
+        """Convert object_ansible_id to object_id for internal use"""
+        if not value:
+            return None
+
+        from ansible_base.resource_registry.models import Resource
+
+        try:
+            resource = Resource.objects.get(ansible_id=value)
+            return resource.object_id
+        except Resource.DoesNotExist:
+            raise serializers.ValidationError("Resource with this ansible_id does not exist.")
+
+
 class DABContentTypeSerializer(serializers.ModelSerializer):
     parent_content_type = serializers.SlugRelatedField(read_only=True, slug_field='api_slug')
 
@@ -25,68 +67,36 @@ class Meta:
         fields = ['api_slug', 'codename', 'content_type', 'name']
 
 
-class ObjectIDAnsibleIDField(serializers.Field):
-    "This is an ansible_id field intended to be used with source pointing to object_id, so, does conversion"
-
-    def to_representation(self, value):
-        "The source for this field is object_id, which is ignored, use content_object instead"
-        assignment = getattr(self, "_this_assignment", None)
-        if not assignment:
-            return None
-        content_object = assignment.content_object
-        if isinstance(content_object, RemoteObject):
-            return None
-        if hasattr(content_object, 'resource'):
-            return str(content_object.resource.ansible_id)
-        return None
-
-    def to_internal_value(self, value):
-        "Targeting object_id, this converts ansible_id into object_id"
-        resource_cls = apps.get_model('dab_resource_registry', 'Resource')
-        resource = resource_cls.objects.get(ansible_id=value)
-        return resource.object_id
-
-
 assignment_common_fields = ('created', 'created_by_ansible_id', 'object_id', 'object_ansible_id', 'content_type', 'role_definition')
 
 
 class BaseAssignmentSerializer(serializers.ModelSerializer):
     content_type = serializers.SlugRelatedField(read_only=True, slug_field='api_slug')
     role_definition = serializers.SlugRelatedField(slug_field='name', queryset=RoleDefinition.objects.all())
     created_by_ansible_id = ActorAnsibleIdField(source='created_by', required=False, allow_null=True)
-    object_ansible_id = ObjectIDAnsibleIDField(source='object_id', required=False, allow_null=True)
+    object_ansible_id = ObjectAnsibleIdField(required=False, allow_null=True)
     object_id = serializers.CharField(allow_blank=True, required=False, allow_null=True)
     from_service = serializers.CharField(write_only=True)
 
-    def to_representation(self, instance):
-        # hack to surface content_object for ObjectIDAnsibleIDField
-        self.fields["object_ansible_id"]._this_assignment = instance
-        return super().to_representation(instance)
-
-    def get_created_by_ansible_id(self, obj):
-        return str(obj.created_by.resource.ansible_id)
-
     def validate(self, attrs):
         """The object_id vs ansible_id is the only dual-write case, where we have to accept either
 
         So this does the mutual validation to assure we have sufficient data.
         """
-        has_oid = bool(self.initial_data.get('object_id'))
-        has_oaid = bool(self.initial_data.get('object_ansible_id'))
-
         rd = attrs['role_definition']
+        has_object_id = 'object_id' in attrs and attrs['object_id']
+        has_object_ansible_id = 'object_ansible_id' in attrs and attrs['object_ansible_id']
+
         if rd.content_type_id:
-            if not self.partial and not has_oid and not has_oaid:
+            if not self.partial and not has_object_id and not has_object_ansible_id:
                 raise serializers.ValidationError("You must provide either 'object_id' or 'object_ansible_id'.")
-            elif not has_oaid:
-                # need to remove blank and null fields or else it can overwrite the non-null non-blank field
-                attrs['object_id'] = self.initial_data['object_id']
+            # If object_ansible_id was provided and converted, use that for object_id
+            if has_object_ansible_id and not has_object_id:
+                attrs['object_id'] = attrs['object_ansible_id']
         else:
-            if has_oaid or has_oid:
+            if has_object_id or has_object_ansible_id:
                 raise serializers.ValidationError("Can not provide either 'object_id' or 'object_ansible_id' for system role")
 
-        # NOTE: right now not enforcing the case you provide both, could check for consistency later
-
         return super().validate(attrs)
 
     def find_existing_assignment(self, queryset):

ansible_base/rbac/service_api/views.py

@@ -1,11 +1,13 @@
 from django.db import transaction
+from django.db.models import OuterRef, Subquery
 from rest_framework import permissions, status, viewsets
 from rest_framework.decorators import action
 from rest_framework.response import Response
 from rest_framework.viewsets import GenericViewSet, mixins
 
 from ansible_base.lib.utils.views.django_app_api import AnsibleBaseDjangoAppApiView
 from ansible_base.lib.utils.views.permissions import try_add_oauth2_scope_permission
+from ansible_base.resource_registry.models import Resource
 from ansible_base.resource_registry.views import HasResourceRegistryPermissions
 from ansible_base.rest_filters.rest_framework import ansible_id_backend
 
@@ -109,10 +111,23 @@ def perform_destroy(self, instance):
                 instance.role_definition.remove_global_permission(instance.actor)
 
 
+def resource_ansible_id_expr(ct_field='content_type_id', oid_field='object_id'):
+    return Subquery(
+        Resource.objects.filter(
+            content_type_id=OuterRef(ct_field),
+            object_id=OuterRef(oid_field),
+        ).values(
+            'ansible_id'
+        )[:1]
+    )
+
+
 class ServiceRoleUserAssignmentViewSet(BaseSerivceRoleAssignmentViewSet):
     """List of user assignments for cross-service communication"""
 
-    queryset = RoleUserAssignment.objects.prefetch_related('user__resource', *prefetch_related)
+    queryset = RoleUserAssignment.objects.prefetch_related('user__resource__content_type', *prefetch_related).annotate(
+        _object_ansible_id_annotation=resource_ansible_id_expr()
+    )
     serializer_class = service_serializers.ServiceRoleUserAssignmentSerializer
     filter_backends = AnsibleBaseDjangoAppApiView.filter_backends + [
         ansible_id_backend.UserAnsibleIdAliasFilterBackend,
@@ -131,7 +146,9 @@ def unassign(self, request):
 class ServiceRoleTeamAssignmentViewSet(BaseSerivceRoleAssignmentViewSet):
     """List of team role assignments for cross-service communication"""
 
-    queryset = RoleTeamAssignment.objects.prefetch_related('team__resource', *prefetch_related)
+    queryset = RoleTeamAssignment.objects.prefetch_related('team__resource__content_type', *prefetch_related).annotate(
+        _object_ansible_id_annotation=resource_ansible_id_expr()
+    )
     serializer_class = service_serializers.ServiceRoleTeamAssignmentSerializer
     filter_backends = AnsibleBaseDjangoAppApiView.filter_backends + [
         ansible_id_backend.TeamAnsibleIdAliasFilterBackend,

ansible_base/resource_registry/models/resource.py

@@ -95,11 +95,11 @@ def summary_fields(self):
 
     @property
     def resource_type(self):
-        return resource_type_cache(self.content_type.pk).name
+        return resource_type_cache(self.content_type_id).name
 
     @property
     def resource_type_obj(self):
-        return resource_type_cache(self.content_type.pk)
+        return resource_type_cache(self.content_type_id)
 
     class Meta:
         unique_together = ('content_type', 'object_id')

test_app/defaults.py

@@ -143,7 +143,7 @@
 
 USE_TZ = True
 
-DEMO_DATA_COUNTS = {'organization': 150, 'user': 379, 'team': 43}
+DEMO_DATA_COUNTS = {'organization': 150, 'user': 379, 'team': 43, 'roledefinition': 100}
 
 ANSIBLE_BASE_TEAM_MODEL = 'test_app.Team'
 ANSIBLE_BASE_ORGANIZATION_MODEL = 'test_app.Organization'

test_app/management/commands/create_demo_data.py

@@ -8,7 +8,7 @@
 from ansible_base.authentication.models import Authenticator, AuthenticatorUser
 from ansible_base.oauth2_provider.models import OAuth2Application
 from ansible_base.rbac import permission_registry
-from ansible_base.rbac.models import DABContentType, RoleDefinition
+from ansible_base.rbac.models import DABContentType, DABPermission, RoleDefinition
 from test_app.models import EncryptionModel, InstanceGroup, Inventory, Organization, Team, User
 
 
@@ -20,6 +20,9 @@ def create_large(self, data_counts):
         start = time.time()
         self.stdout.write('')
         self.stdout.write('About to create large demo data set. This will take a while.')
+
+        # Create standard models first
+        created_org_ids = []
         for cls in (Organization, Team, User):
             count = data_counts[cls._meta.model_name]
             for i in range(count):
@@ -28,9 +31,74 @@ def create_large(self, data_counts):
                 if cls is User:
                     data = {'username': name}
                 elif cls is Team:
-                    data['organization_id'] = i + 1  # fudged, teams fewer than orgs
-                cls.objects.create(**data)
+                    # Use actual created organization IDs, cycling through them
+                    if created_org_ids:
+                        data['organization_id'] = created_org_ids[i % len(created_org_ids)]
+                    else:
+                        raise ValueError("Teams cannot be created before organizations")
+                obj = cls.objects.create(**data)
+                # Collect organization IDs for team creation
+                if cls is Organization:
+                    created_org_ids.append(obj.id)
             self.stdout.write(f'Created {count} {cls._meta.model_name}')
+
+        # Create RoleDefinitions with permissions
+        if 'roledefinition' in data_counts:
+            rd_count = data_counts['roledefinition']
+            org_ct = DABContentType.objects.get_for_model(Organization)
+
+            for i in range(rd_count):
+                # Create some sample permissions for each role definition
+                perm1 = DABPermission.objects.create(name=f'Can view large role {i}', codename=f'view_large_role_{i}', content_type=org_ct)
+                perm2 = DABPermission.objects.create(name=f'Can edit large role {i}', codename=f'edit_large_role_{i}', content_type=org_ct)
+
+                # Create role definition with Organization content type
+                rd = RoleDefinition.objects.create(name=f'Large Role Definition {i}', description=f'Large demo role definition {i}', content_type=org_ct)
+
+                # Add permissions to the role definition
+                rd.permissions.add(perm1, perm2)
+
+            self.stdout.write(f'Created {rd_count} role definitions with permissions')
+
+        # Create permission assignments for users and teams
+        if created_org_ids and 'user' in data_counts and 'team' in data_counts:
+            # Get created users and teams
+            large_users = list(User.objects.filter(username__startswith='large_user_'))
+            large_teams = list(Team.objects.filter(name__startswith='large_team_'))
+            large_orgs = list(Organization.objects.filter(name__startswith='large_organization_'))
+            large_rds = list(RoleDefinition.objects.filter(name__startswith='Large Role Definition'))
+
+            # Give over 25 permissions to users
+            user_permissions_given = 0
+            for user in large_users:
+                for rd in large_rds:
+                    for org in large_orgs:
+                        rd.give_permission(user, org)
+                        user_permissions_given += 1
+                        if user_permissions_given >= 25:
+                            break
+                    if user_permissions_given >= 25:
+                        break
+                if user_permissions_given >= 25:
+                    break
+
+            # Give over 25 permissions to teams
+            team_permissions_given = 0
+            for team in large_teams:
+                for rd in large_rds:
+                    for org in large_orgs:
+                        rd.give_permission(team, org)
+                        team_permissions_given += 1
+                        if team_permissions_given >= 25:
+                            break
+                    if team_permissions_given >= 25:
+                        break
+                if team_permissions_given >= 25:
+                    break
+
+            self.stdout.write(f'Assigned {user_permissions_given} permissions to users')
+            self.stdout.write(f'Assigned {team_permissions_given} permissions to teams')
+
         self.stdout.write(f'Finished creating large demo data in {time.time() - start:.2f} seconds')
 
     def handle(self, *args, **kwargs):