Fix: tag/custom-field document_count scales badly with tag count (#13203)

* Fix (beta): tag/custom-field document_count scales badly with tag count

TagViewSet and CustomFieldViewSet's document_count annotation used a
per-row correlated subquery (annotate_document_count_for_related_queryset),
executed once per tag/custom-field row. Fine at a few dozen rows, but at
~1,000 tags it degrades to a full per-tag GroupAggregate over the tags M2M
table -- 6s+ in production reports, confirmed via EXPLAIN (loops=1000).

Replaced with a single, independent GROUP BY over the through table
(permission filter expressed as a plain WHERE, not an aggregate FILTER),
then injected via Case/When. Also tried a more "obvious" fix -- a plain
Count(filter=Q(id__in=permitted_ids), distinct=True) directly on the M2M
relation -- but that's worse: Postgres fails to plan the id__in check as a
semi-join once a large M2M bridge table is involved, and instead re-checks
subquery membership once per joined row.

Benchmarked against a synthetic corpus (400k documents, 1,000 tags @ ~5/doc,
matching real-world reports in discussion #13161): tag list wall-clock drops
from ~180s to ~8s for a permission-restricted user, ~21s to ~8s for a
superuser. No measurable change at typical home-instance scale (tens of
tags).

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

src/documents/permissions.py

@@ -4,14 +4,14 @@ from django.contrib.auth.models import Group
 from django.contrib.auth.models import Permission
 from django.contrib.auth.models import User
 from django.contrib.contenttypes.models import ContentType
+from django.db.models import Case
 from django.db.models import Count
 from django.db.models import IntegerField
-from django.db.models import OuterRef
 from django.db.models import Q
 from django.db.models import QuerySet
-from django.db.models import Subquery
+from django.db.models import Value
+from django.db.models import When
 from django.db.models.functions import Cast
-from django.db.models.functions import Coalesce
 from guardian.core import ObjectPermissionChecker
 from guardian.models import GroupObjectPermission
 from guardian.models import UserObjectPermission
@@ -204,20 +204,24 @@ def _permitted_document_ids(user):
     ).values_list("id", flat=True)
 
 
-def get_document_count_filter_for_user(user):
+def get_document_count_filter_for_user(user, related_name: str = "documents"):
     """
     Return the Q object used to filter document counts for the given user.
 
     The filter is expressed as an ``id__in`` against a small subquery of permitted
     document IDs to keep the generated SQL simple and avoid large OR clauses.
+
+    ``related_name`` is the ORM path from the annotated model to Document (e.g.
+    ``"documents"`` for Tag's direct M2M, or ``"fields__document"`` for CustomField,
+    which only reaches Document via the CustomFieldInstance through-model).
     """
 
     if getattr(user, "is_superuser", False):
         # Superuser: no permission filtering needed
-        return Q(documents__deleted_at__isnull=True)
+        return Q(**{f"{related_name}__deleted_at__isnull": True})
 
     permitted_ids = _permitted_document_ids(user)
-    return Q(documents__id__in=permitted_ids)
+    return Q(**{f"{related_name}__id__in": permitted_ids})
 
 
 def annotate_document_count_for_related_queryset(
@@ -228,14 +232,33 @@ def annotate_document_count_for_related_queryset(
     user: User | None = None,
 ) -> QuerySet[Any]:
     """
-    Annotate a queryset with permissions-aware document counts using a subquery
-    against a relation table.
+    Annotate a queryset with a permissions-aware document count for a relation
+    to Document that goes through an M2M/through-model table (e.g. Tag via
+    ``Document.tags.through``, or CustomField via ``CustomFieldInstance``).
+
+    Counts are computed via a single, independent GROUP BY over the relation
+    table -- with the permission filter expressed as a plain ``WHERE`` rather
+    than an aggregate ``FILTER`` -- then injected via ``Case``/``When``. This
+    deliberately avoids two slower alternatives found while building this:
+
+    - A per-outer-row correlated subquery (one execution per row of the
+      annotated queryset): fine at a handful of rows, catastrophic once the
+      queryset has hundreds/thousands of rows.
+    - ``Count(..., filter=Q(id__in=permitted_ids), distinct=True)`` applied
+      directly to the M2M relation: Postgres can fail to plan the ``id__in``
+      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
                        or CustomFieldInstance)
-        source_field: field on the relation pointing back to queryset pk
+        related_object_field: field on the relation pointing back to queryset pk
         target_field: field on the relation pointing to Document id
         user: the user for whom to filter permitted document ids
     """
@@ -244,15 +267,27 @@ def annotate_document_count_for_related_queryset(
     counts = (
         through_model.objects.filter(
             **{
-                related_object_field: OuterRef("pk"),
+                f"{related_object_field}__in": queryset.values("pk"),
                 f"{target_field}__in": permitted_ids,
             },
         )
         .values(related_object_field)
-        .annotate(c=Count(target_field))
-        .values("c")
+        .annotate(c=Count(target_field, distinct=True))
+    )
+    counts_by_pk = {row[related_object_field]: row["c"] for row in counts}
+
+    if not counts_by_pk:
+        return queryset.annotate(
+            document_count=Value(0, output_field=IntegerField()),
+        )
+
+    return queryset.annotate(
+        document_count=Case(
+            *(When(pk=pk, then=Value(count)) for pk, count in counts_by_pk.items()),
+            default=Value(0),
+            output_field=IntegerField(),
+        ),
     )
-    return queryset.annotate(document_count=Coalesce(Subquery(counts[:1]), 0))
 
 
 def get_objects_for_user_owner_aware(

src/documents/views.py

@@ -482,7 +482,15 @@ class BulkPermissionMixin:
 class PermissionsAwareDocumentCountMixin(BulkPermissionMixin, PassUserMixin):
     """Mixin to add document count to queryset, permissions-aware if needed"""
 
-    # Default is simple relation path, override for through-table/count specialization.
+    # Direct FK/M2M relation name from this model to Document, used for the
+    # cheap Count(filter=...) path (Correspondent, DocumentType, StoragePath).
+    document_count_related_name: str = "documents"
+
+    # Set both of these instead, for models that only reach Document through
+    # an M2M/through-model table (Tag, CustomField). A plain Count(filter=...)
+    # over such a relation is fine for a direct FK, but forces a much more
+    # expensive plan once an M2M bridge table is involved -- see
+    # annotate_document_count_for_related_queryset() for why.
     document_count_through: type[Model] | None = None
     document_count_source_field: str | None = None
 
@@ -498,12 +506,14 @@ class PermissionsAwareDocumentCountMixin(BulkPermissionMixin, PassUserMixin):
     def get_document_count_filter(self):
         request = getattr(self, "request", None)
         user = getattr(request, "user", None) if request else None
-        return get_document_count_filter_for_user(user)
+        return get_document_count_filter_for_user(
+            user,
+            related_name=self.document_count_related_name,
+        )
 
     def get_queryset(self):
         base_qs = super().get_queryset()
 
-        # Use optimized through-table counting when configured.
         if self.document_count_through:
             user = getattr(getattr(self, "request", None), "user", None)
             return annotate_document_count_for_related_queryset(
@@ -513,10 +523,13 @@ class PermissionsAwareDocumentCountMixin(BulkPermissionMixin, PassUserMixin):
                 user=user,
             )
 
-        # Fallback: simple Count on relation with permission filter.
         filter = self.get_document_count_filter()
         return base_qs.annotate(
-            document_count=Count("documents", filter=filter),
+            document_count=Count(
+                self.document_count_related_name,
+                filter=filter,
+                distinct=True,
+            ),
         )