[scrobbles] Improve query efficency

vrobbler/apps/music/aggregators.py

@@ -16,9 +16,7 @@ def scrobble_counts(user=None):
         now = now_user_timezone(user.profile)
         user_filter = Q(user=user)
 
-    start_of_today = datetime.combine(
-        now.date(), datetime.min.time(), now.tzinfo
-    )
+    start_of_today = datetime.combine(now.date(), datetime.min.time(), now.tzinfo)
     starting_day_of_current_week = now.date() - timedelta(
         days=now.today().isoweekday() % 7
     )
@@ -31,9 +29,7 @@ def scrobble_counts(user=None):
         media_type=Scrobble.MediaType.TRACK,
     )
     data = {}
-    data["today"] = finished_scrobbles_qs.filter(
-        timestamp__gte=start_of_today
-    ).count()
+    data["today"] = finished_scrobbles_qs.filter(timestamp__gte=start_of_today).count()
     data["week"] = finished_scrobbles_qs.filter(
         timestamp__gte=starting_day_of_current_week
     ).count()
@@ -47,9 +43,7 @@ def scrobble_counts(user=None):
     return data
 
 
-def week_of_scrobbles(
-    user=None, start=None, media: str = "tracks"
-) -> dict[str, int]:
+def week_of_scrobbles(user=None, start=None, media: str = "tracks") -> dict[str, int]:
 
     now = timezone.now()
     user_filter = Q()
@@ -101,9 +95,7 @@ def live_charts(
     seven_days_ago = now - timedelta(days=7)
     thirty_days_ago = now - timedelta(days=30)
     start_of_today = datetime.combine(now, datetime.min.time(), tzinfo)
-    start_day_of_week = start_of_today - timedelta(
-        days=now.today().isoweekday() % 7
-    )
+    start_day_of_week = start_of_today - timedelta(days=now.today().isoweekday() % 7)
     start_day_of_month = now.replace(day=1)
     start_day_of_year = now.replace(month=1, day=1)
 
@@ -120,17 +112,13 @@ def live_charts(
     }
 
     time_filter = Q()
-    completion_filter = Q(
-        scrobble__user=user, scrobble__played_to_completion=True
-    )
+    completion_filter = Q(scrobble__user=user, scrobble__played_to_completion=True)
     user_filter = Q(scrobble__user=user)
     count_field = "scrobble"
 
     if media_type == "Artist":
         for period, query_dict in period_queries.items():
-            period_queries[period] = {
-                "track__" + k: v for k, v in query_dict.items()
-            }
+            period_queries[period] = {"track__" + k: v for k, v in query_dict.items()}
         completion_filter = Q(
             track__scrobble__user=user,
             track__scrobble__played_to_completion=True,
@@ -175,9 +163,7 @@ def live_tv_charts(
     seven_days_ago = now - timedelta(days=7)
     thirty_days_ago = now - timedelta(days=30)
     start_of_today = datetime.combine(now, datetime.min.time(), tzinfo)
-    start_day_of_week = start_of_today - timedelta(
-        days=now.today().isoweekday() % 7
-    )
+    start_day_of_week = start_of_today - timedelta(days=now.today().isoweekday() % 7)
     start_day_of_month = now.replace(day=1)
     start_day_of_year = now.replace(month=1, day=1)
 
@@ -232,9 +218,7 @@ def live_youtube_channel_charts(
     seven_days_ago = now - timedelta(days=7)
     thirty_days_ago = now - timedelta(days=30)
     start_of_today = datetime.combine(now, datetime.min.time(), tzinfo)
-    start_day_of_week = start_of_today - timedelta(
-        days=now.today().isoweekday() % 7
-    )
+    start_day_of_week = start_of_today - timedelta(days=now.today().isoweekday() % 7)
     start_day_of_year = now.replace(month=1, day=1)
 
     period_queries = {
@@ -269,3 +253,113 @@ def live_youtube_channel_charts(
         .filter(num_scrobbles__gt=0)
         .order_by("-num_scrobbles")[:limit]
     )
+
+
+def batch_live_charts(
+    user: "User",
+    media_type: str = "Track",
+    limit: int = 15,
+    app_label: str = "music",
+) -> dict[str, list]:
+    from django.db.models import Count, Q
+
+    now = timezone.now()
+    tzinfo = now.tzinfo
+    now = now.date()
+    if user.is_authenticated:
+        now = now_user_timezone(user.profile)
+        tzinfo = now.tzinfo
+
+    seven_days_ago = now - timedelta(days=7)
+    thirty_days_ago = now - timedelta(days=30)
+    start_of_today = datetime.combine(now, datetime.min.time(), tzinfo)
+    start_day_of_year = now.replace(month=1, day=1)
+
+    media_model = apps.get_model(app_label=app_label, model_name=media_type)
+
+    if media_type == "Artist":
+        base_filter = Q(
+            track__scrobble__user=user, track__scrobble__played_to_completion=True
+        )
+        count_field = "track__scrobble"
+    else:
+        base_filter = Q(scrobble__user=user, scrobble__played_to_completion=True)
+        count_field = "scrobble"
+
+    scrobble_filter_prefix = (
+        "track__scrobble__" if media_type == "Artist" else "scrobble__"
+    )
+
+    queryset = media_model.objects.filter(
+        Q(track__scrobble__user=user)
+        if media_type == "Artist"
+        else Q(scrobble__user=user)
+    ).annotate(
+        num_scrobbles_today=Count(
+            count_field,
+            filter=Q(
+                **{
+                    f"{scrobble_filter_prefix}played_to_completion": True,
+                    f"{scrobble_filter_prefix}timestamp__gte": start_of_today,
+                }
+            ),
+            distinct=True,
+        ),
+        num_scrobbles_last7=Count(
+            count_field,
+            filter=Q(
+                **{
+                    f"{scrobble_filter_prefix}played_to_completion": True,
+                    f"{scrobble_filter_prefix}timestamp__gte": seven_days_ago,
+                }
+            ),
+            distinct=True,
+        ),
+        num_scrobbles_last30=Count(
+            count_field,
+            filter=Q(
+                **{
+                    f"{scrobble_filter_prefix}played_to_completion": True,
+                    f"{scrobble_filter_prefix}timestamp__gte": thirty_days_ago,
+                }
+            ),
+            distinct=True,
+        ),
+        num_scrobbles_year=Count(
+            count_field,
+            filter=Q(
+                **{
+                    f"{scrobble_filter_prefix}played_to_completion": True,
+                    f"{scrobble_filter_prefix}timestamp__gte": start_day_of_year,
+                }
+            ),
+            distinct=True,
+        ),
+        num_scrobbles=Count(
+            count_field,
+            filter=Q(
+                **{
+                    f"{scrobble_filter_prefix}played_to_completion": True,
+                }
+            ),
+            distinct=True,
+        ),
+    )
+
+    items = list(queryset.order_by("-num_scrobbles")[:limit])
+    for item in items:
+        item.count = item.num_scrobbles
+
+    def filter_by_count_attr(qs, attr_name):
+        return sorted(
+            [item for item in items if getattr(item, attr_name) > 0],
+            key=lambda x: -getattr(x, attr_name),
+        )[:limit]
+
+    return {
+        "today": filter_by_count_attr(items, "num_scrobbles_today"),
+        "last7": filter_by_count_attr(items, "num_scrobbles_last7"),
+        "last30": filter_by_count_attr(items, "num_scrobbles_last30"),
+        "year": filter_by_count_attr(items, "num_scrobbles_year"),
+        "all": items,
+    }