⚡️(backend) reduce queries made in item access list viewset

The ItemAvvessViewset list endpoint is really complicated and half of
the algo is made to cache the role the current user have on the tree.
This code is now removed and replaced by a queryset annotation allowing
to fetch the user_roles and then determine in the model the max role the
user have on the tree. It reduces the number of sql query by 1 and ease
the code comprehension

Author: lunika

src/backend/core/api/viewsets.py

@@ -1186,21 +1186,21 @@ def list(self, request, *args, **kwargs):
         ancestors_qs = models.Item.objects.filter(
             path__ancestors=self.item.path, ancestors_deleted_at__isnull=True
         )
-        ancestor_items = list(ancestors_qs.order_by("path"))
-
-        accesses_queryset = self.get_queryset().filter(item__in=ancestors_qs)
+        accesses_qs = self.get_queryset().filter(item__in=ancestors_qs)
         if role not in PRIVILEGED_ROLES:
-            accesses_queryset = accesses_queryset.filter(role__in=PRIVILEGED_ROLES)
+            accesses_qs = accesses_qs.filter(role__in=PRIVILEGED_ROLES)
 
-        accesses = list(accesses_queryset.order_by("item__path", "created_at"))
+        accesses_qs = accesses_qs.annotate_user_roles(user)
+        accesses = list(accesses_qs.order_by("item__path", "created_at").iterator())
 
         request_target_keys = {f"user:{user.id}"}
         request_target_keys.update(
             f"team:{team}" for team in getattr(user, "teams", [])
         )
 
         max_role_by_target = {}
-        user_roles_by_path = {}
+        # Keep in this dict the deepest access for each target
+        # This is the access that will be returned to the user
         accesses_by_target = defaultdict(list)
         for access in accesses:
             previous = max_role_by_target.get(access.target_key)
@@ -1215,49 +1215,19 @@ def list(self, request, *args, **kwargs):
                 "item_id": access.item_id,
             }
 
-            accesses_by_target[access.target_key].append(access)
-
-            if access.target_key in request_target_keys:
-                path = str(access.item.path)
-                user_roles_by_path[path] = models.RoleChoices.max(
-                    user_roles_by_path.get(path), access.role
-                )
-
-        user_ancestor_role_by_path = {}
-        user_max_role_by_path = {}
-        for ancestor in ancestor_items:
-            path = str(ancestor.path)
-            parent_path = str(ancestor.path[:-1]) if ancestor.depth > 1 else ""
-            ancestor_role = user_max_role_by_path.get(parent_path)
-            user_ancestor_role_by_path[path] = ancestor_role
-            user_max_role_by_path[path] = models.RoleChoices.max(
-                ancestor_role, user_roles_by_path.get(path)
-            )
+            accesses_by_target[access.target_key] = access
 
-        selected_access_ids = {
-            max(
-                target_accesses,
-                key=lambda item_access: (
-                    item_access.item.depth if item_access.item_id else 0,
-                    item_access.created_at,
-                ),
-            ).pk
-            for target_accesses in accesses_by_target.values()
-            if target_accesses
-        }
+        # Reorder accesses by their corresponding item depth and creation date
+        selected_accesses = sorted(
+            accesses_by_target.values(),
+            key=lambda access: (access.item.depth, access.created_at),
+        )
 
         # serialize and return the response
         context = self.get_serializer_context()
         serializer_class = self.get_serializer_class()
         serialized_data = []
-        for access in accesses:
-            if access.pk not in selected_access_ids:
-                continue
-            path = str(access.item.path)
-            access.set_user_roles_tuple(
-                user_ancestor_role_by_path.get(path),
-                user_roles_by_path.get(path),
-            )
+        for access in selected_accesses:
             serializer = serializer_class(access, context=context)
             serialized_data.append(serializer.data)
 

src/backend/core/models.py

@@ -1103,6 +1103,37 @@ def __str__(self):
         return f"{self.user!s} favorite on item {self.item!s}"
 
 
