83 lines
3.0 KiB
Python
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))"
|
|
)
|
|
)
|