♻️(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.
This commit is contained in:
Samuel Paccoud - DINUM
2025-04-12 11:35:36 +02:00
committed by Anthony LC
parent a7c91f9443
commit df2b953e53
2 changed files with 49 additions and 48 deletions

View File

@@ -404,44 +404,6 @@ class DocumentViewSet(
trashbin_serializer_class = serializers.ListDocumentSerializer trashbin_serializer_class = serializers.ListDocumentSerializer
tree_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): def get_queryset(self):
"""Get queryset performing all annotation and filtering on the document tree structure.""" """Get queryset performing all annotation and filtering on the document tree structure."""
user = self.request.user user = self.request.user
@@ -477,8 +439,9 @@ class DocumentViewSet(
def filter_queryset(self, queryset): def filter_queryset(self, queryset):
"""Override to apply annotations to generic views.""" """Override to apply annotations to generic views."""
queryset = super().filter_queryset(queryset) queryset = super().filter_queryset(queryset)
queryset = self.annotate_is_favorite(queryset) user = self.request.user
queryset = self.annotate_user_roles(queryset) queryset = queryset.annotate_is_favorite(user)
queryset = queryset.annotate_user_roles(user)
return queryset return queryset
def get_response_for_queryset(self, queryset): def get_response_for_queryset(self, queryset):
@@ -502,9 +465,10 @@ class DocumentViewSet(
Additional annotations (e.g., `is_highest_ancestor_for_user`, favorite status) are Additional annotations (e.g., `is_highest_ancestor_for_user`, favorite status) are
applied before ordering and returning the response. applied before ordering and returning the response.
""" """
queryset = ( user = self.request.user
self.get_queryset()
) # Not calling filter_queryset. We do our own cooking. # Not calling filter_queryset. We do our own cooking.
queryset = self.get_queryset()
filterset = ListDocumentFilter( filterset = ListDocumentFilter(
self.request.GET, queryset=queryset, request=self.request self.request.GET, queryset=queryset, request=self.request
@@ -517,7 +481,7 @@ class DocumentViewSet(
for field in ["is_creator_me", "title"]: for field in ["is_creator_me", "title"]:
queryset = filterset.filters[field].filter(queryset, filter_data[field]) 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 # 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. # of each other. In this case we want to keep only the highest ancestors.
@@ -534,7 +498,7 @@ class DocumentViewSet(
) )
# Annotate favorite status and filter if applicable as late as possible # 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 = filterset.filters["is_favorite"].filter(
queryset, filter_data["is_favorite"] queryset, filter_data["is_favorite"]
) )
@@ -702,7 +666,7 @@ class DocumentViewSet(
deleted_at__isnull=False, deleted_at__isnull=False,
deleted_at__gte=models.get_trashbin_cutoff(), 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]) queryset = queryset.filter(user_roles__contains=[models.RoleChoices.OWNER])
return self.get_response_for_queryset(queryset) return self.get_response_for_queryset(queryset)
@@ -896,6 +860,8 @@ class DocumentViewSet(
List ancestors tree above the document. List ancestors tree above the document.
What we need to display is the tree structure opened for the current document. What we need to display is the tree structure opened for the current document.
""" """
user = self.request.user
try: try:
current_document = self.queryset.only("depth", "path").get(pk=pk) current_document = self.queryset.only("depth", "path").get(pk=pk)
except models.Document.DoesNotExist as excpt: except models.Document.DoesNotExist as excpt:
@@ -950,8 +916,8 @@ class DocumentViewSet(
output_field=db.BooleanField(), output_field=db.BooleanField(),
) )
) )
queryset = self.annotate_user_roles(queryset) queryset = queryset.annotate_user_roles(user)
queryset = self.annotate_is_favorite(queryset) queryset = queryset.annotate_is_favorite(user)
# Pass ancestors' links definitions to the serializer as a context variable # Pass ancestors' links definitions to the serializer as a context variable
# in order to allow saving time while computing abilities on the instance # in order to allow saving time while computing abilities on the instance

View File

@@ -464,6 +464,41 @@ class DocumentQuerySet(MP_NodeQuerySet):
return self.filter(link_reach=LinkReachChoices.PUBLIC) 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)): class DocumentManager(MP_NodeManager.from_queryset(DocumentQuerySet)):
""" """