Optimize stats api

config/settings/base.py

@@ -398,6 +398,11 @@
         "task": "documentcloud.users.tasks.permission_digest",
         "schedule": crontab(day_of_week="mon", hour=7, minute=0),
     },
+    "recompute_user_and_org_stats": {
+        "task": "documentcloud.core.tasks.recompute_user_and_org_stats",
+        "schedule": crontab(hour=2, minute=0),
+        "options": {"time_limit": 1800, "soft_time_limit": 1740},
+    },
 }
 
 PERMISSIONS_DIGEST_EMAILS = env.list("PERMISSIONS_DIGEST_EMAILS", default=[])

documentcloud/core/management/commands/recompute_user_and_org_stats.py

@@ -0,0 +1,53 @@
+# Django
+from django.conf import settings
+from django.core.management.base import BaseCommand
+from django.db.models import Count
+from django.utils import timezone
+
+# Standard Library
+from datetime import timedelta
+
+# DocumentCloud
+from documentcloud.documents.models import Document
+from documentcloud.organizations.stats_api.models import OrganizationStats
+from documentcloud.users.stats_api.models import UserStats
+
+BATCH = 1000
+
+
+class Command(BaseCommand):
+    """Recompute stored document counts on stats rows (run nightly)."""
+
+    help = "Recompute total_documents and recent_upload_count on stats rows"
+
+    def handle(self, *args, **options):
+        cutoff = timezone.now() - timedelta(days=settings.UPLOAD_WINDOW_DAYS)
+        self._recompute(UserStats, "user_id", cutoff)
+        self._recompute(OrganizationStats, "organization_id", cutoff)
+
+    def _recompute(self, model, key, cutoff):
+        totals = dict(Document.objects.values_list(key).annotate(c=Count("pk")))
+        recents = dict(
+            Document.objects.filter(created_at__gte=cutoff)
+            .values_list(key)
+            .annotate(c=Count("pk"))
+        )
+        to_update = []
+        updated = 0
+        for row in model.objects.all().iterator(chunk_size=BATCH):
+            k = getattr(row, key)
+            row.total_documents = totals.get(k, 0)
+            row.recent_upload_count = recents.get(k, 0)
+            to_update.append(row)
+            if len(to_update) >= BATCH:
+                model.objects.bulk_update(
+                    to_update, ["total_documents", "recent_upload_count"]
+                )
+                updated += len(to_update)
+                to_update = []
+        if to_update:
+            model.objects.bulk_update(
+                to_update, ["total_documents", "recent_upload_count"]
+            )
+            updated += len(to_update)
+        self.stdout.write(self.style.SUCCESS(f"{model.__name__}: {updated:,} updated"))

documentcloud/core/tasks.py

@@ -0,0 +1,9 @@
+# Django
+from celery import shared_task
+from django.core.management import call_command
+
+
+@shared_task
+def recompute_user_and_org_stats():
+    """Nightly recompute of stored document counts on stats rows."""
+    call_command("recompute_user_and_org_stats")

documentcloud/organizations/migrations/0025_organizationstats_recent_upload_count_and_more.py

@@ -0,0 +1,23 @@
+# Generated by Django 5.2.15 on 2026-09-10 19:12
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("organizations", "0024_organizationstats_last_ai_credit_at"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="organizationstats",
+            name="recent_upload_count",
+            field=models.IntegerField(default=0),
+        ),
+        migrations.AddField(
+            model_name="organizationstats",
+            name="total_documents",
+            field=models.IntegerField(default=0),
+        ),
+    ]

documentcloud/organizations/stats_api/models.py

@@ -11,6 +11,8 @@ class OrganizationStats(models.Model):
     )
     last_upload_at = models.DateTimeField(null=True, blank=True, db_index=True)
     last_ai_credit_at = models.DateTimeField(null=True, blank=True, db_index=True)
+    total_documents = models.IntegerField(default=0)
+    recent_upload_count = models.IntegerField(default=0)
 
     class Meta:
         app_label = "organizations"

documentcloud/organizations/stats_api/serializers.py

@@ -10,9 +10,10 @@ class OrganizationStatsSerializer(serializers.ModelSerializer):
     uuid = serializers.UUIDField(source="organization.uuid", read_only=True)
     total_documents = serializers.IntegerField(read_only=True)
     days_since_last_upload = serializers.SerializerMethodField()
