Files
vrobbler/vrobbler/apps/charts/utils.py
Colin Powell 033239260f
All checks were successful
build / test (push) Successful in 1m57s
[perf] Add caching and lock protections
2026-06-19 13:37:30 -04:00

425 lines
13 KiB
Python

import calendar
import datetime
import logging
from typing import Optional
import pytz
from django.apps import apps
from django.conf import settings
from django.db import transaction
from django.db.models import Count, Q
from django.utils import timezone
logger = logging.getLogger(__name__)
def get_time_filter(
year: Optional[int],
month: Optional[int],
week: Optional[int],
day: Optional[int],
user,
):
"""Build a Q filter for scrobble timestamps based on time period."""
if not year:
return Q()
if user and user.is_authenticated:
tz = pytz.timezone(user.profile.timezone)
else:
tz = pytz.timezone(settings.TIME_ZONE)
if day and month:
start = timezone.make_aware(datetime.datetime(year, month, day, 0, 0, 0))
end = timezone.make_aware(datetime.datetime(year, month, day, 23, 59, 59))
elif week:
period = datetime.date.fromisocalendar(year, week, 1)
start = timezone.make_aware(
datetime.datetime.combine(period, datetime.time.min)
)
end = start + datetime.timedelta(days=6, seconds=86399)
elif month:
last_day = calendar.monthrange(year, month)[1]
start = timezone.make_aware(datetime.datetime(year, month, 1))
end = timezone.make_aware(datetime.datetime(year, month, last_day, 23, 59, 59))
else:
start = timezone.make_aware(datetime.datetime(year, 1, 1))
end = timezone.make_aware(datetime.datetime(year, 12, 31, 23, 59, 59))
return Q(timestamp__gte=start, timestamp__lte=end)
def build_charts(
user,
year: Optional[int] = None,
month: Optional[int] = None,
week: Optional[int] = None,
day: Optional[int] = None,
media_types: Optional[list] = None,
):
"""Build chart records for all or specified media types using upsert."""
ChartRecord = apps.get_model("charts", "ChartRecord")
Scrobble = apps.get_model("scrobbles", "Scrobble")
if media_types is None:
media_types = [
"artist",
"album",
"track",
"tv_series",
"video",
"podcast",
"podcast_episode",
"board_game",
"trail",
"food",
"book",
]
period_filter = Q(year=year)
if day and month:
period_filter = (
period_filter & Q(day=day) & Q(month=month) & Q(week__isnull=True)
)
elif week:
period_filter = (
period_filter & Q(week=week) & Q(day__isnull=True) & Q(month__isnull=True)
)
elif month:
period_filter = (
period_filter & Q(month=month) & Q(week__isnull=True) & Q(day__isnull=True)
)
else:
period_filter = (
period_filter
& Q(month__isnull=True)
& Q(week__isnull=True)
& Q(day__isnull=True)
)
time_filter = get_time_filter(year, month, week, day, user)
base_qs = Scrobble.objects.filter(user=user, played_to_completion=True)
if time_filter:
base_qs = base_qs.filter(time_filter)
media_config = {
"artist": {
"filter": Q(track__isnull=False) & Q(track__artists__isnull=False),
"values": "track__artists",
"annotate": Count("id", distinct=True),
},
"album": {
"filter": Q(track__isnull=False) & Q(track__album__isnull=False),
"values": "track__album",
"annotate": Count("id", distinct=True),
},
"track": {
"filter": Q(track__isnull=False),
"values": "track",
"annotate": Count("id", distinct=True),
},
"tv_series": {
"filter": Q(video__isnull=False) & Q(video__tv_series__isnull=False),
"values": "video__tv_series",
"annotate": Count("id", distinct=True),
},
"video": {
"filter": Q(video__isnull=False),
"values": "video",
"annotate": Count("id", distinct=True),
},
"podcast": {
"filter": Q(podcast_episode__isnull=False)
& Q(podcast_episode__podcast__isnull=False),
"values": "podcast_episode__podcast",
"annotate": Count("id", distinct=True),
},
"podcast_episode": {
"filter": Q(podcast_episode__isnull=False),
"values": "podcast_episode",
"annotate": Count("id", distinct=True),
},
"board_game": {
"filter": Q(board_game__isnull=False),
"values": "board_game",
"annotate": Count("id", distinct=True),
},
"trail": {
"filter": Q(trail__isnull=False),
"values": "trail",
"annotate": Count("id", distinct=True),
},
"geo_location": {
"filter": Q(geo_location__isnull=False),
"values": "geo_location",
"annotate": Count("id", distinct=True),
},
"food": {
"filter": Q(food__isnull=False),
"values": "food",
"annotate": Count("id", distinct=True),
},
"book": {
"filter": Q(book__isnull=False),
"values": "book",
"annotate": Count("id", distinct=True),
},
}
for media_type in media_types:
if media_type not in media_config:
continue
config = media_config[media_type]
results = list(
base_qs.filter(config["filter"])
.values(config["values"])
.annotate(scrobble_count=config["annotate"])
.order_by("-scrobble_count")
)
if not results:
continue
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)}
media_field = f"{media_type}_id"
with transaction.atomic():
records_to_create = []
records_to_update = []
existing = ChartRecord.objects.select_for_update().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,
"month": month,
"week": week,
"day": day,
"rank": ranks[result["scrobble_count"]],
"count": result["scrobble_count"],
}
chart_record_data[media_field] = media_id
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))
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 {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}"
)
def build_yesterdays_charts(user, media_types: Optional[list] = None) -> None:
"""Build charts for yesterday's date."""
tz = pytz.timezone(settings.TIME_ZONE)
if user and user.is_authenticated:
tz = pytz.timezone(user.profile.timezone)
now = timezone.now().astimezone(tz)
yesterday = now - datetime.timedelta(days=1)
logger.info(f"Building charts for {yesterday.date()} for {user}")
build_charts(
user=user,
year=yesterday.year,
month=yesterday.month,
day=yesterday.day,
media_types=media_types,
)
def build_weekly_charts(
user, year: int, week: int, media_types: Optional[list] = None
) -> None:
"""Build weekly charts for a specific ISO week."""
build_charts(
user=user,
year=year,
week=week,
media_types=media_types,
)
def build_monthly_charts(
user, year: int, month: int, media_types: Optional[list] = None
) -> None:
"""Build monthly charts for a specific month."""
build_charts(
user=user,
year=year,
month=month,
media_types=media_types,
)
def build_yearly_charts(user, year: int, media_types: Optional[list] = None) -> None:
"""Build yearly charts for a specific year."""
build_charts(
user=user,
year=year,
media_types=media_types,
)
def build_all_historical_charts(user, media_types: Optional[list] = None) -> None:
"""Build all historical charts for a user."""
Scrobble = apps.get_model("scrobbles", "Scrobble")
first_scrobble = (
Scrobble.objects.filter(user=user, played_to_completion=True)
.order_by("timestamp")
.first()
)
if not first_scrobble:
logger.info(f"No scrobbles found for {user}")
return
now = timezone.now()
current = first_scrobble.timestamp.date()
last_year = None
last_week = None
while current <= now.date():
year = current.year
month = current.month
day = current.day
week = current.isocalendar()[1]
build_monthly_charts(user, year, month, media_types)
build_daily_charts(user, year, month, day, media_types)
if last_week != week:
build_weekly_charts(user, year, week, media_types)
last_week = week
if last_year != year:
build_yearly_charts(user, year, media_types)
last_year = year
current += datetime.timedelta(days=1)
def build_daily_charts(
user, year: int, month: int, day: int, media_types: Optional[list] = None
) -> None:
"""Build daily charts for a specific day."""
build_charts(
user=user,
year=year,
month=month,
day=day,
media_types=media_types,
)
def build_charts_since(user, media_types: Optional[list] = None) -> None:
"""Build/update charts starting from the last known ChartRecord date."""
ChartRecord = apps.get_model("charts", "ChartRecord")
Scrobble = apps.get_model("scrobbles", "Scrobble")
latest_daily = (
ChartRecord.objects.filter(user=user, day__isnull=False)
.order_by("-year", "-month", "-day")
.first()
)
if latest_daily:
start_date = datetime.date(
latest_daily.year, latest_daily.month, latest_daily.day
)
logger.info(
f"Building charts since {start_date} for {user} (last known record)"
)
else:
first_scrobble = (
Scrobble.objects.filter(user=user, played_to_completion=True)
.order_by("timestamp")
.first()
)
if not first_scrobble:
logger.info(f"No scrobbles found for {user}")
return
start_date = first_scrobble.timestamp.date()
logger.info(
f"No existing charts found for {user}, building from "
f"first scrobble: {start_date}"
)
start_date -= datetime.timedelta(days=1)
last_week = start_date.isocalendar()[1]
last_month = start_date.month
last_year = start_date.year
now = timezone.now()
current = start_date + datetime.timedelta(days=1)
while current <= now.date():
year = current.year
month = current.month
day = current.day
week = current.isocalendar()[1]
if last_month != month:
build_monthly_charts(user, year, month, media_types)
last_month = month
build_daily_charts(user, year, month, day, media_types)
if last_week != week:
build_weekly_charts(user, year, week, media_types)
last_week = week
if last_year != year:
build_yearly_charts(user, year, media_types)
last_year = year
current += datetime.timedelta(days=1)
def rebuild_all_charts(user, media_types: Optional[list] = None) -> None:
"""Delete all existing chart records for a user and rebuild them."""
ChartRecord = apps.get_model("charts", "ChartRecord")
deleted_count = ChartRecord.objects.filter(user=user).delete()[0]
logger.info(f"Deleted {deleted_count} existing chart records for {user}")
build_all_historical_charts(user, media_types)