From ce5af999c5805dc8cb8cca7ea47ad3b83ee3a646 Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Wed, 22 Jul 2026 15:58:07 -0700 Subject: [PATCH] Address review feedback: restore validation, scope aggregation to queryset - Restore the document_count_source_field validation helper dropped during the refactor -- misconfiguring a viewset (document_count_through set without document_count_source_field) now fails fast with a clear error again instead of an obscure runtime failure. - Restrict annotate_document_count_for_related_queryset()'s through-table aggregation to rows whose related_object_field is one of the annotated queryset's pks. No-op for the main tag-list call site (queryset is all tags), but avoids unnecessary work for narrower callers like the tag descendants branch in TagViewSet.list(). Co-authored-by: Claude Sonnet 5 --- src/documents/permissions.py | 12 +++++++++++- src/documents/views.py | 13 +++++++++++-- 2 files changed, 22 insertions(+), 3 deletions(-) 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), )