Files
vrobbler/vrobbler/apps/trends/management/commands/compute_trends.py

83 lines
3.0 KiB
Python

import logging
from django.contrib.auth import get_user_model
from django.core.management.base import BaseCommand
from django.utils import timezone
from trends.trends import TREND_REGISTRY
from trends.utils import compute_and_save_trend, get_supported_periods
logger = logging.getLogger(__name__)
User = get_user_model()
class Command(BaseCommand):
help = "Compute trends for all users (or a specific user)"
def add_arguments(self, parser):
parser.add_argument(
"--user-id",
type=int,
help="Compute trends for a specific user only",
)
def handle(self, *args, **options):
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]
else:
users = User.objects.filter(is_active=True)
total_users = len(users)
self.stdout.write(f"Computing trends for {total_users} user(s)...")
overall_start = timezone.now()
ok_count = 0
fail_count = 0
for user in users:
total_trends = len(TREND_REGISTRY)
self.stdout.write(f" {user} ({user.id}): {total_trends} trends...")
user_start = timezone.now()
user_ok = 0
user_fail = 0
for idx, (slug, _) in enumerate(TREND_REGISTRY.items(), start=1):
periods = get_supported_periods(slug)
self.stdout.write(f" [{idx}/{total_trends}] {slug}...\n")
for period in periods:
trend_start = timezone.now()
self.stdout.write(f" {period}... ", ending="")
try:
elapsed = compute_and_save_trend(user, slug, period)
self.stdout.write(self.style.SUCCESS(f"OK ({elapsed:.1f}s)"))
user_ok += 1
except Exception as e:
elapsed = (timezone.now() - trend_start).total_seconds()
self.stdout.write(
self.style.ERROR(f"FAILED after {elapsed:.1f}s: {e}")
)
user_fail += 1
user_elapsed = (timezone.now() - user_start).total_seconds()
self.stdout.write(
self.style.SUCCESS(
f" {user}: {user_ok} OK, {user_fail} failed "
f"({user_elapsed:.1f}s total)"
)
)
ok_count += user_ok
fail_count += user_fail
overall_elapsed = (timezone.now() - overall_start).total_seconds()
self.stdout.write(
self.style.SUCCESS(
f"Done! {ok_count} OK, {fail_count} failed "
f"({overall_elapsed:.1f}s across {total_users} user(s))"
)
)