Files
vrobbler/vrobbler/apps/charts/management/commands/build_charts.py

114 lines
3.6 KiB
Python

import logging
from charts.utils import (
build_all_historical_charts,
build_charts_since,
rebuild_all_charts,
)
from django.core.management.base import BaseCommand
logger = logging.getLogger(__name__)
def get_chart_counts():
from charts.models import ChartRecord
daily = ChartRecord.objects.filter(day__isnull=False).count()
weekly = ChartRecord.objects.filter(week__isnull=False, day__isnull=True).count()
monthly = ChartRecord.objects.filter(
month__isnull=False, week__isnull=True, day__isnull=True
).count()
yearly = ChartRecord.objects.filter(
month__isnull=True, week__isnull=True, day__isnull=True
).count()
total = ChartRecord.objects.count()
return {
"daily": daily,
"weekly": weekly,
"monthly": monthly,
"yearly": yearly,
"total": total,
}
def print_chart_summary(before=None):
counts = get_chart_counts()
if before:
diff = {k: counts[k] - before[k] for k in counts}
print(
f" Daily: {counts['daily']} ({diff['daily']:+d}), "
f"Weekly: {counts['weekly']} ({diff['weekly']:+d}), "
f"Monthly: {counts['monthly']} ({diff['monthly']:+d}), "
f"Yearly: {counts['yearly']} ({diff['yearly']:+d}), "
f"Total: {counts['total']} ({diff['total']:+d})"
)
else:
print(
f" Daily: {counts['daily']}, "
f"Weekly: {counts['weekly']}, "
f"Monthly: {counts['monthly']}, "
f"Yearly: {counts['yearly']}, "
f"Total: {counts['total']}"
)
class Command(BaseCommand):
help = "Build missing charts for all users"
def add_arguments(self, parser):
parser.add_argument(
"--user-id",
type=int,
help="Build charts for a specific user only",
)
parser.add_argument(
"--full-rebuild",
action="store_true",
help="Delete existing chart records and rebuild from scratch",
)
def handle(self, *args, **options):
from django.contrib.auth import get_user_model
User = get_user_model()
full_rebuild = options["full_rebuild"]
if full_rebuild:
self.stdout.write("Full rebuild requested - will delete existing records")
if options["user_id"]:
user = User.objects.filter(id=options["user_id"]).first()
if not user:
self.stderr.write(
self.style.ERROR(f"User with id {options['user_id']} not found")
)
return
users = [user]
action = "Full rebuild for" if full_rebuild else "Progressive build for"
self.stdout.write(f"{action} user: {user}")
else:
users = User.objects.all()
action = "Full rebuild for" if full_rebuild else "Progressive build for"
self.stdout.write(f"{action} charts for {users.count()} users...")
before_counts = get_chart_counts()
self.stdout.write("Chart counts before:")
print_chart_summary()
for user in users:
try:
logger.info(f"Building charts for {user}")
if full_rebuild:
rebuild_all_charts(user)
else:
build_charts_since(user)
self.stdout.write(self.style.SUCCESS(f" OK {user}"))
except Exception as e:
self.stderr.write(f" FAILED {user}: {e}")
self.stdout.write("Chart counts after (diff):")
print_chart_summary(before_counts)
self.stdout.write(self.style.SUCCESS("Done!"))