+class ItemAccessQuerySet(models.QuerySet):
+    """Custom queryset for ItemAccess model with additional methods."""
+
+    def annotate_user_roles(self, user):
+        """
+        Annotate ItemAccess queryset with the roles of the current user
+        on the item or its ancestors.
+        """
+        output_field = ArrayField(base_field=models.CharField())
+
+        if user.is_authenticated:
+            user_roles_subquery = ItemAccess.objects.filter(
+                models.Q(user=user) | models.Q(team__in=user.teams),
+                item__path__ancestors=models.OuterRef("item__path"),
+            ).values_list("role", flat=True)
+
+            return self.annotate(
+                user_roles=models.Func(
+                    user_roles_subquery, function="ARRAY", output_field=output_field
+                )
+            )
+
+        return self.annotate(
+            user_roles=models.Value([], output_field=output_field),
+        )
+
+
+class ItemAccessManager(models.Manager.from_queryset(ItemAccessQuerySet)):
+    """Manager for ItemAccess model."""
+
+
 class ItemAccess(BaseModel):
     """Relation model to give access to an item for a user or a team with a role."""
 
@@ -1122,6 +1153,8 @@ class ItemAccess(BaseModel):
         max_length=20, choices=RoleChoices.choices, default=RoleChoices.READER
     )
 
+    objects = ItemAccessManager()
+
     class Meta:
         db_table = "drive_item_access"
         ordering = ("-created_at",)
@@ -1166,65 +1199,6 @@ def target_key(self):
         """Get a unique key for the actor targeted by the access, without possible conflict."""
         return f"user:{self.user_id!s}" if self.user_id else f"team:{self.team:s}"
 
