diff --git a/src/documents/permissions.py b/src/documents/permissions.py index b70b10c34..06e1e271a 100644 --- a/src/documents/permissions.py +++ b/src/documents/permissions.py @@ -249,6 +249,11 @@ def annotate_document_count_for_related_queryset( check as a semi-join and instead re-checks subquery membership once per row of the (much larger) M2M join -- worse than the correlated subquery. + Aggregation is restricted to rows whose ``related_object_field`` is one of + ``queryset``'s pks, so passing a subset (e.g. a handful of tag descendants) + doesn't pay the cost of counting for every row matching the permission + filter. + Args: queryset: base queryset to annotate (must contain pk) through_model: model representing the relation (e.g., Document.tags.through @@ -260,7 +265,12 @@ def annotate_document_count_for_related_queryset( permitted_ids = _permitted_document_ids(user) counts = ( - through_model.objects.filter(**{f"{target_field}__in": permitted_ids}) + through_model.objects.filter( + **{ + f"{related_object_field}__in": queryset.values("pk"), + f"{target_field}__in": permitted_ids, + }, + ) .values(related_object_field) .annotate(c=Count(target_field, distinct=True)) ) diff --git a/src/documents/views.py b/src/documents/views.py index 1a590f1c7..066f8fd8a 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -494,6 +494,15 @@ class PermissionsAwareDocumentCountMixin(BulkPermissionMixin, PassUserMixin): document_count_through: type[Model] | None = None document_count_source_field: str | None = None + def _get_document_count_source_field(self) -> str: + if self.document_count_source_field is None: + msg = ( + "document_count_source_field must be set when " + "document_count_through is configured" + ) + raise ValueError(msg) + return self.document_count_source_field + def get_document_count_filter(self): request = getattr(self, "request", None) user = getattr(request, "user", None) if request else None @@ -510,7 +519,7 @@ class PermissionsAwareDocumentCountMixin(BulkPermissionMixin, PassUserMixin): return annotate_document_count_for_related_queryset( base_qs, through_model=self.document_count_through, - related_object_field=self.document_count_source_field, + related_object_field=self._get_document_count_source_field(), user=user, ) @@ -613,7 +622,7 @@ class TagViewSet(PermissionsAwareDocumentCountMixin, ModelViewSet[Tag]): pk__in=descendant_pks | {t.pk for t in all_tags}, ).select_related("owner"), through_model=self.document_count_through, - related_object_field=self.document_count_source_field, + related_object_field=self._get_document_count_source_field(), user=user, ).order_by(*ordering), )