104 lines
3.7 KiB
Python
104 lines
3.7 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
|
|
skipped_count = 0
|
|
|
|
for user in users:
|
|
try:
|
|
profile = user.profile
|
|
if profile.trends_disabled:
|
|
self.stdout.write(
|
|
self.style.WARNING(
|
|
f" {user} ({user.id}): trends disabled globally, skipping"
|
|
)
|
|
)
|
|
skipped_count += len(TREND_REGISTRY)
|
|
continue
|
|
disabled_trends = set(profile.disabled_trends or [])
|
|
except Exception:
|
|
disabled_trends = set()
|
|
|
|
active_slugs = [
|
|
s for s in TREND_REGISTRY if s not in disabled_trends
|
|
]
|
|
total_trends = len(active_slugs)
|
|
self.stdout.write(
|
|
f" {user} ({user.id}): {total_trends} trends ("
|
|
f"{len(disabled_trends & set(TREND_REGISTRY.keys()))} disabled)..."
|
|
)
|
|
user_start = timezone.now()
|
|
user_ok = 0
|
|
user_fail = 0
|
|
|
|
for idx, slug in enumerate(active_slugs, 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, {skipped_count} skipped "
|
|
f"({overall_elapsed:.1f}s across {total_users} user(s))"
|
|
)
|
|
)
|