First pass at stats_api

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
@@ -694,3 +696,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)

config/urls.py

@@ -43,13 +43,15 @@
 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,
     ProjectMembershipViewSet,
     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:

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

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"))

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_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
+    view) rather than via a post_save signal. Updates existing stats
+    rows only.
+    """
+    when = when or timezone.now()
+    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
+        )
+
+
+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
+        )

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_id=user.pk,
+            organization_id=user.organization.pk,
+        )
     return HttpResponse("OK")
 
 

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"}],

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_id=self.request.user.pk,
+            organization_id=organization.pk,
+        )
+
         for document, file_url, force_ocr, ocr_engine in zip(
             documents, file_urls, force_ocrs, ocr_engines
         ):

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