♻️(backend) factorize document query set annotation

The methods to annotate a document queryset were factorized on the
viewset but the correct place is the custom queryset itself now that
we have one.

Author: sampaccoud

src/backend/core/api/viewsets.py

@@ -400,44 +400,6 @@ class DocumentViewSet(
     trashbin_serializer_class = serializers.ListDocumentSerializer
     tree_serializer_class = serializers.ListDocumentSerializer
 
-    def annotate_is_favorite(self, queryset):
-        """
-        Annotate document queryset with the favorite status for the current user.
-        """
-        user = self.request.user
-
-        if user.is_authenticated:
-            favorite_exists_subquery = models.DocumentFavorite.objects.filter(
-                document_id=db.OuterRef("pk"), user=user
-            )
-            return queryset.annotate(is_favorite=db.Exists(favorite_exists_subquery))
-
-        return queryset.annotate(is_favorite=db.Value(False))
-
-    def annotate_user_roles(self, queryset):
-        """
-        Annotate document queryset with the roles of the current user
-        on the document or its ancestors.
-        """
-        user = self.request.user
-        output_field = ArrayField(base_field=db.CharField())
-
-        if user.is_authenticated:
-            user_roles_subquery = models.DocumentAccess.objects.filter(
-                db.Q(user=user) | db.Q(team__in=user.teams),
-                document__path=Left(db.OuterRef("path"), Length("document__path")),
-            ).values_list("role", flat=True)
-
-            return queryset.annotate(
-                user_roles=db.Func(
-                    user_roles_subquery, function="ARRAY", output_field=output_field
-                )
-            )
-
-        return queryset.annotate(
-            user_roles=db.Value([], output_field=output_field),
-        )
-
     def get_queryset(self):
         """Get queryset performing all annotation and filtering on the document tree structure."""
         user = self.request.user
@@ -473,8 +435,9 @@ def get_queryset(self):
     def filter_queryset(self, queryset):
         """Override to apply annotations to generic views."""
         queryset = super().filter_queryset(queryset)
-        queryset = self.annotate_is_favorite(queryset)
-        queryset = self.annotate_user_roles(queryset)
+        user = self.request.user
+        queryset = queryset.annotate_is_favorite(user)
+        queryset = queryset.annotate_user_roles(user)
         return queryset
 
     def get_response_for_queryset(self, queryset):
@@ -498,9 +461,10 @@ def list(self, request, *args, **kwargs):
         Additional annotations (e.g., `is_highest_ancestor_for_user`, favorite status) are
         applied before ordering and returning the response.
         """
-        queryset = (
-            self.get_queryset()
-        )  # Not calling filter_queryset. We do our own cooking.
+        user = self.request.user
+
+        # Not calling filter_queryset. We do our own cooking.
+        queryset = self.get_queryset()
 
         filterset = ListDocumentFilter(
             self.request.GET, queryset=queryset, request=self.request
@@ -513,7 +477,7 @@ def list(self, request, *args, **kwargs):
         for field in ["is_creator_me", "title"]:
             queryset = filterset.filters[field].filter(queryset, filter_data[field])
 
-        queryset = self.annotate_user_roles(queryset)
+        queryset = queryset.annotate_user_roles(user)
 
         # Among the results, we may have documents that are ancestors/descendants
         # of each other. In this case we want to keep only the highest ancestors.
@@ -530,7 +494,7 @@ def list(self, request, *args, **kwargs):
         )
 
         # Annotate favorite status and filter if applicable as late as possible
-        queryset = self.annotate_is_favorite(queryset)
+        queryset = queryset.annotate_is_favorite(user)
         queryset = filterset.filters["is_favorite"].filter(
             queryset, filter_data["is_favorite"]
         )
@@ -613,7 +577,7 @@ def trashbin(self, request, *args, **kwargs):
             deleted_at__isnull=False,
             deleted_at__gte=models.get_trashbin_cutoff(),
         )
-        queryset = self.annotate_user_roles(queryset)
+        queryset = queryset.annotate_user_roles(self.request.user)
         queryset = queryset.filter(user_roles__contains=[models.RoleChoices.OWNER])
 
         return self.get_response_for_queryset(queryset)
@@ -793,6 +757,8 @@ def tree(self, request, pk, *args, **kwargs):
         List ancestors tree above the document.
         What we need to display is the tree structure opened for the current document.
         """
+        user = self.request.user
+
         try:
             current_document = self.queryset.only("depth", "path").get(pk=pk)
         except models.Document.DoesNotExist as excpt:
@@ -847,8 +813,8 @@ def tree(self, request, pk, *args, **kwargs):
                 output_field=db.BooleanField(),
             )
         )
-        queryset = self.annotate_user_roles(queryset)
-        queryset = self.annotate_is_favorite(queryset)
+        queryset = queryset.annotate_user_roles(user)
+        queryset = queryset.annotate_is_favorite(user)
 
         # Pass ancestors' links definitions to the serializer as a context variable
         # in order to allow saving time while computing abilities on the instance

src/backend/core/models.py

@@ -464,6 +464,41 @@ def readable_per_se(self, user):
 
         return self.filter(link_reach=LinkReachChoices.PUBLIC)
 
+    def annotate_is_favorite(self, user):
+        """
+        Annotate document queryset with the favorite status for the current user.
+        """
+        if user.is_authenticated:
+            favorite_exists_subquery = DocumentFavorite.objects.filter(
+                document_id=models.OuterRef("pk"), user=user
+            )
+            return self.annotate(is_favorite=models.Exists(favorite_exists_subquery))
+
+        return self.annotate(is_favorite=models.Value(False))
+
+    def annotate_user_roles(self, user):
+        """
+        Annotate document queryset with the roles of the current user
+        on the document or its ancestors.
+        """
+        output_field = ArrayField(base_field=models.CharField())
+
+        if user.is_authenticated:
+            user_roles_subquery = DocumentAccess.objects.filter(
+                models.Q(user=user) | models.Q(team__in=user.teams),
+                document__path=Left(models.OuterRef("path"), Length("document__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 DocumentManager(MP_NodeManager.from_queryset(DocumentQuerySet)):
     """