vrobbler/vrobbler/apps/scrobbles/management/commands/recompute_long_play_seconds.py

from django.core.management.base import BaseCommand
from django.db import connection

from scrobbles.constants import LONG_PLAY_MEDIA
from scrobbles.models import Scrobble

BATCH_SIZE = 1000


class Command(BaseCommand):
    help = (
        "Backfill long_play_last_scrobble FK chains, then recompute "
        "long_play_seconds by walking forward through scrobbles in "
        "timestamp order with a running accumulator."
    )

    def add_arguments(self, parser):
        parser.add_argument(
            "--dry-run",
            action="store_true",
            help="Show what would be changed without making changes",
        )
        parser.add_argument(
            "--media-type",
            type=str,
            help="Only process a specific media type (e.g., Book, VideoGame)",
        )

    def handle(self, *args, **options):
        dry_run = options["dry_run"]
        media_type = options.get("media_type")

        media_types = list(LONG_PLAY_MEDIA.values())
        if media_type:
            if media_type not in media_types:
                self.stdout.write(
                    self.style.ERROR(
                        f"Invalid media type '{media_type}'. "
                        f"Valid: {', '.join(media_types)}"
                    )
                )
                return
            media_types = [media_type]

        # Step 1: backfill long_play_last_scrobble
        self.stdout.write("Step 1: Backfilling long_play_last_scrobble chains...")
        total_backfilled = 0
        for mt in media_types:
            n = self._backfill_chain(mt, dry_run)
            total_backfilled += n
            self.stdout.write(f"  {mt}: {n} scrobbles linked")

        if dry_run:
            self.stdout.write(
                self.style.WARNING(
                    f"Would backfill {total_backfilled} scrobbles total. "
                    "Run without --dry-run to apply."
                )
            )
        else:
            self.stdout.write(
                self.style.SUCCESS(f"Backfilled {total_backfilled} scrobbles")
            )

        # Step 2: recompute long_play_seconds
        self.stdout.write(
            "\nStep 2: Recomputing long_play_seconds in timestamp order..."
        )

        total_updated = 0
        for mt in media_types:
            n = self._recompute_for_media_type(mt, dry_run)
            total_updated += n
            self.stdout.write(f"  {mt}: {n} scrobbles updated")

        if dry_run:
            self.stdout.write(
                self.style.WARNING(
                    f"Dry run: would update {total_updated} scrobbles total. "
                    "Use without --dry-run to apply."
                )
            )
        else:
            self.stdout.write(
                self.style.SUCCESS(f"Updated {total_updated} scrobbles")
            )

    def _recompute_for_media_type(self, media_type: str, dry_run: bool) -> int:
        """Process scrobbles for a single media type in timestamp order with a
        running accumulator, avoiding O(n2) FK chain walks."""
        fk = _media_type_to_fk(media_type)
        fk_id = f"{fk}_id"

        scrobbles = Scrobble.objects.filter(
            media_type=media_type,
            **{f"{fk}__isnull": False},
            playback_position_seconds__isnull=False,
        ).order_by(fk_id, "user_id", "timestamp")

        total = scrobbles.count()
        if not total:
            return 0

        updated = 0
        batch = []
        last_key = None
        running_total = 0

        for scrobble in scrobbles.iterator():
            key = (getattr(scrobble, fk_id), scrobble.user_id)

            if key != last_key:
                running_total = 0
                last_key = key

            running_total += scrobble.playback_position_seconds or 0

            if scrobble.long_play_seconds != running_total:
                updated += 1
                if not dry_run:
                    scrobble.long_play_seconds = running_total
                    batch.append(scrobble)
                    if len(batch) >= BATCH_SIZE:
                        Scrobble.objects.bulk_update(batch, ["long_play_seconds"])
                        batch = []

            if scrobble.long_play_complete:
                running_total = 0

        if batch:
            Scrobble.objects.bulk_update(batch, ["long_play_seconds"])

        return updated

    def _backfill_chain(self, media_type: str, dry_run: bool) -> int:
        """Set long_play_last_scrobble on each scrobble to the previous
        scrobble for the same media+user using a single UPDATE with a
        correlated subquery."""
        fk = _media_type_to_fk(media_type)
        table = Scrobble._meta.db_table

        if dry_run:
            with connection.cursor() as cursor:
                cursor.execute(
                    f"SELECT COUNT(*) FROM {table} "
                    f"WHERE long_play_last_scrobble_id IS NULL "
                    f"AND {fk}_id IS NOT NULL"
                )
                return cursor.fetchone()[0]

        with connection.cursor() as cursor:
            cursor.execute(
                f"UPDATE {table} "
                f"SET long_play_last_scrobble_id = ("
                f"  SELECT id FROM {table} AS prev "
                f"  WHERE prev.{fk}_id = {table}.{fk}_id "
                f"    AND prev.user_id = {table}.user_id "
                f"    AND prev.timestamp < {table}.timestamp "
                f"  ORDER BY prev.timestamp DESC LIMIT 1"
                f") "
                f"WHERE long_play_last_scrobble_id IS NULL "
                f"AND {fk}_id IS NOT NULL"
            )
            return cursor.rowcount


def _media_type_to_fk(media_type):
    mapping = {
        "VideoGame": "video_game",
        "Book": "book",
        "BrickSet": "brick_set",
        "Task": "task",
    }
    return mapping.get(media_type)