[charts] Upsert records instead of deleting

vrobbler/apps/charts/utils.py

@@ -56,7 +56,7 @@ def build_charts(
     day: Optional[int] = None,
     media_types: Optional[list] = None,
 ):
-    """Build chart records for all or specified media types."""
+    """Build chart records for all or specified media types using upsert."""
     ChartRecord = apps.get_model("charts", "ChartRecord")
     Scrobble = apps.get_model("scrobbles", "Scrobble")
 
@@ -96,8 +96,6 @@ def build_charts(
             & Q(day__isnull=True)
         )
 
-    ChartRecord.objects.filter(period_filter, user=user).delete()
-
     time_filter = get_time_filter(year, month, week, day, user)
     base_qs = Scrobble.objects.filter(user=user, played_to_completion=True)
 
@@ -168,8 +166,6 @@ def build_charts(
         },
     }
 
-    BATCH_SIZE = 500
-
     for media_type in media_types:
         if media_type not in media_config:
             continue
@@ -189,9 +185,23 @@ def build_charts(
         unique_counts = sorted(set(r["scrobble_count"] for r in results), reverse=True)
         ranks = {count: rank for rank, count in enumerate(unique_counts, start=1)}
 
-        chart_records = []
-        total_created = 0
+        media_field = f"{media_type}_id"
+        records_to_create = []
+        records_to_update = []
+
+        existing = ChartRecord.objects.filter(
+            period_filter, user=user, **{media_field + "__isnull": False}
+        )
+        existing_by_media_id = {getattr(r, media_field): r for r in existing}
+        found_media_ids = set()
+
         for result in results:
+            media_id = result[config["values"]]
+            if media_id is None:
+                continue
+
+            found_media_ids.add(media_id)
+
             chart_record_data = {
                 "user_id": user.id,
                 "year": year,
@@ -201,20 +211,33 @@ def build_charts(
                 "rank": ranks[result["scrobble_count"]],
                 "count": result["scrobble_count"],
             }
-            chart_record_data[f"{media_type}_id"] = result[config["values"]]
-            chart_records.append(ChartRecord(**chart_record_data))
+            chart_record_data[media_field] = media_id
 
-            if len(chart_records) >= BATCH_SIZE:
-                ChartRecord.objects.bulk_create(chart_records, batch_size=BATCH_SIZE)
-                total_created += len(chart_records)
-                chart_records = []
+            if media_id in existing_by_media_id:
+                existing_record = existing_by_media_id[media_id]
+                existing_record.rank = chart_record_data["rank"]
+                existing_record.count = chart_record_data["count"]
+                records_to_update.append(existing_record)
+            else:
+                records_to_create.append(ChartRecord(**chart_record_data))
 
-        if chart_records:
-            ChartRecord.objects.bulk_create(chart_records, batch_size=BATCH_SIZE)
-            total_created += len(chart_records)
+        ids_to_delete = [
+            r.id for r in existing if getattr(r, media_field) not in found_media_ids
+        ]
+        if ids_to_delete:
+            ChartRecord.objects.filter(id__in=ids_to_delete).delete()
+
+        if records_to_update:
+            ChartRecord.objects.bulk_update(
+                records_to_update, ["rank", "count"], batch_size=500
+            )
+
+        if records_to_create:
+            ChartRecord.objects.bulk_create(records_to_create, batch_size=500)
 
         logger.info(
-            f"Built {total_created} chart records for {media_type}, period "
+            f"Built {len(records_to_create)} new, {len(records_to_update)} updated "
+            f"chart records for {media_type}, period "
             f"{year}-{month or ''}-{week or ''}-{day or ''} for {user}"
         )