-    recent_upload_count = serializers.SerializerMethodField(
-        help_text="Documents uploaded by the org within the configured window "
-        "(UPLOAD_WINDOW_DAYS, defaults to 90)."
+    recent_upload_count = serializers.IntegerField(
+        read_only=True,
+        help_text="Number of documents uploaded within the configured recent window "
+        "(UPLOAD_WINDOW_DAYS, currently defaults to 90).",
     )
     ai_credits = serializers.SerializerMethodField()
 
@@ -34,9 +35,6 @@ def get_days_since_last_upload(self, obj):
             return None
         return (timezone.now() - obj.last_upload_at).days
 
-    def get_recent_upload_count(self, obj):
-        return getattr(obj, "recent_upload_count", None)
-
     def get_ai_credits(self, obj):
         org = obj.organization
         return {

documentcloud/organizations/stats_api/tests.py

@@ -1,5 +1,6 @@
 # Django
 from django.conf import settings
+from django.core.management import call_command
 from django.test import TestCase
 from django.utils import timezone
 from rest_framework import status
@@ -107,7 +108,7 @@ def test_list_unauthenticated(self, client):
         assert response.status_code == status.HTTP_403_FORBIDDEN
 
     def test_list_excludes_individual_orgs(self, client):
-        """Users pull individual orgs in on the users endpoint"""
+        """Individual orgs are excluded from the org stats endpoint."""
         client.force_authenticate(user=self._admin())
         collective = OrganizationFactory.create_batch(3, individual=False)
         individual = OrganizationFactory(individual=True)
@@ -120,11 +121,12 @@ def test_list_excludes_individual_orgs(self, client):
         assert str(individual.uuid) not in uuids
 
     def test_retrieve_populates_enriched_fields(self, client):
-        """Regression test. org detail view populates the annotated counts, which
-        were previously only set in paginate_queryset (list view)."""
+        """Detail view serves the stored document counts, populated by the
+        recompute command."""
         client.force_authenticate(user=self._admin())
         org = OrganizationFactory(individual=False)
         DocumentFactory.create_batch(2, organization=org)
+        call_command("recompute_user_and_org_stats")
 
         response = client.get(f"/stats_api/organizations/{org.uuid}/")
         assert response.status_code == status.HTTP_200_OK
@@ -157,11 +159,8 @@ def test_document_upload_bumps_org_watermark(
         user, org = user_with_collective_org
         client.force_authenticate(user=user)
 
-        print("fixture org:", org.pk)
-        print("user.organization:", user.organization.pk)  # same as org.pk?
-
         response = client.post("/api/documents/", {"title": "t"})
-        print("status:", response.status_code, response.content[:200])  # created?
+        assert response.status_code == status.HTTP_201_CREATED
 
         stats = OrganizationStats.objects.get(organization=org)
         assert stats.last_upload_at is not None
@@ -170,7 +169,6 @@ def test_ai_credit_charge_bumps_org_watermark(
         self, client, user_with_collective_org
     ):
         """use_ai_credits' record_ai_credit_use call should bump the org watermark."""
-
         user, org = user_with_collective_org
         # give the org credits to spend
         org.monthly_ai_credits = 5

documentcloud/organizations/stats_api/views.py

@@ -1,6 +1,6 @@
 # Django
 from django.conf import settings
-from django.db.models import Count, Q
+from django.db.models import Exists, OuterRef
 from django.utils import timezone
 from django.utils.dateparse import parse_datetime
 from rest_framework import mixins, viewsets
@@ -16,7 +16,7 @@
 
 # DocumentCloud
 from documentcloud.core.pagination import CursorPagination
-from documentcloud.documents.choices import Status
+from documentcloud.documents.models import Document
 from documentcloud.organizations.stats_api.models import OrganizationStats
 from documentcloud.organizations.stats_api.serializers import (
     OrganizationStatsSerializer,
@@ -58,64 +58,40 @@ class Meta:
     filterset_class = Filter
 
     def get_queryset(self):
-        return OrganizationStats.objects.select_related("organization").filter(
-            organization__individual=False
-        )
-
-    def _annotate_and_prefetch(self, queryset):
-        cutoff = timezone.now() - timedelta(days=settings.UPLOAD_WINDOW_DAYS)
         return (
-            queryset.select_related("organization", "organization__parent")
-            .prefetch_related("organization__groups")
-            .annotate(
-                total_documents=Count(
-                    "organization__documents",
-                    filter=~Q(organization__documents__status=Status.deleted),
-                    distinct=True,
-                ),
-                recent_upload_count=Count(
-                    "organization__documents",
-                    filter=Q(organization__documents__created_at__gte=cutoff)
-                    & ~Q(organization__documents__status=Status.deleted),
-                    distinct=True,
-                ),
+            OrganizationStats.objects.select_related(
+                "organization", "organization__parent"
             )
+            .prefetch_related("organization__groups")
+            .filter(organization__individual=False)
         )
 
-    def paginate_queryset(self, queryset):
-        page = super().paginate_queryset(queryset)
-        annotated = self._annotate_and_prefetch(
-            OrganizationStats.objects.filter(pk__in=[o.pk for o in page])
-        ).order_by("pk")
-        return list(annotated)
-
-    def get_object(self):
-        obj = super().get_object()
-        return self._annotate_and_prefetch(
-            OrganizationStats.objects.filter(pk=obj.pk)
-        ).get()
-
     @action(detail=False, methods=["get"])
     def aged_out(self, request):
-        """Orgs with a document that crossed the window boundary since `since`,
-        so their upload count has dropped and needs re-syncing."""
+        """Orgs with a document that crossed the recent-upload window boundary
+        since `since`, so their recent_upload_count has dropped without any event.
+        Lets the caller (Squarelet) know which orgs to re-sync.
+        """
         since = request.query_params.get("since")
         if not since:
             return Response({"error": "since query param is required"}, status=400)
         since_dt = parse_datetime(since)
         if since_dt is None:
             return Response({"error": "since must be an ISO 8601 datetime"}, status=400)
+        if timezone.is_naive(since_dt):
+            since_dt = timezone.make_aware(since_dt, timezone.utc)
 
         win = timedelta(days=settings.UPLOAD_WINDOW_DAYS)
         now = timezone.now()
-        qs = (
-            self.get_queryset()
-            .filter(
-                organization__documents__created_at__gte=since_dt - win,
-                organization__documents__created_at__lt=now - win,
-            )
-            .distinct()
+
+        # Exists() short-circuits per org instead of joining + distinct over a
+        # heavy org's whole document set which can time out.
+        aged_doc = Document.objects.filter(
+            organization_id=OuterRef("organization_id"),
+            created_at__gte=since_dt - win,
+            created_at__lt=now - win,
         )
+        qs = self.get_queryset().filter(Exists(aged_doc))
 
         page = self.paginate_queryset(qs)
         return self.get_paginated_response(self.get_serializer(page, many=True).data)

documentcloud/users/migrations/0015_userstats_recent_upload_count_and_more.py

@@ -0,0 +1,23 @@
+# Generated by Django 5.2.15 on 2026-09-10 19:12
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("users", "0014_userstats_last_ai_credit_at"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="userstats",
+            name="recent_upload_count",
+            field=models.IntegerField(default=0),
+        ),
+        migrations.AddField(
+            model_name="userstats",
+            name="total_documents",
+            field=models.IntegerField(default=0),
+        ),
+    ]

documentcloud/users/stats_api/models.py

@@ -11,6 +11,8 @@ class UserStats(models.Model):
     )
     last_upload_at = models.DateTimeField(null=True, blank=True, db_index=True)
     last_ai_credit_at = models.DateTimeField(null=True, blank=True, db_index=True)
+    total_documents = models.IntegerField(default=0)
+    recent_upload_count = models.IntegerField(default=0)
 
     class Meta:
         app_label = "users"

documentcloud/users/stats_api/serializers.py

@@ -8,14 +8,15 @@
 
 class UserStatsSerializer(serializers.ModelSerializer):
     uuid = serializers.UUIDField(source="user.uuid", read_only=True)
-    total_documents = serializers.IntegerField(read_only=True)
     days_since_last_upload = serializers.SerializerMethodField(
         help_text="Number of days since the last time the user uploaded a document"
     )
     last_login_at = serializers.DateTimeField(source="user.last_login", read_only=True)
-    recent_upload_count = serializers.SerializerMethodField(
+    total_documents = serializers.IntegerField(read_only=True)
+    recent_upload_count = serializers.IntegerField(
+        read_only=True,
         help_text="Number of documents uploaded within the configured recent window "
-        "(UPLOAD_WINDOW_DAYS, currently defaults to 90)."
+        "(UPLOAD_WINDOW_DAYS, currently defaults to 90).",
     )
     individual_ai_credits = serializers.SerializerMethodField()
 
@@ -38,9 +39,6 @@ def get_days_since_last_upload(self, obj):
             return None
         return (timezone.now() - obj.last_upload_at).days
 
-    def get_recent_upload_count(self, obj):
-        return getattr(obj, "recent_upload_count", None)
-
     def get_individual_ai_credits(self, obj):
         orgs = getattr(obj.user, "individual_orgs", [])
         if not orgs:

documentcloud/users/stats_api/tests.py

@@ -1,4 +1,5 @@
 # Django
+from django.core.management import call_command
 from django.utils import timezone
 from rest_framework import status
 
@@ -40,9 +41,8 @@ def test_list_unauthenticated(self, client):
         assert response.status_code == status.HTTP_403_FORBIDDEN
 
     def test_retrieve_populates_enriched_fields(self, client):
-        """Regression: detail view must populate individual_ai_credits and the
-        annotated counts, not just the list view (they were only set in
-        paginate_queryset before)."""
+        """Detail view serves the stored document counts (populated by the
+        recompute command) and the live-computed individual_ai_credits."""
         admin = self._admin()
         client.force_authenticate(user=admin)
 
@@ -52,8 +52,9 @@ def test_retrieve_populates_enriched_fields(self, client):
         org.monthly_ai_credits = 5
         org.ai_credits_per_month = 10
         org.save()
-        # give them documents so the counts are non-zero
+        # give them documents, then recompute so the stored counts reflect them
         DocumentFactory.create_batch(2, user=target, organization=org)
+        call_command("recompute_user_and_org_stats")
 
         response = client.get(f"/stats_api/users/{target.uuid}/")
         assert response.status_code == status.HTTP_200_OK