diff --git a/vrobbler/apps/charts/utils.py b/vrobbler/apps/charts/utils.py index d8d6718..140ef5d 100644 --- a/vrobbler/apps/charts/utils.py +++ b/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}" )