129 lines
4.7 KiB
Python
129 lines
4.7 KiB
Python
import logging
|
|
|
|
from django.core.management.base import BaseCommand
|
|
from django.db import models, transaction
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = "Enrich YouTube and Twitch channel metadata from upstream APIs"
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument(
|
|
"--force",
|
|
action="store_true",
|
|
help="Overwrite existing channel name and cover image",
|
|
)
|
|
parser.add_argument(
|
|
"--dry-run",
|
|
action="store_true",
|
|
help="Show what would be done without making changes",
|
|
)
|
|
parser.add_argument(
|
|
"--youtube-only",
|
|
action="store_true",
|
|
help="Only process channels with a youtube_id",
|
|
)
|
|
parser.add_argument(
|
|
"--twitch-only",
|
|
action="store_true",
|
|
help="Only process channels with a twitch_id",
|
|
)
|
|
parser.add_argument(
|
|
"--needs-metadata",
|
|
action="store_true",
|
|
help="Only process channels missing youtube_id, twitch_id, cover image, or with broken cover image",
|
|
)
|
|
|
|
def _has_broken_image(self, channel) -> bool:
|
|
if not channel.cover_image or not channel.cover_image.name:
|
|
return False
|
|
try:
|
|
return not channel.cover_image.storage.exists(channel.cover_image.name)
|
|
except Exception:
|
|
return True
|
|
|
|
def handle(self, *args, **options):
|
|
from videos.models import Channel
|
|
|
|
force = options["force"]
|
|
dry_run = options["dry_run"]
|
|
youtube_only = options["youtube_only"]
|
|
twitch_only = options["twitch_only"]
|
|
needs_metadata = options["needs_metadata"]
|
|
|
|
qs = Channel.objects.all()
|
|
|
|
if youtube_only:
|
|
qs = qs.exclude(youtube_id__isnull=True).exclude(youtube_id="")
|
|
elif twitch_only:
|
|
qs = qs.exclude(twitch_id__isnull=True).exclude(twitch_id="")
|
|
elif needs_metadata:
|
|
no_id = models.Q(youtube_id__isnull=True) & models.Q(twitch_id__isnull=True)
|
|
no_id |= models.Q(youtube_id="") & models.Q(twitch_id="")
|
|
no_id |= models.Q(youtube_id__isnull=True) & models.Q(twitch_id="")
|
|
no_id |= models.Q(youtube_id="") & models.Q(twitch_id__isnull=True)
|
|
qs = qs.filter(
|
|
no_id
|
|
| models.Q(cover_image__isnull=True)
|
|
| models.Q(cover_image="")
|
|
)
|
|
else:
|
|
qs = qs.filter(
|
|
models.Q(youtube_id__isnull=False) | models.Q(twitch_id__isnull=False)
|
|
).exclude(youtube_id="", twitch_id="")
|
|
|
|
total = qs.count()
|
|
self.stdout.write(f"Processing {total} channels from DB filter")
|
|
|
|
broken_channels = []
|
|
if needs_metadata:
|
|
broken_qs = Channel.objects.filter(
|
|
cover_image__isnull=False,
|
|
).exclude(
|
|
cover_image="",
|
|
)
|
|
if youtube_only:
|
|
broken_qs = broken_qs.exclude(youtube_id__isnull=True).exclude(youtube_id="")
|
|
elif twitch_only:
|
|
broken_qs = broken_qs.exclude(twitch_id__isnull=True).exclude(twitch_id="")
|
|
for channel in broken_qs.iterator():
|
|
if self._has_broken_image(channel):
|
|
broken_channels.append(channel)
|
|
|
|
all_channels = list(qs) + broken_channels
|
|
total = len(all_channels)
|
|
self.stdout.write(f"Total channels to process: {total}")
|
|
|
|
if dry_run:
|
|
for channel in all_channels:
|
|
source = "youtube" if channel.youtube_id else "twitch"
|
|
identifier = channel.youtube_id or channel.twitch_id
|
|
status = f"({source}: {identifier})"
|
|
if self._has_broken_image(channel):
|
|
status += " [image BROKEN]"
|
|
self.stdout.write(
|
|
f" [DRY RUN] Would fix {channel.name} {status}"
|
|
)
|
|
return
|
|
|
|
updated = 0
|
|
errors = 0
|
|
for channel in all_channels:
|
|
try:
|
|
with transaction.atomic():
|
|
channel.fix_metadata(force=force or self._has_broken_image(channel))
|
|
updated += 1
|
|
source = "youtube" if channel.youtube_id else "twitch"
|
|
self.stdout.write(f" [{updated}/{total}] {channel.name} ({source})")
|
|
except Exception as e:
|
|
errors += 1
|
|
self.stdout.write(
|
|
self.style.ERROR(f" Error updating channel {channel.name}: {e}")
|
|
)
|
|
|
|
self.stdout.write(
|
|
self.style.SUCCESS(f"\nDone! {updated} channels updated, {errors} errors")
|
|
)
|