From 02dfba160cedd8a37aecf2affe4f3c7c54fd3c7a Mon Sep 17 00:00:00 2001
From: Jacob Nesbitt <jjnesbitt2@gmail.com>
Date: Mon, 30 Dec 2024 12:41:17 -0500
Subject: [PATCH] Fix N+1 query problem for dandiset star info

To fix this behavior for the dandiset list endpoint, create a new method
on the DandisetViewSet class, `_get_dandiset_star_context`, which
collects dandiset star information. This information is then passed to
the DandisetListSerializer as context.
---
 dandiapi/api/views/dandiset.py    | 54 +++++++++++++++++++++++++++++--
 dandiapi/api/views/serializers.py |  6 ++++
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/dandiapi/api/views/dandiset.py b/dandiapi/api/views/dandiset.py
index 1cf6e0650..f0a6af3aa 100644
--- a/dandiapi/api/views/dandiset.py
+++ b/dandiapi/api/views/dandiset.py
@@ -26,6 +26,7 @@
 from dandiapi.api.asset_paths import get_root_paths_many
 from dandiapi.api.mail import send_ownership_change_emails
 from dandiapi.api.models import Dandiset, Version
+from dandiapi.api.models.dandiset import DandisetStar
 from dandiapi.api.services import audit
 from dandiapi.api.services.dandiset import (
     create_dandiset,
@@ -196,6 +197,36 @@ def get_object(self):
 
         return dandiset
 
+    def _get_dandiset_star_context(self, dandisets):
+        # Default value for all relevant dandisets
+        dandisets_to_stars = {
+            d.id: {'total': 0, 'starred_by_current_user': False} for d in dandisets
+        }
+
+        # Group the stars for these dandisets by the dandiset ID,
+        # yielding pairs of (Dandiset ID, Star Count)
+        dandiset_stars = (
+            DandisetStar.objects.filter(dandiset__in=dandisets)
+            .values_list('dandiset')
+            .annotate(star_count=Count('id'))
+            .order_by()
+        )
+        for dandiset_id, star_count in dandiset_stars:
+            dandisets_to_stars[dandiset_id]['total'] = star_count
+
+        # Only annotate dandisets as starred by current user if user is logged in
+        user = self.request.user
+        if user.is_anonymous:
+            return dandisets_to_stars
+        user = typing.cast(User, user)
+
+        # Filter previous query to current user stars
+        user_starred_dandisets = dandiset_stars.filter(user=user)
+        for dandiset_id, _ in user_starred_dandisets:
+            dandisets_to_stars[dandiset_id]['starred_by_current_user'] = True
+
+        return dandisets_to_stars
+
     @staticmethod
     def _get_dandiset_to_version_map(dandisets):
         """Map Dandiset IDs to that dandiset's draft and most recently published version."""
@@ -266,6 +297,7 @@ def search(self, request, *args, **kwargs):
         qs = self.get_queryset()
         dandisets = self.filter_queryset(qs).filter(id__in=relevant_assets.values('dandiset_id'))
         dandisets = self.paginate_queryset(dandisets)
+        dandiset_stars = self._get_dandiset_star_context(dandisets)
         dandisets_to_versions = self._get_dandiset_to_version_map(dandisets)
         dandisets_to_asset_counts = (
             AssetSearch.objects.values('dandiset_id')
@@ -287,7 +319,11 @@ def search(self, request, *args, **kwargs):
         serializer = DandisetSearchResultListSerializer(
             dandisets,
             many=True,
-            context={'dandisets': dandisets_to_versions, 'asset_counts': dandisets_to_asset_counts},
+            context={
+                'dandisets': dandisets_to_versions,
+                'asset_counts': dandisets_to_asset_counts,
+                'stars': dandiset_stars,
+            },
         )
         return self.get_paginated_response(serializer.data)
 
@@ -299,8 +335,14 @@ def list(self, request, *args, **kwargs):
         qs = self.get_queryset()
         dandisets = self.paginate_queryset(self.filter_queryset(qs))
         dandisets_to_versions = self._get_dandiset_to_version_map(dandisets)
+        dandiset_stars = self._get_dandiset_star_context(dandisets)
         serializer = DandisetListSerializer(
-            dandisets, many=True, context={'dandisets': dandisets_to_versions}
+            dandisets,
+            many=True,
+            context={
+                'dandisets': dandisets_to_versions,
+                'stars': dandiset_stars,
+            },
         )
         return self.get_paginated_response(serializer.data)
 
@@ -565,7 +607,13 @@ def starred(self, request) -> Response:
         dandisets = Dandiset.objects.filter(stars__user=request.user).order_by('-stars__created')
         dandisets = self.paginate_queryset(dandisets)
         dandisets_to_versions = self._get_dandiset_to_version_map(dandisets)
+        dandiset_stars = self._get_dandiset_star_context(dandisets)
         serializer = DandisetListSerializer(
-            dandisets, many=True, context={'dandisets': dandisets_to_versions}
+            dandisets,
+            many=True,
+            context={
+                'dandisets': dandisets_to_versions,
+                'stars': dandiset_stars,
+            },
         )
         return self.get_paginated_response(serializer.data)
diff --git a/dandiapi/api/views/serializers.py b/dandiapi/api/views/serializers.py
index a410f4c27..e8e15665f 100644
--- a/dandiapi/api/views/serializers.py
+++ b/dandiapi/api/views/serializers.py
@@ -169,6 +169,12 @@ def get_contact_person(self, dandiset):
 
         return contact
 
+    def get_star_count(self, dandiset):
+        return self.context['stars'][dandiset.id]['total']
+
+    def get_is_starred(self, dandiset):
+        return self.context['stars'][dandiset.id]['starred_by_current_user']
+
     most_recent_published_version = serializers.SerializerMethodField()
     draft_version = serializers.SerializerMethodField()