Aurora-Core/app/scripts/bulk_import.py

"""One-off ops script: enqueue an ingestion job for every id in TMDb's
daily export, for one or more entity types -- a full-catalog bulk sync.

    python -m app.scripts.bulk_import movie series person collection
    python -m app.scripts.bulk_import movie --dry-run
    python -m app.scripts.bulk_import series --no-seasons

Every id goes through the exact same IngestionJob + arq task pipeline the
admin API uses for a single id (get_or_create_ingestion_job is
idempotency-keyed on provider+entity_type+external_id), so interrupting
and re-running this script is safe -- already-created jobs are reused,
never duplicated, and re-enqueuing an already-succeeded job just re-syncs
that entity.

This does not itself rate-limit against TMDb -- the actual API calls
happen later, inside the worker tasks this enqueues, so TMDb's request
rate is governed by WORKER_CONCURRENCY (how many workers process the
queue concurrently), not by how fast this script can enqueue. Enqueueing
does hit Postgres once per id (via get_or_create_ingestion_job), so this
runs a bounded number of ids concurrently to avoid one id's DB round-trip
serializing the next -- see _CONCURRENCY below.

Expect this to run for hours across a full catalog (TMDb has roughly 1M+
movies, 200K+ TV series, 3M+ people, 150K+ collections) -- run it inside
the worker container detached (e.g. `docker compose -f
docker-compose.prod.yml exec -d worker python -m app.scripts.bulk_import
...`) or under nohup/tmux/screen so it survives a disconnect, and expect
substantial TMDb API traffic, database growth, and (since artwork
self-hosting cascades on every successful ingestion, per
docs/database.md's artwork storage section) disk usage for downloaded
images while it runs.
"""

from __future__ import annotations

import argparse
import asyncio

import httpx
from arq import ArqRedis

from app.core.database import dispose_engine
from app.core.logging import configure_logging, get_logger
from app.core.queue import close_arq_pool, get_arq_pool
from app.services.bulk_import import fetch_daily_export_ids
from app.services.ingestion.runner import get_or_create_ingestion_job

logger = get_logger(__name__)

_TASK_NAME: dict[str, str] = {
    "movie": "ingest_movie_task",
    "series": "ingest_series_task",
    "person": "ingest_person_task",
    "collection": "ingest_collection_task",
}
_JOB_TYPE: dict[str, str] = {
    "movie": "movie_sync",
    "series": "series_sync",
    "person": "person_sync",
    "collection": "collection_sync",
}

# How many ids to create IngestionJob rows / enqueue for concurrently.
# Bounded so this doesn't try to open millions of DB round-trips at once,
# but high enough that per-id latency isn't fully serialized.
_CONCURRENCY = 50


async def _enqueue_one(
    queue: ArqRedis, entity_type: str, external_id: str, *, include_seasons: bool
) -> None:
    job_id = await get_or_create_ingestion_job(_JOB_TYPE[entity_type], entity_type, external_id)
    task_name = _TASK_NAME[entity_type]
    if entity_type == "series":
        await queue.enqueue_job(task_name, str(job_id), external_id, include_seasons)
    else:
        await queue.enqueue_job(task_name, str(job_id), external_id)


async def _bulk_enqueue(entity_type: str, *, dry_run: bool, include_seasons: bool) -> None:
    async with httpx.AsyncClient(timeout=120.0) as client:
        ids = await fetch_daily_export_ids(entity_type, client=client)
    logger.info("bulk_import_fetched_ids", entity_type=entity_type, count=len(ids))

    if dry_run:
        logger.info("bulk_import_dry_run", entity_type=entity_type, would_enqueue=len(ids))
        return

    queue = await get_arq_pool()
    # A fixed pool of _CONCURRENCY consumer tasks pulling from a plain
    # asyncio.Queue -- NOT asyncio.gather() over one Task per id. gather()
    # eagerly wraps every argument into a Task before awaiting any of them,
    # so with a catalog in the millions that means millions of live Task +
    # coroutine objects created up front regardless of how tightly a
    # semaphore bounds actual concurrent execution -- a real bottleneck
    # (and memory risk) discovered running this against the real 1.24M-
    # movie export, where it looked hung for minutes just getting through
    # that setup before any work could start.
    id_queue: asyncio.Queue[str] = asyncio.Queue()
    for external_id in ids:
        id_queue.put_nowait(external_id)

    enqueued = 0
    lock = asyncio.Lock()

    async def _consumer() -> None:
        nonlocal enqueued
        while True:
            try:
                external_id = id_queue.get_nowait()
            except asyncio.QueueEmpty:
                return
            await _enqueue_one(queue, entity_type, external_id, include_seasons=include_seasons)
            async with lock:
                enqueued += 1
                if enqueued % 5000 == 0:
                    logger.info(
                        "bulk_import_progress",
                        entity_type=entity_type,
                        enqueued=enqueued,
                        total=len(ids),
                    )

    await asyncio.gather(*(_consumer() for _ in range(_CONCURRENCY)))
    logger.info("bulk_import_complete", entity_type=entity_type, enqueued=enqueued)


async def _main(entity_types: list[str], *, dry_run: bool, include_seasons: bool) -> None:
    configure_logging()
    try:
        for entity_type in entity_types:
            await _bulk_enqueue(entity_type, dry_run=dry_run, include_seasons=include_seasons)
    finally:
        await close_arq_pool()
        await dispose_engine()


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Bulk-enqueue ingestion jobs from TMDb's daily ID export files."
    )
    parser.add_argument(
        "entity_types",
        nargs="+",
        choices=list(_TASK_NAME),
        help="One or more entity types to bulk-import.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Fetch and count ids without enqueueing anything.",
    )
    parser.add_argument(
        "--no-seasons",
        action="store_true",
        help="For 'series': don't cascade into per-season/episode ingestion "
        "(cuts total TMDb calls roughly in half for TV, at the cost of "
        "season/episode data). Cascading is the default for a full pull.",
    )
    args = parser.parse_args()
    asyncio.run(
        _main(args.entity_types, dry_run=args.dry_run, include_seasons=not args.no_seasons)
    )


if __name__ == "__main__":
    main()