From c58dd9d242cb7cc6b4402544be91a5fa0a932a00 Mon Sep 17 00:00:00 2001 From: duckduckgrayduck <102841251+duckduckgrayduck@users.noreply.github.com> Date: Mon, 3 Aug 2026 14:22:54 -0500 Subject: [PATCH 1/4] First pass at stats_api --- config/settings/base.py | 3 + config/urls.py | 7 + .../0023_organizationstats_and_more.py | 264 ++++++++++++++++++ documentcloud/organizations/models.py | 3 + .../organizations/stats_api/__init__.py | 0 .../organizations/stats_api/models.py | 18 ++ .../organizations/stats_api/serializers.py | 35 +++ .../organizations/stats_api/views.py | 103 +++++++ ...s_alter_user_document_language_and_more.py | 264 ++++++++++++++++++ documentcloud/users/models.py | 3 + documentcloud/users/signals.py | 26 ++ documentcloud/users/stats_api/__init__.py | 0 documentcloud/users/stats_api/models.py | 18 ++ documentcloud/users/stats_api/serializers.py | 39 +++ documentcloud/users/stats_api/tests.py | 0 documentcloud/users/stats_api/views.py | 126 +++++++++ 16 files changed, 909 insertions(+) create mode 100644 documentcloud/organizations/migrations/0023_organizationstats_and_more.py create mode 100644 documentcloud/organizations/stats_api/__init__.py create mode 100644 documentcloud/organizations/stats_api/models.py create mode 100644 documentcloud/organizations/stats_api/serializers.py create mode 100644 documentcloud/organizations/stats_api/views.py create mode 100644 documentcloud/users/migrations/0013_userstats_alter_user_document_language_and_more.py create mode 100644 documentcloud/users/stats_api/__init__.py create mode 100644 documentcloud/users/stats_api/models.py create mode 100644 documentcloud/users/stats_api/serializers.py create mode 100644 documentcloud/users/stats_api/tests.py create mode 100644 documentcloud/users/stats_api/views.py diff --git a/config/settings/base.py b/config/settings/base.py index 970acd14..469e2953 100644 --- a/config/settings/base.py +++ b/config/settings/base.py @@ -694,3 +694,6 @@ # ------------------------------------------------------------------------------ MAX_PAGES = env.int("MAX_PAGES", default=50) GRAFT_DEBUG = env.bool("GRAFT_DEBUG", default=False) + +# STATS API +UPLOAD_WINDOW_DAYS = env.int("UPLOAD_WINDOW_DAYS", default=90) diff --git a/config/urls.py b/config/urls.py index 0727eded..656159c6 100644 --- a/config/urls.py +++ b/config/urls.py @@ -43,6 +43,7 @@ from documentcloud.drf_bulk.routers import BulkDefaultRouter, BulkRouterMixin from documentcloud.entities.views import EntityOccurrenceViewSet, EntityViewSet from documentcloud.flatpages.views import FlatPageViewSet +from documentcloud.organizations.stats_api.views import OrganizationStatsViewSet from documentcloud.organizations.views import OrganizationViewSet from documentcloud.projects.views import ( CollaborationViewSet, @@ -50,6 +51,7 @@ ProjectViewSet, ) from documentcloud.statistics.views import StatisticsViewSet +from documentcloud.users.stats_api.views import UserStatsViewSet from documentcloud.users.views import MessageView, UserViewSet @@ -94,6 +96,10 @@ class BulkNestedDefaultRouter(BulkRouterMixin, NestedDefaultRouter): router.register("documents/search/saved", SavedSearchViewSet, basename="saved_search") +stats_router = BulkDefaultRouter() +stats_router.register("users", UserStatsViewSet, basename="user-stats") +stats_router.register("organizations", OrganizationStatsViewSet, basename="org-stats") + urlpatterns = [ path("", RedirectView.as_view(url="/api/"), name="index"), path(settings.ADMIN_URL, admin.site.urls), @@ -138,6 +144,7 @@ class BulkNestedDefaultRouter(BulkRouterMixin, NestedDefaultRouter): path( "addons/dashboard/scraper/", scraper_dashboard, name="addon-scraper-dashboard" ), + path("stats_api/", include(stats_router.urls)), ] if "debug_toolbar" in settings.INSTALLED_APPS: diff --git a/documentcloud/organizations/migrations/0023_organizationstats_and_more.py b/documentcloud/organizations/migrations/0023_organizationstats_and_more.py new file mode 100644 index 00000000..dff626ca --- /dev/null +++ b/documentcloud/organizations/migrations/0023_organizationstats_and_more.py @@ -0,0 +1,264 @@ +# Generated by Django 5.2.15 on 2026-08-03 15:06 + +import django.db.models.deletion +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("organizations", "0022_organization_created_at_organization_updated_at"), + ] + + operations = [ + migrations.CreateModel( + name="OrganizationStats", + fields=[ + ( + "organization", + models.OneToOneField( + on_delete=django.db.models.deletion.CASCADE, + primary_key=True, + related_name="stats", + serialize=False, + to=settings.SQUARELET_ORGANIZATION_MODEL, + ), + ), + ( + "last_upload_at", + models.DateTimeField(blank=True, db_index=True, null=True), + ), + ], + ), + migrations.AlterField( + model_name="organization", + name="document_language", + field=models.CharField( + blank=True, + choices=[ + ("afr", "Afrikaans"), + ("amh", "Amharic"), + ("ara", "Arabic"), + ("asm", "Assamese"), + ("aze", "Azerbaijani"), + ("aze_cyrl", "Azerbaijani - Cyrillic"), + ("bel", "Belarusian"), + ("ben", "Bengali"), + ("bod", "Tibetan"), + ("bos", "Bosnian"), + ("bul", "Bulgarian"), + ("cat", "Catalan; Valencian"), + ("ceb", "Cebuano"), + ("ces", "Czech"), + ("zho", "chi_sim"), + ("tra", "chi_tra"), + ("chr", "Cherokee"), + ("cym", "Welsh"), + ("dan", "Danish"), + ("deu", "German"), + ("dzo", "Dzongkha"), + ("ell", "Greek"), + ("eng", "English"), + ("enm", "Middle English"), + ("epo", "Esperanto"), + ("est", "Estonian"), + ("eus", "Basque"), + ("fas", "Persian"), + ("fin", "Finnish"), + ("fra", "French"), + ("frk", "German Fraktur"), + ("frm", "Middle French"), + ("gle", "Irish"), + ("glg", "Galician"), + ("grc", "Ancient Greek"), + ("guj", "Gujarati"), + ("hat", "Haitian; Haitian Creole"), + ("heb", "Hebrew"), + ("hin", "Hindi"), + ("hrv", "Croatian"), + ("hun", "Hungarian"), + ("iku", "Inuktitut"), + ("ind", "Indonesian"), + ("isl", "Icelandic"), + ("ita", "Italian"), + ("ita_old", "Italian - Old"), + ("jav", "Javanese"), + ("jpn", "Japanese"), + ("kan", "Kannada"), + ("kat", "Georgian"), + ("kat_old", "Georgian - Old"), + ("kaz", "Kazakh"), + ("khm", "Central Khmer"), + ("kir", "Kirghiz; Kyrgyz"), + ("kor", "Korean"), + ("kur", "Kurdish"), + ("lao", "Lao"), + ("lat", "Latin"), + ("lav", "Latvian"), + ("lit", "Lithuanian"), + ("mal", "Malayalam"), + ("mar", "Marathi"), + ("mkd", "Macedonian"), + ("mlt", "Maltese"), + ("msa", "Malay"), + ("mya", "Burmese"), + ("nep", "Nepali"), + ("nld", "Dutch; Flemish"), + ("nor", "Norwegian"), + ("ori", "Oriya"), + ("pan", "Panjabi; Punjabi"), + ("pol", "Polish"), + ("por", "Portuguese"), + ("pus", "Pushto; Pashto"), + ("ron", "Romanian; Moldavian; Moldovan"), + ("rus", "Russian"), + ("san", "Sanskrit"), + ("sin", "Sinhala; Sinhalese"), + ("slk", "Slovak"), + ("slv", "Slovenian"), + ("spa", "Spanish; Castilian"), + ("spa_old", "Spanish; Castilian - Old"), + ("sqi", "Albanian"), + ("srp", "Serbian"), + ("srp_latn", "Serbian - Latin"), + ("swa", "Swahili"), + ("swe", "Swedish"), + ("syr", "Syriac"), + ("tam", "Tamil"), + ("tel", "Telugu"), + ("tgk", "Tajik"), + ("tgl", "Tagalog"), + ("tha", "Thai"), + ("tir", "Tigrinya"), + ("tur", "Turkish"), + ("uig", "Uighur; Uyghur"), + ("ukr", "Ukrainian"), + ("urd", "Urdu"), + ("uzb", "Uzbek"), + ("uzb_cyrl", "Uzbek - Cyrillic"), + ("vie", "Vietnamese"), + ("yid", "Yiddish"), + ], + default="eng", + help_text="The default document language for user's in this organization", + max_length=8, + verbose_name="document language", + ), + ), + migrations.AlterField( + model_name="organization", + name="language", + field=models.CharField( + blank=True, + choices=[ + ("afr", "Afrikaans"), + ("amh", "Amharic"), + ("ara", "Arabic"), + ("asm", "Assamese"), + ("aze", "Azerbaijani"), + ("aze_cyrl", "Azerbaijani - Cyrillic"), + ("bel", "Belarusian"), + ("ben", "Bengali"), + ("bod", "Tibetan"), + ("bos", "Bosnian"), + ("bul", "Bulgarian"), + ("cat", "Catalan; Valencian"), + ("ceb", "Cebuano"), + ("ces", "Czech"), + ("zho", "chi_sim"), + ("tra", "chi_tra"), + ("chr", "Cherokee"), + ("cym", "Welsh"), + ("dan", "Danish"), + ("deu", "German"), + ("dzo", "Dzongkha"), + ("ell", "Greek"), + ("eng", "English"), + ("enm", "Middle English"), + ("epo", "Esperanto"), + ("est", "Estonian"), + ("eus", "Basque"), + ("fas", "Persian"), + ("fin", "Finnish"), + ("fra", "French"), + ("frk", "German Fraktur"), + ("frm", "Middle French"), + ("gle", "Irish"), + ("glg", "Galician"), + ("grc", "Ancient Greek"), + ("guj", "Gujarati"), + ("hat", "Haitian; Haitian Creole"), + ("heb", "Hebrew"), + ("hin", "Hindi"), + ("hrv", "Croatian"), + ("hun", "Hungarian"), + ("iku", "Inuktitut"), + ("ind", "Indonesian"), + ("isl", "Icelandic"), + ("ita", "Italian"), + ("ita_old", "Italian - Old"), + ("jav", "Javanese"), + ("jpn", "Japanese"), + ("kan", "Kannada"), + ("kat", "Georgian"), + ("kat_old", "Georgian - Old"), + ("kaz", "Kazakh"), + ("khm", "Central Khmer"), + ("kir", "Kirghiz; Kyrgyz"), + ("kor", "Korean"), + ("kur", "Kurdish"), + ("lao", "Lao"), + ("lat", "Latin"), + ("lav", "Latvian"), + ("lit", "Lithuanian"), + ("mal", "Malayalam"), + ("mar", "Marathi"), + ("mkd", "Macedonian"), + ("mlt", "Maltese"), + ("msa", "Malay"), + ("mya", "Burmese"), + ("nep", "Nepali"), + ("nld", "Dutch; Flemish"), + ("nor", "Norwegian"), + ("ori", "Oriya"), + ("pan", "Panjabi; Punjabi"), + ("pol", "Polish"), + ("por", "Portuguese"), + ("pus", "Pushto; Pashto"), + ("ron", "Romanian; Moldavian; Moldovan"), + ("rus", "Russian"), + ("san", "Sanskrit"), + ("sin", "Sinhala; Sinhalese"), + ("slk", "Slovak"), + ("slv", "Slovenian"), + ("spa", "Spanish; Castilian"), + ("spa_old", "Spanish; Castilian - Old"), + ("sqi", "Albanian"), + ("srp", "Serbian"), + ("srp_latn", "Serbian - Latin"), + ("swa", "Swahili"), + ("swe", "Swedish"), + ("syr", "Syriac"), + ("tam", "Tamil"), + ("tel", "Telugu"), + ("tgk", "Tajik"), + ("tgl", "Tagalog"), + ("tha", "Thai"), + ("tir", "Tigrinya"), + ("tur", "Turkish"), + ("uig", "Uighur; Uyghur"), + ("ukr", "Ukrainian"), + ("urd", "Urdu"), + ("uzb", "Uzbek"), + ("uzb_cyrl", "Uzbek - Cyrillic"), + ("vie", "Vietnamese"), + ("yid", "Yiddish"), + ], + default="eng", + help_text="The default interface language for user's in this organization", + max_length=8, + verbose_name="language", + ), + ), + ] diff --git a/documentcloud/organizations/models.py b/documentcloud/organizations/models.py index 422a7663..24470a5e 100644 --- a/documentcloud/organizations/models.py +++ b/documentcloud/organizations/models.py @@ -19,6 +19,9 @@ from documentcloud.core.fields import AutoCreatedField, AutoLastModifiedField from documentcloud.organizations.exceptions import InsufficientAICreditsError from documentcloud.organizations.querysets import OrganizationQuerySet +from documentcloud.organizations.stats_api.models import ( # pylint: disable=unused-import + OrganizationStats, +) logger = logging.getLogger(__name__) diff --git a/documentcloud/organizations/stats_api/__init__.py b/documentcloud/organizations/stats_api/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/documentcloud/organizations/stats_api/models.py b/documentcloud/organizations/stats_api/models.py new file mode 100644 index 00000000..df5510c6 --- /dev/null +++ b/documentcloud/organizations/stats_api/models.py @@ -0,0 +1,18 @@ +# Django +from django.db import models + + +class OrganizationStats(models.Model): + organization = models.OneToOneField( + "organizations.Organization", + on_delete=models.CASCADE, + related_name="stats", + primary_key=True, + ) + last_upload_at = models.DateTimeField(null=True, blank=True, db_index=True) + + class Meta: + app_label = "organizations" + + def __str__(self): + return f"Stats for organization {self.organization_id}" diff --git a/documentcloud/organizations/stats_api/serializers.py b/documentcloud/organizations/stats_api/serializers.py new file mode 100644 index 00000000..c5372002 --- /dev/null +++ b/documentcloud/organizations/stats_api/serializers.py @@ -0,0 +1,35 @@ +# Django +from django.utils import timezone +from rest_framework import serializers + +# DocumentCloud +from documentcloud.organizations.stats_api.models import OrganizationStats + + +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)." + ) + + class Meta: + model = OrganizationStats + fields = [ + "uuid", + "total_documents", + "last_upload_at", + "days_since_last_upload", + "recent_upload_count", + ] + read_only_fields = fields + + def get_days_since_last_upload(self, obj): + if obj.last_upload_at is None: + return None + return (timezone.now() - obj.last_upload_at).days + + def get_recent_upload_count(self, obj): + return getattr(obj, "recent_upload_count", None) diff --git a/documentcloud/organizations/stats_api/views.py b/documentcloud/organizations/stats_api/views.py new file mode 100644 index 00000000..4bab06a5 --- /dev/null +++ b/documentcloud/organizations/stats_api/views.py @@ -0,0 +1,103 @@ +# Django +from django.conf import settings +from django.db.models import Count, Q +from django.utils import timezone +from django.utils.dateparse import parse_datetime +from rest_framework import mixins, viewsets +from rest_framework.decorators import action +from rest_framework.permissions import IsAdminUser +from rest_framework.response import Response + +# Standard Library +from datetime import timedelta + +# Third Party +from django_filters import rest_framework as django_filters + +# DocumentCloud +from documentcloud.core.pagination import CursorCountPagination +from documentcloud.documents.choices import Status +from documentcloud.organizations.stats_api.models import OrganizationStats +from documentcloud.organizations.stats_api.serializers import ( + OrganizationStatsSerializer, +) + + +class OrganizationStatsViewSet( + mixins.ListModelMixin, + mixins.RetrieveModelMixin, + viewsets.GenericViewSet, +): + serializer_class = OrganizationStatsSerializer + permission_classes = [IsAdminUser] + filter_backends = [django_filters.DjangoFilterBackend] + pagination_class = CursorCountPagination + lookup_field = "organization__uuid" + lookup_url_kwarg = "uuid" + + class Filter(django_filters.FilterSet): + uploaded_within_days = django_filters.NumberFilter( + method="filter_uploaded_within_days", + label="Uploaded in last N days", + help_text=( + "Return orgs whose most recent upload was within the last N days." + ), + ) + + def filter_uploaded_within_days(self, queryset, _name, value): + days = int(value) + if days < 0: + return queryset.none() + cutoff = timezone.now() - timedelta(days=days) + return queryset.filter(last_upload_at__gte=cutoff) + + class Meta: + model = OrganizationStats + fields = [] + + filterset_class = Filter + + def get_queryset(self): + cutoff = timezone.now() - timedelta(days=settings.UPLOAD_WINDOW_DAYS) + return ( + OrganizationStats.objects.select_related("organization") + .filter(organization__individual=False) + .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, + ), + ) + ) + + @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.""" + 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) + + 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() + ) + + page = self.paginate_queryset(qs) + return Response(self.get_serializer(page, many=True).data) diff --git a/documentcloud/users/migrations/0013_userstats_alter_user_document_language_and_more.py b/documentcloud/users/migrations/0013_userstats_alter_user_document_language_and_more.py new file mode 100644 index 00000000..e2ac55ec --- /dev/null +++ b/documentcloud/users/migrations/0013_userstats_alter_user_document_language_and_more.py @@ -0,0 +1,264 @@ +# Generated by Django 5.2.15 on 2026-08-01 21:22 + +import django.db.models.deletion +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("users", "0012_default_pinned_projects"), + ] + + operations = [ + migrations.CreateModel( + name="UserStats", + fields=[ + ( + "user", + models.OneToOneField( + on_delete=django.db.models.deletion.CASCADE, + primary_key=True, + related_name="stats", + serialize=False, + to=settings.AUTH_USER_MODEL, + ), + ), + ( + "last_upload_at", + models.DateTimeField(blank=True, db_index=True, null=True), + ), + ], + ), + migrations.AlterField( + model_name="user", + name="document_language", + field=models.CharField( + blank=True, + choices=[ + ("afr", "Afrikaans"), + ("amh", "Amharic"), + ("ara", "Arabic"), + ("asm", "Assamese"), + ("aze", "Azerbaijani"), + ("aze_cyrl", "Azerbaijani - Cyrillic"), + ("bel", "Belarusian"), + ("ben", "Bengali"), + ("bod", "Tibetan"), + ("bos", "Bosnian"), + ("bul", "Bulgarian"), + ("cat", "Catalan; Valencian"), + ("ceb", "Cebuano"), + ("ces", "Czech"), + ("zho", "chi_sim"), + ("tra", "chi_tra"), + ("chr", "Cherokee"), + ("cym", "Welsh"), + ("dan", "Danish"), + ("deu", "German"), + ("dzo", "Dzongkha"), + ("ell", "Greek"), + ("eng", "English"), + ("enm", "Middle English"), + ("epo", "Esperanto"), + ("est", "Estonian"), + ("eus", "Basque"), + ("fas", "Persian"), + ("fin", "Finnish"), + ("fra", "French"), + ("frk", "German Fraktur"), + ("frm", "Middle French"), + ("gle", "Irish"), + ("glg", "Galician"), + ("grc", "Ancient Greek"), + ("guj", "Gujarati"), + ("hat", "Haitian; Haitian Creole"), + ("heb", "Hebrew"), + ("hin", "Hindi"), + ("hrv", "Croatian"), + ("hun", "Hungarian"), + ("iku", "Inuktitut"), + ("ind", "Indonesian"), + ("isl", "Icelandic"), + ("ita", "Italian"), + ("ita_old", "Italian - Old"), + ("jav", "Javanese"), + ("jpn", "Japanese"), + ("kan", "Kannada"), + ("kat", "Georgian"), + ("kat_old", "Georgian - Old"), + ("kaz", "Kazakh"), + ("khm", "Central Khmer"), + ("kir", "Kirghiz; Kyrgyz"), + ("kor", "Korean"), + ("kur", "Kurdish"), + ("lao", "Lao"), + ("lat", "Latin"), + ("lav", "Latvian"), + ("lit", "Lithuanian"), + ("mal", "Malayalam"), + ("mar", "Marathi"), + ("mkd", "Macedonian"), + ("mlt", "Maltese"), + ("msa", "Malay"), + ("mya", "Burmese"), + ("nep", "Nepali"), + ("nld", "Dutch; Flemish"), + ("nor", "Norwegian"), + ("ori", "Oriya"), + ("pan", "Panjabi; Punjabi"), + ("pol", "Polish"), + ("por", "Portuguese"), + ("pus", "Pushto; Pashto"), + ("ron", "Romanian; Moldavian; Moldovan"), + ("rus", "Russian"), + ("san", "Sanskrit"), + ("sin", "Sinhala; Sinhalese"), + ("slk", "Slovak"), + ("slv", "Slovenian"), + ("spa", "Spanish; Castilian"), + ("spa_old", "Spanish; Castilian - Old"), + ("sqi", "Albanian"), + ("srp", "Serbian"), + ("srp_latn", "Serbian - Latin"), + ("swa", "Swahili"), + ("swe", "Swedish"), + ("syr", "Syriac"), + ("tam", "Tamil"), + ("tel", "Telugu"), + ("tgk", "Tajik"), + ("tgl", "Tagalog"), + ("tha", "Thai"), + ("tir", "Tigrinya"), + ("tur", "Turkish"), + ("uig", "Uighur; Uyghur"), + ("ukr", "Ukrainian"), + ("urd", "Urdu"), + ("uzb", "Uzbek"), + ("uzb_cyrl", "Uzbek - Cyrillic"), + ("vie", "Vietnamese"), + ("yid", "Yiddish"), + ], + default="eng", + help_text="The default language for documents uploaded by this user", + max_length=8, + verbose_name="document language", + ), + ), + migrations.AlterField( + model_name="user", + name="language", + field=models.CharField( + blank=True, + choices=[ + ("afr", "Afrikaans"), + ("amh", "Amharic"), + ("ara", "Arabic"), + ("asm", "Assamese"), + ("aze", "Azerbaijani"), + ("aze_cyrl", "Azerbaijani - Cyrillic"), + ("bel", "Belarusian"), + ("ben", "Bengali"), + ("bod", "Tibetan"), + ("bos", "Bosnian"), + ("bul", "Bulgarian"), + ("cat", "Catalan; Valencian"), + ("ceb", "Cebuano"), + ("ces", "Czech"), + ("zho", "chi_sim"), + ("tra", "chi_tra"), + ("chr", "Cherokee"), + ("cym", "Welsh"), + ("dan", "Danish"), + ("deu", "German"), + ("dzo", "Dzongkha"), + ("ell", "Greek"), + ("eng", "English"), + ("enm", "Middle English"), + ("epo", "Esperanto"), + ("est", "Estonian"), + ("eus", "Basque"), + ("fas", "Persian"), + ("fin", "Finnish"), + ("fra", "French"), + ("frk", "German Fraktur"), + ("frm", "Middle French"), + ("gle", "Irish"), + ("glg", "Galician"), + ("grc", "Ancient Greek"), + ("guj", "Gujarati"), + ("hat", "Haitian; Haitian Creole"), + ("heb", "Hebrew"), + ("hin", "Hindi"), + ("hrv", "Croatian"), + ("hun", "Hungarian"), + ("iku", "Inuktitut"), + ("ind", "Indonesian"), + ("isl", "Icelandic"), + ("ita", "Italian"), + ("ita_old", "Italian - Old"), + ("jav", "Javanese"), + ("jpn", "Japanese"), + ("kan", "Kannada"), + ("kat", "Georgian"), + ("kat_old", "Georgian - Old"), + ("kaz", "Kazakh"), + ("khm", "Central Khmer"), + ("kir", "Kirghiz; Kyrgyz"), + ("kor", "Korean"), + ("kur", "Kurdish"), + ("lao", "Lao"), + ("lat", "Latin"), + ("lav", "Latvian"), + ("lit", "Lithuanian"), + ("mal", "Malayalam"), + ("mar", "Marathi"), + ("mkd", "Macedonian"), + ("mlt", "Maltese"), + ("msa", "Malay"), + ("mya", "Burmese"), + ("nep", "Nepali"), + ("nld", "Dutch; Flemish"), + ("nor", "Norwegian"), + ("ori", "Oriya"), + ("pan", "Panjabi; Punjabi"), + ("pol", "Polish"), + ("por", "Portuguese"), + ("pus", "Pushto; Pashto"), + ("ron", "Romanian; Moldavian; Moldovan"), + ("rus", "Russian"), + ("san", "Sanskrit"), + ("sin", "Sinhala; Sinhalese"), + ("slk", "Slovak"), + ("slv", "Slovenian"), + ("spa", "Spanish; Castilian"), + ("spa_old", "Spanish; Castilian - Old"), + ("sqi", "Albanian"), + ("srp", "Serbian"), + ("srp_latn", "Serbian - Latin"), + ("swa", "Swahili"), + ("swe", "Swedish"), + ("syr", "Syriac"), + ("tam", "Tamil"), + ("tel", "Telugu"), + ("tgk", "Tajik"), + ("tgl", "Tagalog"), + ("tha", "Thai"), + ("tir", "Tigrinya"), + ("tur", "Turkish"), + ("uig", "Uighur; Uyghur"), + ("ukr", "Ukrainian"), + ("urd", "Urdu"), + ("uzb", "Uzbek"), + ("uzb_cyrl", "Uzbek - Cyrillic"), + ("vie", "Vietnamese"), + ("yid", "Yiddish"), + ], + default="eng", + help_text="The interface language for this user", + max_length=8, + verbose_name="language", + ), + ), + ] diff --git a/documentcloud/users/models.py b/documentcloud/users/models.py index f455f06e..1e0d6d51 100644 --- a/documentcloud/users/models.py +++ b/documentcloud/users/models.py @@ -11,6 +11,9 @@ # DocumentCloud from documentcloud.core.choices import Language from documentcloud.users.managers import UserManager +from documentcloud.users.stats_api.models import ( # pylint: disable=unused-import + UserStats, +) class User(SAUser): diff --git a/documentcloud/users/signals.py b/documentcloud/users/signals.py index 4a7cdd0f..0a8d2b61 100644 --- a/documentcloud/users/signals.py +++ b/documentcloud/users/signals.py @@ -1,6 +1,8 @@ # Django from django.contrib.auth.signals import user_logged_in +from django.db.models.signals import post_save from django.dispatch import receiver +from django.utils import timezone # DocumentCloud from documentcloud.addons.models import AddOn @@ -14,3 +16,27 @@ def default_addons(sender, user, request, **kwargs): if not user.active_addons.exists(): user.active_addons.set(AddOn.objects.filter(default=True)) + + +@receiver(post_save, dispatch_uid="documentcloud.users.signals.touch_last_upload") +def touch_last_upload(sender, instance, created, **kwargs): + """Record upload recency when a user creates a new document""" + # DocumentCloud + from documentcloud.documents.models import Document + from documentcloud.organizations.stats_api.models import OrganizationStats + from documentcloud.users.stats_api.models import UserStats + + if sender is not Document or not created: + return + + now = timezone.now() + UserStats.objects.update_or_create( + user_id=instance.user_id, + defaults={"last_upload_at": now}, + ) + + if not instance.organization.individual: + OrganizationStats.objects.update_or_create( + organization_id=instance.organization_id, + defaults={"last_upload_at": now}, + ) diff --git a/documentcloud/users/stats_api/__init__.py b/documentcloud/users/stats_api/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/documentcloud/users/stats_api/models.py b/documentcloud/users/stats_api/models.py new file mode 100644 index 00000000..5f7cb50d --- /dev/null +++ b/documentcloud/users/stats_api/models.py @@ -0,0 +1,18 @@ +# Django +from django.db import models + + +class UserStats(models.Model): + user = models.OneToOneField( + "users.User", + on_delete=models.CASCADE, + related_name="stats", + primary_key=True, + ) + last_upload_at = models.DateTimeField(null=True, blank=True, db_index=True) + + class Meta: + app_label = "users" + + def __str__(self): + return f"Stats for user {self.user_id}" diff --git a/documentcloud/users/stats_api/serializers.py b/documentcloud/users/stats_api/serializers.py new file mode 100644 index 00000000..6cdc6726 --- /dev/null +++ b/documentcloud/users/stats_api/serializers.py @@ -0,0 +1,39 @@ +# Django +from django.utils import timezone +from rest_framework import serializers + +# DocumentCloud +from documentcloud.users.stats_api.models import UserStats + + +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( + help_text="Number of documents uploaded within the configured recent window " + "(UPLOAD_WINDOW_DAYS, currently defaults to 90)." + ) + + class Meta: + model = UserStats + fields = [ + "uuid", + "total_documents", + "last_upload_at", + "days_since_last_upload", + "last_login_at", + "recent_upload_count", + ] + read_only_fields = fields + + def get_days_since_last_upload(self, obj): + if obj.last_upload_at is None: + return None + return (timezone.now() - obj.last_upload_at).days + + def get_recent_upload_count(self, obj): + return getattr(obj, "recent_upload_count", None) diff --git a/documentcloud/users/stats_api/tests.py b/documentcloud/users/stats_api/tests.py new file mode 100644 index 00000000..e69de29b diff --git a/documentcloud/users/stats_api/views.py b/documentcloud/users/stats_api/views.py new file mode 100644 index 00000000..f58c6b3b --- /dev/null +++ b/documentcloud/users/stats_api/views.py @@ -0,0 +1,126 @@ +# Django +from django.conf import settings +from django.db.models import Count, Q +from django.utils import timezone +from django.utils.dateparse import parse_datetime +from rest_framework import mixins, viewsets +from rest_framework.decorators import action +from rest_framework.permissions import IsAdminUser +from rest_framework.response import Response + +# Standard Library +from datetime import timedelta + +# Third Party +from django_filters import rest_framework as django_filters + +# DocumentCloud +from documentcloud.core.pagination import CursorCountPagination +from documentcloud.documents.choices import Status +from documentcloud.users.stats_api.models import UserStats +from documentcloud.users.stats_api.serializers import UserStatsSerializer + + +class UserStatsViewSet( + mixins.ListModelMixin, + mixins.RetrieveModelMixin, + viewsets.GenericViewSet, +): + serializer_class = UserStatsSerializer + permission_classes = [IsAdminUser] + filter_backends = [django_filters.DjangoFilterBackend] + pagination_class = CursorCountPagination + lookup_field = "user__uuid" + lookup_url_kwarg = "uuid" + + class Filter(django_filters.FilterSet): + SINCE_FIELDS = ["last_upload_at", "user__last_login"] + + active_within_days = django_filters.NumberFilter( + method="filter_active_within_days", + label="Active in last N days (upload or login)", + help_text="Return users who uploaded or logged in within the last N days.", + ) + uploaded_within_days = django_filters.NumberFilter( + method="filter_uploaded_within_days", + label="Uploaded in last N days", + help_text="Return users who uploaded a document in the last N days.", + ) + logged_in_within_days = django_filters.NumberFilter( + method="filter_logged_in_within_days", + label="Logged in within last N days", + help_text=( + "Return users whose most recent login was within the last N days." + ), + ) + + def filter_active_within_days(self, queryset, _name, value): + days = int(value) + if days < 0: + return queryset.none() + cutoff = timezone.now() - timedelta(days=days) + query = Q() + for field in self.SINCE_FIELDS: + query |= Q(**{f"{field}__gt": cutoff}) + return queryset.filter(query) + + def filter_uploaded_within_days(self, queryset, _name, value): + days = int(value) + if days < 0: + return queryset.none() + cutoff = timezone.now() - timedelta(days=days) + return queryset.filter(last_upload_at__gte=cutoff) + + def filter_logged_in_within_days(self, queryset, _name, value): + days = int(value) + if days < 0: + return queryset.none() + cutoff = timezone.now() - timedelta(days=days) + return queryset.filter(user__last_login__gte=cutoff) + + class Meta: + model = UserStats + fields = [] + + filterset_class = Filter + + def get_queryset(self): + cutoff = timezone.now() - timedelta(days=settings.UPLOAD_WINDOW_DAYS) + return UserStats.objects.select_related("user").annotate( + total_documents=Count( + "user__documents", + filter=~Q(user__documents__status=Status.deleted), + distinct=True, + ), + recent_upload_count=Count( + "user__documents", + filter=Q(user__documents__created_at__gte=cutoff) + & ~Q(user__documents__status=Status.deleted), + distinct=True, + ), + ) + + @action(detail=False, methods=["get"]) + def aged_out(self, request): + """Users with a document that crossed the window boundary since `since`, + so their upload count has dropped and needs re-syncing.""" + 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) + + win = timedelta(days=settings.UPLOAD_WINDOW_DAYS) + now = timezone.now() + qs = ( + self.get_queryset() + .filter( + user__documents__created_at__gte=since_dt - win, + user__documents__created_at__lt=now - win, + ) + .distinct() + ) + + page = self.paginate_queryset(qs) + return Response(self.get_serializer(page, many=True).data) From 8df4b6ad61669401625cd728ecef30e71fb85ced Mon Sep 17 00:00:00 2001 From: duckduckgrayduck <102841251+duckduckgrayduck@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:24:30 -0500 Subject: [PATCH 2/4] Fix test --- documentcloud/organizations/tests/test_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/documentcloud/organizations/tests/test_models.py b/documentcloud/organizations/tests/test_models.py index 9916230c..ae84d5f1 100644 --- a/documentcloud/organizations/tests/test_models.py +++ b/documentcloud/organizations/tests/test_models.py @@ -71,7 +71,7 @@ def test_merge_fks(self): if f.is_relation and f.auto_created ] ) - == 8 + == 9 ) # Many to many relations defined on the Organization model assert ( From 412c21b669402d029ce0f81cb9c782fc029a1171 Mon Sep 17 00:00:00 2001 From: duckduckgrayduck <102841251+duckduckgrayduck@users.noreply.github.com> Date: Mon, 10 Aug 2026 09:43:45 -0500 Subject: [PATCH 3/4] Move upload stats from using a signal to explicit calls in document and mailgun views. Add a Django management command to backfill missing user/org stats. Optimize the get_total methods for AI credits to make them pre-fetch friendly. Paginate the querysets. Add AI credit balances to the APIs. --- .../management/commands/backfill_stats.py | 58 ++++++++++++ documentcloud/core/utils.py | 37 ++++++++ documentcloud/core/views.py | 12 +++ documentcloud/documents/tests/test_views.py | 2 +- documentcloud/documents/views.py | 12 ++- documentcloud/organizations/apps.py | 4 + ...024_organizationstats_last_ai_credit_at.py | 18 ++++ documentcloud/organizations/models.py | 34 ++++--- documentcloud/organizations/signals.py | 24 +++++ .../organizations/stats_api/models.py | 1 + .../organizations/stats_api/serializers.py | 11 +++ .../organizations/stats_api/tests.py | 90 +++++++++++++++++++ .../organizations/stats_api/utils.py | 17 ++++ .../organizations/stats_api/views.py | 21 +++-- .../0014_userstats_last_ai_credit_at.py | 18 ++++ documentcloud/users/signals.py | 40 ++++----- documentcloud/users/stats_api/models.py | 1 + documentcloud/users/stats_api/serializers.py | 14 +++ documentcloud/users/stats_api/views.py | 68 ++++++++++---- 19 files changed, 420 insertions(+), 62 deletions(-) create mode 100644 documentcloud/core/management/commands/backfill_stats.py create mode 100644 documentcloud/organizations/migrations/0024_organizationstats_last_ai_credit_at.py create mode 100644 documentcloud/organizations/signals.py create mode 100644 documentcloud/organizations/stats_api/tests.py create mode 100644 documentcloud/organizations/stats_api/utils.py create mode 100644 documentcloud/users/migrations/0014_userstats_last_ai_credit_at.py diff --git a/documentcloud/core/management/commands/backfill_stats.py b/documentcloud/core/management/commands/backfill_stats.py new file mode 100644 index 00000000..3ba8d838 --- /dev/null +++ b/documentcloud/core/management/commands/backfill_stats.py @@ -0,0 +1,58 @@ +# Django +from django.core.management.base import BaseCommand + +# DocumentCloud +from documentcloud.organizations.models import Organization +from documentcloud.organizations.stats_api.models import OrganizationStats +from documentcloud.users.models import User +from documentcloud.users.stats_api.models import UserStats + +BATCH_SIZE = 500 + + +class Command(BaseCommand): + """Backfill stats rows for existing users and organizations. + + The post_save signals only create stats rows for users/orgs created after + they were deployed, so every pre-existing record lacks a row. This command + creates the missing rows. + + Individual organizations are skipped, matching the org stats endpoint (which + only surfaces collective orgs) and the create_organization_stats signal. + Info about AI credit balances on individual orgs are pulled on the user + record instead. + """ + + help = "Create stats rows for existing users and collective organizations" + + def handle(self, *args, **options): + # pylint: disable=unused-argument + self._backfill( + "user", + User.objects.filter(stats__isnull=True).values_list("pk", flat=True), + lambda pk: UserStats(user_id=pk), + UserStats, + ) + self._backfill( + "organization", + Organization.objects.filter( + individual=False, stats__isnull=True + ).values_list("pk", flat=True), + lambda pk: OrganizationStats(organization_id=pk), + OrganizationStats, + ) + + def _backfill(self, label, pk_iterable, build, model): + batch = [] + total = 0 + for pk in pk_iterable.iterator(chunk_size=BATCH_SIZE): + batch.append(build(pk)) + if len(batch) >= BATCH_SIZE: + model.objects.bulk_create(batch, ignore_conflicts=True) + total += len(batch) + batch = [] + self.stdout.write(f"{label}: {total:,} created...") + if batch: + model.objects.bulk_create(batch, ignore_conflicts=True) + total += len(batch) + self.stdout.write(self.style.SUCCESS(f"{label}: done, {total:,} processed")) diff --git a/documentcloud/core/utils.py b/documentcloud/core/utils.py index 55d7e605..eb9dec48 100644 --- a/documentcloud/core/utils.py +++ b/documentcloud/core/utils.py @@ -1,4 +1,5 @@ # Django +from django.utils import timezone from django.utils.text import slugify as django_slugify # Standard Library @@ -8,6 +9,10 @@ from drf_spectacular.extensions import OpenApiAuthenticationExtension from unidecode import unidecode +# DocumentCloud +from documentcloud.organizations.stats_api.models import OrganizationStats +from documentcloud.users.stats_api.models import UserStats + class ProcessingTokenAuthenticationScheme(OpenApiAuthenticationExtension): target_class = "documentcloud.core.authentication.ProcessingTokenAuthentication" @@ -55,3 +60,35 @@ def format_date(date): if date is None: return None return date.replace(tzinfo=None).isoformat() + "Z" + + +def record_uploads(user_ids=None, organization_ids=None, when=None): + """ + Bump the upload watermark for the given uploaders. + Called explicitly at document-creation sites (perform_create and the mailgun + view) rather than via a post_save signal. Updates existing stats + rows only. + """ + when = when or timezone.now() + if user_ids: + UserStats.objects.filter(user_id__in=user_ids).update(last_upload_at=when) + if organization_ids: + OrganizationStats.objects.filter(organization_id__in=organization_ids).update( + last_upload_at=when + ) + + +def record_ai_credit_use(user_id=None, organization_id=None, when=None): + """ + Bump the AI-credit-use watermark on the user and org stats rows. + Called explicitly from Organization.use_ai_credits. + Balances are read live via get_total_* calls, so this only records when + credits were last used. + """ + when = when or timezone.now() + if user_id: + UserStats.objects.filter(user_id=user_id).update(last_ai_credit_at=when) + if organization_id: + OrganizationStats.objects.filter(organization_id=organization_id).update( + last_ai_credit_at=when + ) diff --git a/documentcloud/core/views.py b/documentcloud/core/views.py index a5ab3a93..6bec3bfb 100644 --- a/documentcloud/core/views.py +++ b/documentcloud/core/views.py @@ -30,6 +30,7 @@ from documentcloud.common.environment import storage from documentcloud.common.extensions import EXTENSIONS from documentcloud.core.choices import Language +from documentcloud.core.utils import record_uploads from documentcloud.documents.choices import Access from documentcloud.documents.models import Document from documentcloud.documents.tasks import fetch_file_url @@ -104,6 +105,7 @@ def mailgun(request): attachments = json.loads(request.POST.get("attachments", "[]")) + created_any = False for attachment in attachments: with transaction.atomic(): title, original_extension = os.path.splitext(attachment["name"]) @@ -118,6 +120,7 @@ def mailgun(request): title=title, original_extension=original_extension, ) + created_any = True document.index_on_commit() transaction.on_commit( lambda a=attachment, d=document: fetch_file_url.delay( @@ -128,6 +131,15 @@ def mailgun(request): auth=("api", settings.MAILGUN_API_KEY), ) ) + + # All attachments in a message share one uploader (the mailkey user / their + # org). Bump the upload watermark once, explicitly, rather than via a signal. + # Only when at least one valid attachment was actually created. + if created_any: + record_uploads( + user_ids=[user.pk], + organization_ids=[user.organization.pk], + ) return HttpResponse("OK") diff --git a/documentcloud/documents/tests/test_views.py b/documentcloud/documents/tests/test_views.py index 93a4cdbf..a160ea26 100644 --- a/documentcloud/documents/tests/test_views.py +++ b/documentcloud/documents/tests/test_views.py @@ -307,7 +307,7 @@ def test_create_bad_ocr_engine(self, client, user): def test_bulk_create(self, client, user, django_assert_num_queries): """Create multiple documents""" client.force_authenticate(user=user) - with django_assert_num_queries(11): + with django_assert_num_queries(13): response = client.post( "/api/documents/", [{"title": "Test 1"}, {"title": "Test 2"}, {"title": "Test 3"}], diff --git a/documentcloud/documents/views.py b/documentcloud/documents/views.py index 091cb698..5378997d 100644 --- a/documentcloud/documents/views.py +++ b/documentcloud/documents/views.py @@ -41,6 +41,7 @@ ) from documentcloud.core.utils import ( # pylint:disable=unused-import ProcessingTokenAuthenticationScheme, + record_uploads, ) from documentcloud.documents.choices import Access, EntityKind, OccurrenceKind, Status from documentcloud.documents.constants import DATA_KEY_REGEX @@ -877,13 +878,18 @@ def perform_create(self, serializer): force_ocrs = [serializer.validated_data.pop("force_ocr", False)] ocr_engines = [serializer.validated_data.pop("ocr_engine", "tess4")] - documents = serializer.save( - user=self.request.user, organization=self.request.user.organization - ) + organization = self.request.user.organization + documents = serializer.save(user=self.request.user, organization=organization) if not bulk: documents = [documents] + # Update the stat records for last uploads + record_uploads( + user_ids=[self.request.user.pk], + organization_ids=[organization.pk], + ) + for document, file_url, force_ocr, ocr_engine in zip( documents, file_urls, force_ocrs, ocr_engines ): diff --git a/documentcloud/organizations/apps.py b/documentcloud/organizations/apps.py index e5a19087..89acb939 100644 --- a/documentcloud/organizations/apps.py +++ b/documentcloud/organizations/apps.py @@ -4,3 +4,7 @@ class OrganizationsConfig(AppConfig): name = "documentcloud.organizations" + + def ready(self): + # DocumentCloud + import documentcloud.organizations.signals # pylint: disable=unused-import diff --git a/documentcloud/organizations/migrations/0024_organizationstats_last_ai_credit_at.py b/documentcloud/organizations/migrations/0024_organizationstats_last_ai_credit_at.py new file mode 100644 index 00000000..79287b9f --- /dev/null +++ b/documentcloud/organizations/migrations/0024_organizationstats_last_ai_credit_at.py @@ -0,0 +1,18 @@ +# Generated by Django 5.2.15 on 2026-08-11 01:39 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("organizations", "0023_organizationstats_and_more"), + ] + + operations = [ + migrations.AddField( + model_name="organizationstats", + name="last_ai_credit_at", + field=models.DateTimeField(blank=True, db_index=True, null=True), + ), + ] diff --git a/documentcloud/organizations/models.py b/documentcloud/organizations/models.py index 24470a5e..22725ca4 100644 --- a/documentcloud/organizations/models.py +++ b/documentcloud/organizations/models.py @@ -17,6 +17,7 @@ # DocumentCloud from documentcloud.core.choices import Language from documentcloud.core.fields import AutoCreatedField, AutoLastModifiedField +from documentcloud.core.utils import record_ai_credit_use from documentcloud.organizations.exceptions import InsufficientAICreditsError from documentcloud.organizations.querysets import OrganizationQuerySet from documentcloud.organizations.stats_api.models import ( # pylint: disable=unused-import @@ -252,15 +253,29 @@ def deduct_credits(amount, organization, field): note=note, ) + record_ai_credit_use( + user_id=user_id, + organization_id=organization.pk, + ) + return ai_credit_count + # We iterate self.groups.all() and filter in Python on the following three methods + # rather than using self.groups.filter(share_resources=True) + # so this uses the prefetch cache. + # .filter() on a prefetched relation re-queries. .all() reads the cached + # groups. This keeps the org stats endpoint's credit lookups prefetch-friendly + # (no per-org N+1). Behavior is identical either way. + # https://docs.djangoproject.com/en/6.1/ref/models/querysets/#prefetch-related + def get_total_number_ai_credits(self): """Get total number AI credits including parent and groups""" number_ai_credits = self.number_ai_credits if self.parent and self.parent.share_resources: number_ai_credits += self.parent.number_ai_credits - for group in self.groups.filter(share_resources=True): - number_ai_credits += group.number_ai_credits + for group in self.groups.all(): + if group.share_resources: + number_ai_credits += group.number_ai_credits return number_ai_credits def get_total_monthly_ai_credits(self): @@ -268,8 +283,9 @@ def get_total_monthly_ai_credits(self): monthly_ai_credits = self.monthly_ai_credits if self.parent and self.parent.share_resources: monthly_ai_credits += self.parent.monthly_ai_credits - for group in self.groups.filter(share_resources=True): - monthly_ai_credits += group.monthly_ai_credits + for group in self.groups.all(): + if group.share_resources: + monthly_ai_credits += group.monthly_ai_credits return monthly_ai_credits def get_total_monthly_ai_credits_allowance(self): @@ -278,15 +294,11 @@ def get_total_monthly_ai_credits_allowance(self): This is the amount that monthly_credits will reset to each month. """ total = self.ai_credits_per_month - - # Include parent if it shares resources if self.parent and self.parent.share_resources: total += self.parent.ai_credits_per_month - - # Include groups that share resources - for group in self.groups.filter(share_resources=True): - total += group.ai_credits_per_month - + for group in self.groups.all(): + if group.share_resources: + total += group.ai_credits_per_month return total diff --git a/documentcloud/organizations/signals.py b/documentcloud/organizations/signals.py new file mode 100644 index 00000000..004e4f97 --- /dev/null +++ b/documentcloud/organizations/signals.py @@ -0,0 +1,24 @@ +# Django +from django.db.models.signals import post_save +from django.dispatch import receiver + +# DocumentCloud +from documentcloud.organizations.models import Organization +from documentcloud.organizations.stats_api.models import OrganizationStats + + +@receiver( + post_save, + sender=Organization, + dispatch_uid="documentcloud.organizations.signals.create_organization_stats", +) +def create_organization_stats(sender, instance, created, **kwargs): + """Create an OrganizationStats row when a new collective org is created. + + Same rationale as UserStats: post_save is the only hook that catches every path. + Skips individual orgs, since the org stats endpoint only surfaces collective + ones. Creates only, never updates, so no stats logic lives here. + """ + # pylint: disable=unused-argument + if created and not instance.individual: + OrganizationStats.objects.get_or_create(organization=instance) diff --git a/documentcloud/organizations/stats_api/models.py b/documentcloud/organizations/stats_api/models.py index df5510c6..751b1077 100644 --- a/documentcloud/organizations/stats_api/models.py +++ b/documentcloud/organizations/stats_api/models.py @@ -10,6 +10,7 @@ class OrganizationStats(models.Model): primary_key=True, ) 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) class Meta: app_label = "organizations" diff --git a/documentcloud/organizations/stats_api/serializers.py b/documentcloud/organizations/stats_api/serializers.py index c5372002..558576d4 100644 --- a/documentcloud/organizations/stats_api/serializers.py +++ b/documentcloud/organizations/stats_api/serializers.py @@ -14,6 +14,7 @@ class OrganizationStatsSerializer(serializers.ModelSerializer): help_text="Documents uploaded by the org within the configured window " "(UPLOAD_WINDOW_DAYS, defaults to 90)." ) + ai_credits = serializers.SerializerMethodField() class Meta: model = OrganizationStats @@ -23,6 +24,8 @@ class Meta: "last_upload_at", "days_since_last_upload", "recent_upload_count", + "ai_credits", + "last_ai_credit_at", ] read_only_fields = fields @@ -33,3 +36,11 @@ def get_days_since_last_upload(self, obj): def get_recent_upload_count(self, obj): return getattr(obj, "recent_upload_count", None) + + def get_ai_credits(self, obj): + org = obj.organization + return { + "ai_credits_per_month": org.get_total_monthly_ai_credits_allowance(), + "monthly_ai_credits": org.get_total_monthly_ai_credits(), + "number_ai_credits": org.get_total_number_ai_credits(), + } diff --git a/documentcloud/organizations/stats_api/tests.py b/documentcloud/organizations/stats_api/tests.py new file mode 100644 index 00000000..74c76ecd --- /dev/null +++ b/documentcloud/organizations/stats_api/tests.py @@ -0,0 +1,90 @@ +# Django +from django.test import TestCase + +# DocumentCloud +from documentcloud.organizations.models import ( # confirm extra required fields + Organization, +) +from documentcloud.organizations.tests.factories import OrganizationFactory +from documentcloud.users.tests.factories import UserFactory + + +class OrgTotalCreditsQueryTest(TestCase): + @classmethod + def setUpTestData(cls): + cls.parent = OrganizationFactory( + share_resources=True, + number_ai_credits=100, + monthly_ai_credits=100, + ai_credits_per_month=100, + ) + cls.org = OrganizationFactory( + parent=cls.parent, + share_resources=True, + number_ai_credits=10, + monthly_ai_credits=10, + ai_credits_per_month=10, + ) + cls.sharing_group = OrganizationFactory( + share_resources=True, + number_ai_credits=5, + monthly_ai_credits=5, + ai_credits_per_month=5, + ) + cls.non_sharing_group = OrganizationFactory( + share_resources=False, + number_ai_credits=999, + monthly_ai_credits=999, + ai_credits_per_month=999, + ) + cls.org.groups.set([cls.sharing_group, cls.non_sharing_group]) + cls.user = UserFactory() + + def test_total_credits_use_prefetch_cache_no_n_plus_one(self): + # Mirror the stats endpoint's fetch. + org = ( + Organization.objects.select_related("parent") + .prefetch_related("groups") + .get(pk=self.org.pk) + ) + # parent (select_related) + groups (prefetch) already loaded; the methods + # filter groups in Python via self.groups.all(). All three calls => 0 queries. + # Revert any method to self.groups.filter(share_resources=True) and this + # becomes 3, failing with the offending SQL. + with self.assertNumQueries(0): + org.get_total_number_ai_credits() + org.get_total_monthly_ai_credits() + org.get_total_monthly_ai_credits_allowance() + + def test_totals_are_correct(self): + org = ( + Organization.objects.select_related("parent") + .prefetch_related("groups") + .get(pk=self.org.pk) + ) + # self + sharing parent + sharing group; non_sharing_group's 999 excluded. + self.assertEqual(org.get_total_number_ai_credits(), 10 + 100 + 5) + self.assertEqual(org.get_total_monthly_ai_credits(), 10 + 100 + 5) + self.assertEqual(org.get_total_monthly_ai_credits_allowance(), 10 + 100 + 5) + + def test_python_filter_matches_orm_filter(self): + """'all()+Python == .filter(share_resources=True).""" + org = Organization.objects.prefetch_related("groups").get(pk=self.org.pk) + python_side = {g.pk for g in org.groups.all() if g.share_resources} + orm_side = set( + org.groups.filter(share_resources=True).values_list("pk", flat=True) + ) + self.assertEqual(python_side, orm_side) + + def test_totals_match_use_ai_credits_consumption(self): + org = Organization.objects.get(pk=self.org.pk) + total_monthly = org.get_total_monthly_ai_credits() + total_number = org.get_total_number_ai_credits() + + consumed = org.use_ai_credits( + amount=total_monthly + total_number, + user_id=self.user.pk, + note="test drain", + ) + self.assertEqual(consumed["monthly"], total_monthly) + self.assertEqual(consumed["regular"], total_number) diff --git a/documentcloud/organizations/stats_api/utils.py b/documentcloud/organizations/stats_api/utils.py new file mode 100644 index 00000000..14f9db57 --- /dev/null +++ b/documentcloud/organizations/stats_api/utils.py @@ -0,0 +1,17 @@ +# Django +from django.utils import timezone + + +def record_upload(user_id, organization_id, organization_individual): + # DocumentCloud + from documentcloud.organizations.stats_api.models import OrganizationStats + from documentcloud.users.stats_api.models import UserStats + + now = timezone.now() + UserStats.objects.update_or_create( + user_id=user_id, defaults={"last_upload_at": now} + ) + if not organization_individual: + OrganizationStats.objects.update_or_create( + organization_id=organization_id, defaults={"last_upload_at": now} + ) diff --git a/documentcloud/organizations/stats_api/views.py b/documentcloud/organizations/stats_api/views.py index 4bab06a5..b1d97947 100644 --- a/documentcloud/organizations/stats_api/views.py +++ b/documentcloud/organizations/stats_api/views.py @@ -15,7 +15,7 @@ from django_filters import rest_framework as django_filters # DocumentCloud -from documentcloud.core.pagination import CursorCountPagination +from documentcloud.core.pagination import CursorPagination from documentcloud.documents.choices import Status from documentcloud.organizations.stats_api.models import OrganizationStats from documentcloud.organizations.stats_api.serializers import ( @@ -31,7 +31,7 @@ class OrganizationStatsViewSet( serializer_class = OrganizationStatsSerializer permission_classes = [IsAdminUser] filter_backends = [django_filters.DjangoFilterBackend] - pagination_class = CursorCountPagination + pagination_class = CursorPagination lookup_field = "organization__uuid" lookup_url_kwarg = "uuid" @@ -58,10 +58,17 @@ class Meta: filterset_class = Filter def get_queryset(self): + return OrganizationStats.objects.select_related("organization").filter( + organization__individual=False + ) + + def paginate_queryset(self, queryset): + page = super().paginate_queryset(queryset) cutoff = timezone.now() - timedelta(days=settings.UPLOAD_WINDOW_DAYS) - return ( - OrganizationStats.objects.select_related("organization") - .filter(organization__individual=False) + annotated = ( + OrganizationStats.objects.filter(pk__in=[o.pk for o in page]) + .select_related("organization", "organization__parent") + .prefetch_related("organization__groups") .annotate( total_documents=Count( "organization__documents", @@ -75,7 +82,9 @@ def get_queryset(self): distinct=True, ), ) + .order_by("pk") ) + return list(annotated) @action(detail=False, methods=["get"]) def aged_out(self, request): @@ -100,4 +109,4 @@ def aged_out(self, request): ) page = self.paginate_queryset(qs) - return Response(self.get_serializer(page, many=True).data) + return self.get_paginated_response(self.get_serializer(page, many=True).data) diff --git a/documentcloud/users/migrations/0014_userstats_last_ai_credit_at.py b/documentcloud/users/migrations/0014_userstats_last_ai_credit_at.py new file mode 100644 index 00000000..46266e4d --- /dev/null +++ b/documentcloud/users/migrations/0014_userstats_last_ai_credit_at.py @@ -0,0 +1,18 @@ +# Generated by Django 5.2.15 on 2026-08-11 01:39 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("users", "0013_userstats_alter_user_document_language_and_more"), + ] + + operations = [ + migrations.AddField( + model_name="userstats", + name="last_ai_credit_at", + field=models.DateTimeField(blank=True, db_index=True, null=True), + ), + ] diff --git a/documentcloud/users/signals.py b/documentcloud/users/signals.py index 0a8d2b61..63e95990 100644 --- a/documentcloud/users/signals.py +++ b/documentcloud/users/signals.py @@ -2,10 +2,11 @@ from django.contrib.auth.signals import user_logged_in from django.db.models.signals import post_save from django.dispatch import receiver -from django.utils import timezone # DocumentCloud from documentcloud.addons.models import AddOn +from documentcloud.users.models import User +from documentcloud.users.stats_api.models import UserStats @receiver(user_logged_in, dispatch_uid="documentcloud.user.signals.default_addons") @@ -18,25 +19,18 @@ def default_addons(sender, user, request, **kwargs): user.active_addons.set(AddOn.objects.filter(default=True)) -@receiver(post_save, dispatch_uid="documentcloud.users.signals.touch_last_upload") -def touch_last_upload(sender, instance, created, **kwargs): - """Record upload recency when a user creates a new document""" - # DocumentCloud - from documentcloud.documents.models import Document - from documentcloud.organizations.stats_api.models import OrganizationStats - from documentcloud.users.stats_api.models import UserStats - - if sender is not Document or not created: - return - - now = timezone.now() - UserStats.objects.update_or_create( - user_id=instance.user_id, - defaults={"last_upload_at": now}, - ) - - if not instance.organization.individual: - OrganizationStats.objects.update_or_create( - organization_id=instance.organization_id, - defaults={"last_upload_at": now}, - ) +@receiver( + post_save, sender=User, dispatch_uid="documentcloud.users.signals.create_user_stats" +) +def create_user_stats(sender, instance, created, **kwargs): + """ + Create a UserStats row when a new user is created. + Users are created by squarelet sync, + so post_save is the only path that catches every creation. + get_or_create avoids a duplicate-row error if two requests create the same + user's stats at once. The receiver only ever creates, never updates, + so no stats logic lives in the signal. + """ + # pylint: disable=unused-argument + if created: + UserStats.objects.get_or_create(user=instance) diff --git a/documentcloud/users/stats_api/models.py b/documentcloud/users/stats_api/models.py index 5f7cb50d..ecd384be 100644 --- a/documentcloud/users/stats_api/models.py +++ b/documentcloud/users/stats_api/models.py @@ -10,6 +10,7 @@ class UserStats(models.Model): primary_key=True, ) 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) class Meta: app_label = "users" diff --git a/documentcloud/users/stats_api/serializers.py b/documentcloud/users/stats_api/serializers.py index 6cdc6726..a7452f93 100644 --- a/documentcloud/users/stats_api/serializers.py +++ b/documentcloud/users/stats_api/serializers.py @@ -17,6 +17,7 @@ class UserStatsSerializer(serializers.ModelSerializer): help_text="Number of documents uploaded within the configured recent window " "(UPLOAD_WINDOW_DAYS, currently defaults to 90)." ) + individual_ai_credits = serializers.SerializerMethodField() class Meta: model = UserStats @@ -27,6 +28,8 @@ class Meta: "days_since_last_upload", "last_login_at", "recent_upload_count", + "individual_ai_credits", + "last_ai_credit_at", ] read_only_fields = fields @@ -37,3 +40,14 @@ def get_days_since_last_upload(self, obj): 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: + return None + org = orgs[0] + return { + "ai_credits_per_month": org.ai_credits_per_month, + "monthly_ai_credits": org.monthly_ai_credits, + "number_ai_credits": org.number_ai_credits, + } diff --git a/documentcloud/users/stats_api/views.py b/documentcloud/users/stats_api/views.py index f58c6b3b..d6c847d3 100644 --- a/documentcloud/users/stats_api/views.py +++ b/documentcloud/users/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 Count, Prefetch, Q from django.utils import timezone from django.utils.dateparse import parse_datetime from rest_framework import mixins, viewsets @@ -15,8 +15,9 @@ from django_filters import rest_framework as django_filters # DocumentCloud -from documentcloud.core.pagination import CursorCountPagination +from documentcloud.core.pagination import CursorPagination from documentcloud.documents.choices import Status +from documentcloud.organizations.models import Organization from documentcloud.users.stats_api.models import UserStats from documentcloud.users.stats_api.serializers import UserStatsSerializer @@ -29,7 +30,7 @@ class UserStatsViewSet( serializer_class = UserStatsSerializer permission_classes = [IsAdminUser] filter_backends = [django_filters.DjangoFilterBackend] - pagination_class = CursorCountPagination + pagination_class = CursorPagination lookup_field = "user__uuid" lookup_url_kwarg = "uuid" @@ -39,7 +40,9 @@ class Filter(django_filters.FilterSet): active_within_days = django_filters.NumberFilter( method="filter_active_within_days", label="Active in last N days (upload or login)", - help_text="Return users who uploaded or logged in within the last N days.", + help_text=( + "Return users who uploaded " "or logged in within the last N days." + ), ) uploaded_within_days = django_filters.NumberFilter( method="filter_uploaded_within_days", @@ -53,6 +56,11 @@ class Filter(django_filters.FilterSet): "Return users whose most recent login was within the last N days." ), ) + used_ai_credits_within_days = django_filters.NumberFilter( + method="filter_used_ai_credits_within_days", + label="Used AI credits within last N days", + help_text="Return users who used AI credits within the last N days.", + ) def filter_active_within_days(self, queryset, _name, value): days = int(value) @@ -61,7 +69,7 @@ def filter_active_within_days(self, queryset, _name, value): cutoff = timezone.now() - timedelta(days=days) query = Q() for field in self.SINCE_FIELDS: - query |= Q(**{f"{field}__gt": cutoff}) + query |= Q(**{f"{field}__gte": cutoff}) return queryset.filter(query) def filter_uploaded_within_days(self, queryset, _name, value): @@ -78,6 +86,13 @@ def filter_logged_in_within_days(self, queryset, _name, value): cutoff = timezone.now() - timedelta(days=days) return queryset.filter(user__last_login__gte=cutoff) + def filter_used_ai_credits_within_days(self, queryset, _name, value): + days = int(value) + if days < 0: + return queryset.none() + cutoff = timezone.now() - timedelta(days=days) + return queryset.filter(last_ai_credit_at__gte=cutoff) + class Meta: model = UserStats fields = [] @@ -85,20 +100,37 @@ class Meta: filterset_class = Filter def get_queryset(self): + return UserStats.objects.select_related("user") + + def paginate_queryset(self, queryset): + page = super().paginate_queryset(queryset) cutoff = timezone.now() - timedelta(days=settings.UPLOAD_WINDOW_DAYS) - return UserStats.objects.select_related("user").annotate( - total_documents=Count( - "user__documents", - filter=~Q(user__documents__status=Status.deleted), - distinct=True, - ), - recent_upload_count=Count( - "user__documents", - filter=Q(user__documents__created_at__gte=cutoff) - & ~Q(user__documents__status=Status.deleted), - distinct=True, - ), + annotated = ( + UserStats.objects.filter(pk__in=[u.pk for u in page]) + .select_related("user") + .prefetch_related( + Prefetch( + "user__organizations", + queryset=Organization.objects.filter(individual=True), + to_attr="individual_orgs", + ) + ) + .annotate( + total_documents=Count( + "user__documents", + filter=~Q(user__documents__status=Status.deleted), + distinct=True, + ), + recent_upload_count=Count( + "user__documents", + filter=Q(user__documents__created_at__gte=cutoff) + & ~Q(user__documents__status=Status.deleted), + distinct=True, + ), + ) + .order_by("pk") ) + return list(annotated) @action(detail=False, methods=["get"]) def aged_out(self, request): @@ -123,4 +155,4 @@ def aged_out(self, request): ) page = self.paginate_queryset(qs) - return Response(self.get_serializer(page, many=True).data) + return self.get_paginated_response(self.get_serializer(page, many=True).data) From 7918eaf4fddb6adb6c2042550da7651b0d5fdede Mon Sep 17 00:00:00 2001 From: duckduckgrayduck <102841251+duckduckgrayduck@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:12:10 -0500 Subject: [PATCH 4/4] Enrich detail view, remove leftover util add user/org stats to INSTALLED_APPS and add tests --- config/settings/base.py | 2 + documentcloud/conftest.py | 11 ++ documentcloud/core/utils.py | 10 +- documentcloud/core/views.py | 4 +- documentcloud/documents/views.py | 4 +- documentcloud/organizations/models.py | 3 - documentcloud/organizations/stats_api/apps.py | 7 + .../organizations/stats_api/tests.py | 167 ++++++++++++++++-- .../organizations/stats_api/utils.py | 17 -- .../organizations/stats_api/views.py | 21 ++- documentcloud/users/models.py | 3 - documentcloud/users/stats_api/apps.py | 7 + documentcloud/users/stats_api/tests.py | 164 +++++++++++++++++ documentcloud/users/stats_api/views.py | 17 +- 14 files changed, 378 insertions(+), 59 deletions(-) create mode 100644 documentcloud/organizations/stats_api/apps.py delete mode 100644 documentcloud/organizations/stats_api/utils.py create mode 100644 documentcloud/users/stats_api/apps.py diff --git a/config/settings/base.py b/config/settings/base.py index 469e2953..08e0cb3f 100644 --- a/config/settings/base.py +++ b/config/settings/base.py @@ -95,6 +95,8 @@ "documentcloud.statistics.apps.StatisticsConfig", "documentcloud.users.apps.UsersConfig", "documentcloud.entities.apps.EntitiesConfig", + "documentcloud.organizations.stats_api", + "documentcloud.users.stats_api", ] # https://docs.djangoproject.com/en/dev/ref/settings/#installed-apps INSTALLED_APPS = DJANGO_APPS + THIRD_PARTY_APPS + LOCAL_APPS diff --git a/documentcloud/conftest.py b/documentcloud/conftest.py index 9a81b0a0..550ac2ea 100644 --- a/documentcloud/conftest.py +++ b/documentcloud/conftest.py @@ -90,3 +90,14 @@ def entity(): @pytest.fixture def entity_occurrence(): return EntityOccurrenceFactory() + + +@pytest.fixture +def user_with_collective_org(): + member = UserFactory() + org = OrganizationFactory(individual=False, members=[member]) + # UserFactory gives the user an individual org as their active org; make the + # collective org active so user.organization (used by perform_create) returns it + member.memberships.filter(organization=org).update(active=True) + member.memberships.exclude(organization=org).update(active=False) + return member, org diff --git a/documentcloud/core/utils.py b/documentcloud/core/utils.py index eb9dec48..cea83e73 100644 --- a/documentcloud/core/utils.py +++ b/documentcloud/core/utils.py @@ -62,7 +62,7 @@ def format_date(date): return date.replace(tzinfo=None).isoformat() + "Z" -def record_uploads(user_ids=None, organization_ids=None, when=None): +def record_uploads(user_id=None, organization_id=None, when=None): """ Bump the upload watermark for the given uploaders. Called explicitly at document-creation sites (perform_create and the mailgun @@ -70,10 +70,10 @@ def record_uploads(user_ids=None, organization_ids=None, when=None): rows only. """ when = when or timezone.now() - if user_ids: - UserStats.objects.filter(user_id__in=user_ids).update(last_upload_at=when) - if organization_ids: - OrganizationStats.objects.filter(organization_id__in=organization_ids).update( + if user_id: + UserStats.objects.filter(user_id=user_id).update(last_upload_at=when) + if organization_id: + OrganizationStats.objects.filter(organization_id=organization_id).update( last_upload_at=when ) diff --git a/documentcloud/core/views.py b/documentcloud/core/views.py index 6bec3bfb..08d2d299 100644 --- a/documentcloud/core/views.py +++ b/documentcloud/core/views.py @@ -137,8 +137,8 @@ def mailgun(request): # Only when at least one valid attachment was actually created. if created_any: record_uploads( - user_ids=[user.pk], - organization_ids=[user.organization.pk], + user_id=user.pk, + organization_id=user.organization.pk, ) return HttpResponse("OK") diff --git a/documentcloud/documents/views.py b/documentcloud/documents/views.py index 5378997d..007d1d5b 100644 --- a/documentcloud/documents/views.py +++ b/documentcloud/documents/views.py @@ -886,8 +886,8 @@ def perform_create(self, serializer): # Update the stat records for last uploads record_uploads( - user_ids=[self.request.user.pk], - organization_ids=[organization.pk], + user_id=self.request.user.pk, + organization_id=organization.pk, ) for document, file_url, force_ocr, ocr_engine in zip( diff --git a/documentcloud/organizations/models.py b/documentcloud/organizations/models.py index 22725ca4..232a9932 100644 --- a/documentcloud/organizations/models.py +++ b/documentcloud/organizations/models.py @@ -20,9 +20,6 @@ from documentcloud.core.utils import record_ai_credit_use from documentcloud.organizations.exceptions import InsufficientAICreditsError from documentcloud.organizations.querysets import OrganizationQuerySet -from documentcloud.organizations.stats_api.models import ( # pylint: disable=unused-import - OrganizationStats, -) logger = logging.getLogger(__name__) diff --git a/documentcloud/organizations/stats_api/apps.py b/documentcloud/organizations/stats_api/apps.py new file mode 100644 index 00000000..24d969ca --- /dev/null +++ b/documentcloud/organizations/stats_api/apps.py @@ -0,0 +1,7 @@ +# Django +from django.apps import AppConfig + + +class OrganizationStatsApiConfig(AppConfig): + name = "documentcloud.organizations.stats_api" + label = "organizations_stats_api" diff --git a/documentcloud/organizations/stats_api/tests.py b/documentcloud/organizations/stats_api/tests.py index 74c76ecd..e5c096ad 100644 --- a/documentcloud/organizations/stats_api/tests.py +++ b/documentcloud/organizations/stats_api/tests.py @@ -1,10 +1,23 @@ # Django +from django.conf import settings from django.test import TestCase +from django.utils import timezone +from rest_framework import status + +# Standard Library +import json +from datetime import timedelta + +# Third Party +import pytest # DocumentCloud +from documentcloud.documents.models import Document +from documentcloud.documents.tests.factories import DocumentFactory from documentcloud.organizations.models import ( # confirm extra required fields Organization, ) +from documentcloud.organizations.stats_api.models import OrganizationStats from documentcloud.organizations.tests.factories import OrganizationFactory from documentcloud.users.tests.factories import UserFactory @@ -40,17 +53,15 @@ def setUpTestData(cls): cls.org.groups.set([cls.sharing_group, cls.non_sharing_group]) cls.user = UserFactory() - def test_total_credits_use_prefetch_cache_no_n_plus_one(self): + def test_total_credits_use_prefetch_cache(self): # Mirror the stats endpoint's fetch. org = ( Organization.objects.select_related("parent") .prefetch_related("groups") .get(pk=self.org.pk) ) - # parent (select_related) + groups (prefetch) already loaded; the methods - # filter groups in Python via self.groups.all(). All three calls => 0 queries. - # Revert any method to self.groups.filter(share_resources=True) and this - # becomes 3, failing with the offending SQL. + # Regression test against N+1 behavior in get_total methods. + # parent (select_related) + groups (prefetch) already loaded. with self.assertNumQueries(0): org.get_total_number_ai_credits() org.get_total_monthly_ai_credits() @@ -62,25 +73,16 @@ def test_totals_are_correct(self): .prefetch_related("groups") .get(pk=self.org.pk) ) - # self + sharing parent + sharing group; non_sharing_group's 999 excluded. + # self + sharing parent + sharing group + # non_sharing_group's 999 excluded. self.assertEqual(org.get_total_number_ai_credits(), 10 + 100 + 5) self.assertEqual(org.get_total_monthly_ai_credits(), 10 + 100 + 5) self.assertEqual(org.get_total_monthly_ai_credits_allowance(), 10 + 100 + 5) - def test_python_filter_matches_orm_filter(self): - """'all()+Python == .filter(share_resources=True).""" - org = Organization.objects.prefetch_related("groups").get(pk=self.org.pk) - python_side = {g.pk for g in org.groups.all() if g.share_resources} - orm_side = set( - org.groups.filter(share_resources=True).values_list("pk", flat=True) - ) - self.assertEqual(python_side, orm_side) - def test_totals_match_use_ai_credits_consumption(self): org = Organization.objects.get(pk=self.org.pk) total_monthly = org.get_total_monthly_ai_credits() total_number = org.get_total_number_ai_credits() - consumed = org.use_ai_credits( amount=total_monthly + total_number, user_id=self.user.pk, @@ -88,3 +90,136 @@ def test_totals_match_use_ai_credits_consumption(self): ) self.assertEqual(consumed["monthly"], total_monthly) self.assertEqual(consumed["regular"], total_number) + + +@pytest.mark.django_db() +class TestOrganizationStatsAPI: + def _admin(self): + return UserFactory(is_staff=True) + + def test_list_requires_admin(self, client): + client.force_authenticate(user=UserFactory()) + response = client.get("/stats_api/organizations/") + assert response.status_code == status.HTTP_403_FORBIDDEN + + def test_list_unauthenticated(self, client): + response = client.get("/stats_api/organizations/") + 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""" + client.force_authenticate(user=self._admin()) + collective = OrganizationFactory.create_batch(3, individual=False) + individual = OrganizationFactory(individual=True) + response = client.get("/stats_api/organizations/") + assert response.status_code == status.HTTP_200_OK + uuids = {r["uuid"] for r in json.loads(response.content)["results"]} + for org in collective: + assert str(org.uuid) in uuids + # get_queryset filters individual=False + 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).""" + client.force_authenticate(user=self._admin()) + org = OrganizationFactory(individual=False) + DocumentFactory.create_batch(2, organization=org) + + response = client.get(f"/stats_api/organizations/{org.uuid}/") + assert response.status_code == status.HTTP_200_OK + body = json.loads(response.content) + assert body["total_documents"] == 2 + assert body["recent_upload_count"] == 2 + + def test_filter_uploaded_within_days(self, client): + client.force_authenticate(user=self._admin()) + now = timezone.now() + + recent = OrganizationFactory(individual=False) + OrganizationStats.objects.filter(organization=recent).update( + last_upload_at=now - timedelta(days=1) + ) + old = OrganizationFactory(individual=False) + OrganizationStats.objects.filter(organization=old).update( + last_upload_at=now - timedelta(days=30) + ) + + response = client.get("/stats_api/organizations/", {"uploaded_within_days": 7}) + assert response.status_code == status.HTTP_200_OK + uuids = {r["uuid"] for r in json.loads(response.content)["results"]} + assert str(recent.uuid) in uuids + assert str(old.uuid) not in uuids + + def test_document_upload_bumps_org_watermark( + self, client, user_with_collective_org + ): + 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? + + stats = OrganizationStats.objects.get(organization=org) + assert stats.last_upload_at is not None + + 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 + org.save() + assert OrganizationStats.objects.get(organization=org).last_ai_credit_at is None + + response = client.post( + f"/api/organizations/{org.pk}/ai_credits/", + {"ai_credits": 1, "user_id": user.pk}, + HTTP_AUTHORIZATION=f"processing-token {settings.PROCESSING_TOKEN}", + ) + assert response.status_code == status.HTTP_200_OK + + stats = OrganizationStats.objects.get(organization=org) + assert stats.last_ai_credit_at is not None + + def test_aged_out_requires_since(self, client): + client.force_authenticate(user=self._admin()) + response = client.get("/stats_api/organizations/aged_out/") + assert response.status_code == status.HTTP_400_BAD_REQUEST + + def test_aged_out_invalid_since(self, client): + client.force_authenticate(user=self._admin()) + response = client.get( + "/stats_api/organizations/aged_out/", {"since": "not-a-date"} + ) + assert response.status_code == status.HTTP_400_BAD_REQUEST + + def test_aged_out_returns_boundary_crossers(self, client): + """Test that aged_out catches boundary crossers""" + client.force_authenticate(user=self._admin()) + now = timezone.now() + win = timedelta(days=settings.UPLOAD_WINDOW_DAYS) + + # org with a doc that aged out of the window since `since` + aged = OrganizationFactory(individual=False) + doc = DocumentFactory(organization=aged) + # created_at is auto_now_add + Document.objects.filter(pk=doc.pk).update( + created_at=now - win - timedelta(days=1) + ) + + # org with only a fresh doc — still in window, should NOT appear + fresh = OrganizationFactory(individual=False) + DocumentFactory(organization=fresh) # created_at = now + + since = (now - timedelta(days=2)).isoformat() + response = client.get("/stats_api/organizations/aged_out/", {"since": since}) + assert response.status_code == status.HTTP_200_OK + uuids = {r["uuid"] for r in json.loads(response.content)["results"]} + assert str(aged.uuid) in uuids + assert str(fresh.uuid) not in uuids diff --git a/documentcloud/organizations/stats_api/utils.py b/documentcloud/organizations/stats_api/utils.py deleted file mode 100644 index 14f9db57..00000000 --- a/documentcloud/organizations/stats_api/utils.py +++ /dev/null @@ -1,17 +0,0 @@ -# Django -from django.utils import timezone - - -def record_upload(user_id, organization_id, organization_individual): - # DocumentCloud - from documentcloud.organizations.stats_api.models import OrganizationStats - from documentcloud.users.stats_api.models import UserStats - - now = timezone.now() - UserStats.objects.update_or_create( - user_id=user_id, defaults={"last_upload_at": now} - ) - if not organization_individual: - OrganizationStats.objects.update_or_create( - organization_id=organization_id, defaults={"last_upload_at": now} - ) diff --git a/documentcloud/organizations/stats_api/views.py b/documentcloud/organizations/stats_api/views.py index b1d97947..d8d6ed2d 100644 --- a/documentcloud/organizations/stats_api/views.py +++ b/documentcloud/organizations/stats_api/views.py @@ -62,12 +62,10 @@ def get_queryset(self): organization__individual=False ) - def paginate_queryset(self, queryset): - page = super().paginate_queryset(queryset) + def _annotate_and_prefetch(self, queryset): cutoff = timezone.now() - timedelta(days=settings.UPLOAD_WINDOW_DAYS) - annotated = ( - OrganizationStats.objects.filter(pk__in=[o.pk for o in page]) - .select_related("organization", "organization__parent") + return ( + queryset.select_related("organization", "organization__parent") .prefetch_related("organization__groups") .annotate( total_documents=Count( @@ -82,10 +80,21 @@ def paginate_queryset(self, queryset): distinct=True, ), ) - .order_by("pk") ) + + 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`, diff --git a/documentcloud/users/models.py b/documentcloud/users/models.py index 1e0d6d51..f455f06e 100644 --- a/documentcloud/users/models.py +++ b/documentcloud/users/models.py @@ -11,9 +11,6 @@ # DocumentCloud from documentcloud.core.choices import Language from documentcloud.users.managers import UserManager -from documentcloud.users.stats_api.models import ( # pylint: disable=unused-import - UserStats, -) class User(SAUser): diff --git a/documentcloud/users/stats_api/apps.py b/documentcloud/users/stats_api/apps.py new file mode 100644 index 00000000..8f69f79d --- /dev/null +++ b/documentcloud/users/stats_api/apps.py @@ -0,0 +1,7 @@ +# Django +from django.apps import AppConfig + + +class UsersStatsApiConfig(AppConfig): + name = "documentcloud.users.stats_api" + label = "users_stats_api" diff --git a/documentcloud/users/stats_api/tests.py b/documentcloud/users/stats_api/tests.py index e69de29b..2a93fbbb 100644 --- a/documentcloud/users/stats_api/tests.py +++ b/documentcloud/users/stats_api/tests.py @@ -0,0 +1,164 @@ +# Django +from django.utils import timezone +from rest_framework import status + +# Standard Library +import json +from datetime import timedelta + +# Third Party +import pytest + +# DocumentCloud +from documentcloud.documents.tests.factories import DocumentFactory +from documentcloud.users.stats_api.models import UserStats +from documentcloud.users.tests.factories import UserFactory + + +@pytest.mark.django_db() +class TestUserStatsAPI: + def _admin(self): + return UserFactory(is_staff=True) + + def test_list(self, client): + admin = self._admin() + client.force_authenticate(user=admin) + UserFactory.create_batch(3) + response = client.get("/stats_api/users/") + assert response.status_code == status.HTTP_200_OK + body = json.loads(response.content) + # admin + 3 created + any individual-org users; assert at least the ones we made + assert len(body["results"]) >= 4 + + def test_list_requires_admin(self, client): + client.force_authenticate(user=UserFactory()) # non-staff + response = client.get("/stats_api/users/") + assert response.status_code == status.HTTP_403_FORBIDDEN + + def test_list_unauthenticated(self, client): + response = client.get("/stats_api/users/") + 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).""" + admin = self._admin() + client.force_authenticate(user=admin) + + target = UserFactory() + # give the individual org credits so individual_ai_credits is non-zero + org = target.organization + org.monthly_ai_credits = 5 + org.ai_credits_per_month = 10 + org.save() + # give them documents so the counts are non-zero + DocumentFactory.create_batch(2, user=target, organization=org) + + response = client.get(f"/stats_api/users/{target.uuid}/") + assert response.status_code == status.HTTP_200_OK + body = json.loads(response.content) + + assert body["total_documents"] == 2 + assert body["recent_upload_count"] == 2 + # the field the review flagged — must be present and reflect the org + assert body["individual_ai_credits"]["monthly_ai_credits"] == 5 + assert body["individual_ai_credits"]["ai_credits_per_month"] == 10 + + def test_filter_uploaded_within_days(self, client): + admin = self._admin() + client.force_authenticate(user=admin) + now = timezone.now() + + recent = UserFactory() + UserStats.objects.filter(user=recent).update( + last_upload_at=now - timedelta(days=1) + ) + old = UserFactory() + UserStats.objects.filter(user=old).update( + last_upload_at=now - timedelta(days=30) + ) + + response = client.get("/stats_api/users/", {"uploaded_within_days": 7}) + assert response.status_code == status.HTTP_200_OK + uuids = {r["uuid"] for r in json.loads(response.content)["results"]} + assert str(recent.uuid) in uuids + assert str(old.uuid) not in uuids + + def test_filter_logged_in_within_days(self, client): + admin = self._admin() + client.force_authenticate(user=admin) + now = timezone.now() + + recent = UserFactory(last_login=now - timedelta(days=1)) + old = UserFactory(last_login=now - timedelta(days=30)) + + response = client.get("/stats_api/users/", {"logged_in_within_days": 7}) + uuids = {r["uuid"] for r in json.loads(response.content)["results"]} + assert str(recent.uuid) in uuids + assert str(old.uuid) not in uuids + + def test_filter_used_ai_credits_within_days(self, client): + admin = self._admin() + client.force_authenticate(user=admin) + now = timezone.now() + + recent = UserFactory() + UserStats.objects.filter(user=recent).update( + last_ai_credit_at=now - timedelta(days=1) + ) + old = UserFactory() + UserStats.objects.filter(user=old).update( + last_ai_credit_at=now - timedelta(days=30) + ) + + response = client.get("/stats_api/users/", {"used_ai_credits_within_days": 7}) + uuids = {r["uuid"] for r in json.loads(response.content)["results"]} + assert str(recent.uuid) in uuids + assert str(old.uuid) not in uuids + + def test_active_within_days_excludes_ai_credit_only(self, client): + """active = upload OR login, NOT ai-credit use.""" + admin = self._admin() + client.force_authenticate(user=admin) + now = timezone.now() + + uploader = UserFactory() + UserStats.objects.filter(user=uploader).update( + last_upload_at=now - timedelta(days=1) + ) + ai_only = UserFactory(last_login=now - timedelta(days=90)) + UserStats.objects.filter(user=ai_only).update( + last_ai_credit_at=now - timedelta(days=1) + ) + + response = client.get("/stats_api/users/", {"active_within_days": 7}) + uuids = {r["uuid"] for r in json.loads(response.content)["results"]} + assert str(uploader.uuid) in uuids + # AI-credit use alone is NOT activity + assert str(ai_only.uuid) not in uuids + + def test_record_uploads_sets_watermark(self): + # DocumentCloud + from documentcloud.core.utils import record_uploads + + user = UserFactory() + stats = UserStats.objects.get(user=user) + assert stats.last_upload_at is None + + record_uploads(user_id=user.pk) + + stats.refresh_from_db() + assert stats.last_upload_at is not None + + def test_record_ai_credit_use_sets_watermark(self): + # DocumentCloud + from documentcloud.core.utils import record_ai_credit_use + + user = UserFactory() + stats = UserStats.objects.get(user=user) + assert stats.last_ai_credit_at is None + + record_ai_credit_use(user_id=user.pk) + stats.refresh_from_db() + assert stats.last_ai_credit_at is not None diff --git a/documentcloud/users/stats_api/views.py b/documentcloud/users/stats_api/views.py index d6c847d3..de681510 100644 --- a/documentcloud/users/stats_api/views.py +++ b/documentcloud/users/stats_api/views.py @@ -104,10 +104,19 @@ def get_queryset(self): def paginate_queryset(self, queryset): page = super().paginate_queryset(queryset) - cutoff = timezone.now() - timedelta(days=settings.UPLOAD_WINDOW_DAYS) - annotated = ( + annotated = self._annotate_and_prefetch( UserStats.objects.filter(pk__in=[u.pk for u in page]) - .select_related("user") + ).order_by("pk") + return list(annotated) + + def get_object(self): + obj = super().get_object() + return self._annotate_and_prefetch(UserStats.objects.filter(pk=obj.pk)).get() + + def _annotate_and_prefetch(self, queryset): + cutoff = timezone.now() - timedelta(days=settings.UPLOAD_WINDOW_DAYS) + return ( + queryset.select_related("user") .prefetch_related( Prefetch( "user__organizations", @@ -128,9 +137,7 @@ def paginate_queryset(self, queryset): distinct=True, ), ) - .order_by("pk") ) - return list(annotated) @action(detail=False, methods=["get"]) def aged_out(self, request):