-    def set_user_roles_tuple(self, ancestors_role, current_role):
-        """
-        Set a precomputed (ancestor_role, current_role) tuple for this instance.
-
-        This avoids querying the database in `get_roles_tuple()` and is useful
-        when roles are already known, such as in bulk serialization.
-
-        Args:
-            ancestor_role (str | None): Highest role on any ancestor document.
-            current_role (str | None): Role on the current document.
-        """
-        # pylint: disable=attribute-defined-outside-init
-        self._prefetched_user_roles_tuple = (ancestors_role, current_role)
-
-    def get_user_roles_tuple(self, user):
-        """
-        Return a tuple of:
-        - the highest role the user has on any ancestor of the item
-        - the role the user has on the current item
-
-        If roles have been explicitly set using `set_user_roles_tuple()`,
-        those will be returned instead of querying the database.
-
-        This allows viewsets or serializers to precompute roles for performance
-        when handling multiple items at once.
-
-        Args:
-            user (User): The user whose roles are being evaluated.
-
-        Returns:
-            tuple[str | None, str | None]: (max_ancestor_role, current_document_role)
-        """
-        if not user.is_authenticated:
-            return None, None
-
-        try:
-            return self._prefetched_user_roles_tuple
-        except AttributeError:
-            pass
-
-        ancestors = (
-            self.item.ancestors() | Item.objects.filter(pk=self.item_id)
-        ).filter(ancestors_deleted_at__isnull=True)
-
-        access_tuples = ItemAccess.objects.filter(
-            models.Q(user=user) | models.Q(team__in=user.teams),
-            item__in=ancestors,
-        ).values_list("item_id", "role")
-
-        ancestors_roles = []
-        current_role = None
-        for item_id, role in access_tuples:
-            if item_id == self.item_id:
-                current_role = role
-            else:
-                ancestors_roles.append(role)
-
-        return RoleChoices.max(*ancestors_roles), current_role
-
     def _compute_max_ancestors_role(self):
         """
         Compute the max ancestors role for this instance.
@@ -1281,16 +1255,30 @@ def max_ancestors_role_item_id(self, max_ancestors_role_item_id):
         """Cache the max_ancestors_role_item_id."""
         self._max_ancestors_role_item_id = max_ancestors_role_item_id
 
+    def get_role(self, user):
+        """Return the role a user has on an item related to this access.."""
+        if not user.is_authenticated:
+            return None
+
+        try:
+            roles = self.user_roles or []
+        except AttributeError:
+            roles = ItemAccess.objects.filter(
+                models.Q(user=user) | models.Q(team__in=user.teams),
+                item__path__ancestors=self.item.path,
+            ).values_list("role", flat=True)
+
+        return RoleChoices.max(*roles)
+
     def get_abilities(self, user, is_explicit=True):
         """
         Compute and return abilities for a given user on the item access.
         """
-        ancestors_role, current_role = self.get_user_roles_tuple(user)
-        role = RoleChoices.max(ancestors_role, current_role)
-        is_owner_or_admin = role in PRIVILEGED_ROLES
+        user_role = self.get_role(user)
+        is_owner_or_admin = user_role in PRIVILEGED_ROLES
 
         if self.role == RoleChoices.OWNER:
-            can_delete = role == RoleChoices.OWNER and (
+            can_delete = user_role == RoleChoices.OWNER and (
                 # check if item is not root trying to avoid an extra query
                 self.item.depth > 1
                 or ItemAccess.objects.filter(
@@ -1306,7 +1294,7 @@ def get_abilities(self, user, is_explicit=True):
                 set_role_to.extend(
                     [RoleChoices.READER, RoleChoices.EDITOR, RoleChoices.ADMIN]
                 )
-            if role == RoleChoices.OWNER:
+            if user_role == RoleChoices.OWNER:
                 set_role_to.append(RoleChoices.OWNER)
 
         # Filter out roles that would be lower than the one the user already has

src/backend/core/tests/items/test_api_item_accesses.py

@@ -145,7 +145,7 @@ def test_api_item_accesses_list_authenticated_related_non_privileged(
     other_access = factories.UserItemAccessFactory(user=user)
     factories.UserItemAccessFactory(item=other_access.item)
 
-    with django_assert_num_queries(4):
+    with django_assert_num_queries(3):
         response = client.get(f"/api/v1.0/items/{item.id!s}/accesses/")
 
     assert response.status_code == 200
@@ -253,7 +253,7 @@ def test_api_item_accesses_list_authenticated_related_privileged(
     other_access = factories.UserItemAccessFactory(user=user)
     factories.UserItemAccessFactory(item=other_access.item)
 
-    with django_assert_num_queries(4):
+    with django_assert_num_queries(3):
         response = client.get(f"/api/v1.0/items/{item.id!s}/accesses/")
 
     assert response.status_code == 200
@@ -1492,7 +1492,7 @@ def test_api_item_accesses_delete_owners_last_owner_child_team(
 ## Realistic case.
 
 
-def test_api_item_accesses_explicit():
+def test_api_item_accesses_explicit(django_assert_num_queries):
     """
     test case with a combination of explicit accesses and inherited accesses.
     An explicit access id added on the root item with a "weak" role (editor)
@@ -1517,7 +1517,8 @@ def test_api_item_accesses_explicit():
 
     assert item.get_role(user) == "owner"
 
-    response = client.get(f"/api/v1.0/items/{item.id!s}/accesses/")
+    with django_assert_num_queries(3):
+        response = client.get(f"/api/v1.0/items/{item.id!s}/accesses/")
     assert response.status_code == 200
     content = response.json()
     assert content == [
@@ -1616,7 +1617,9 @@ def test_api_item_accesses_explicit():
     factories.UserItemAccessFactory(item=parent, user=user, role="administrator")
 
     response = client.get(f"/api/v1.0/items/{item.id!s}/accesses/")
-    assert response.status_code == 200
+    with django_assert_num_queries(3):
+        response = client.get(f"/api/v1.0/items/{item.id!s}/accesses/")
+
     content = response.json()
     assert content == [
         {