Aurora-Core/app/services/ingestion/runner.py

"""Orchestrates one ingestion job end-to-end: create/track the `IngestionJob`
row, run the provider fetch + normalize + upsert pipeline, and record
success/failure. Used by both the admin trigger endpoint (section 17) and
the arq worker task (section 15) so every entry point shares identical
job-lifecycle and idempotency behavior.
"""

from __future__ import annotations

import uuid
from collections.abc import Awaitable, Callable

from arq import Retry
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.database import session_scope
from app.core.logging import get_logger
from app.models.enums import IngestionJobStatus
from app.models.provider import IngestionJob
from app.providers.base import MetadataProvider
from app.providers.tmdb import TMDbProvider
from app.repositories.polymorphic import get_primary_artwork_ids
from app.repositories.provider_repository import (
    create_ingestion_job,
    get_ingestion_job,
    get_or_create_provider,
    mark_job_failed,
    mark_job_running,
    mark_job_succeeded,
)
from app.repositories.tv_repository import get_seasons_for_series
from app.services.ingestion.movie_ingestion import ingest_movie
from app.services.ingestion.person_ingestion import ingest_person
from app.services.ingestion.series_ingestion import ingest_season, ingest_series

logger = get_logger(__name__)


async def _require_job(session: AsyncSession, job_id: uuid.UUID) -> IngestionJob:
    job = await get_ingestion_job(session, job_id)
    if job is None:
        # Only possible if the row was deleted out-of-band; the job was
        # created moments earlier in this same call chain.
        raise RuntimeError(f"ingestion job {job_id} vanished mid-run")
    return job


async def get_or_create_ingestion_job(job_type: str, entity_type: str, external_id: str) -> uuid.UUID:
    """Create the DB-tracked `IngestionJob` row (or return the existing one
    for this provider+entity+external_id, keyed by `idempotency_key`)
    without running the pipeline.

    Called synchronously from the admin endpoint, *before* enqueuing to arq,
    specifically so the HTTP response can return a `job_id` that
    `GET /admin/ingest/jobs/{job_id}` can immediately resolve — arq's own
    queue-internal job id is a different, unrelated identifier and was
    never valid here.
    """

    idempotency_key = f"tmdb:{entity_type}:{external_id}"
    async with session_scope() as session:
        provider_row = await get_or_create_provider(session, "tmdb", "The Movie Database (TMDb)")
        job = await create_ingestion_job(
            session,
            provider_id=provider_row.id,
            job_type=job_type,
            idempotency_key=idempotency_key,
            entity_type=entity_type,
            external_id=external_id,
        )
        return job.id


async def _run_ingestion(
    job_id: uuid.UUID,
    label: str,
    work: Callable[[AsyncSession, MetadataProvider], Awaitable[uuid.UUID]],
) -> uuid.UUID:
    """Shared job-lifecycle wrapper: mark running, run `work` in its own
    session/transaction, mark succeeded/failed. `work` does the actual
    fetch+normalize+upsert and returns the resulting Aurora entity id."""

    async with session_scope() as session:
        job = await _require_job(session, job_id)
        await mark_job_running(session, job)

    provider = TMDbProvider()
    try:
        async with session_scope() as session:
            result_id = await work(session, provider)
        async with session_scope() as session:
            job = await _require_job(session, job_id)
            await mark_job_succeeded(session, job, result_id)
        return result_id
    except Exception as exc:
        logger.error(f"{label}_failed", job_id=str(job_id), error=str(exc))
        async with session_scope() as session:
            job = await _require_job(session, job_id)
            await mark_job_failed(session, job, str(exc))
            attempts = job.attempts
            still_retryable = job.status == IngestionJobStatus.RETRYING
        if still_retryable:
            # arq only auto-retries on its own Retry (or CancelledError),
            # never on an arbitrary exception — mark_job_failed setting the
            # DB status to RETRYING was, on its own, purely cosmetic; this
            # is what actually gets the job back on the queue instead of
            # leaving it permanently "retrying" without ever being retried.
            raise Retry(defer=_retry_backoff_seconds(attempts)) from exc
        raise
    finally:
        await provider.aclose()


def _retry_backoff_seconds(attempt: int) -> float:
    return min(60.0, 2.0**attempt)


async def run_movie_ingestion(job_id: uuid.UUID, external_id: str) -> uuid.UUID:
    return await _run_ingestion(
        job_id,
        "movie_ingestion",
        lambda session, provider: ingest_movie(
            session, provider, external_id, ingestion_job_id=job_id
        ),
    )


async def run_person_ingestion(job_id: uuid.UUID, external_id: str) -> uuid.UUID:
    return await _run_ingestion(
        job_id,
        "person_ingestion",
        lambda session, provider: ingest_person(
            session, provider, external_id, ingestion_job_id=job_id
        ),
    )


async def run_series_ingestion(job_id: uuid.UUID, external_id: str) -> uuid.UUID:
    return await _run_ingestion(
        job_id,
        "series_ingestion",
        lambda session, provider: ingest_series(
            session, provider, external_id, ingestion_job_id=job_id
        ),
    )


async def run_season_ingestion(
    job_id: uuid.UUID, series_external_id: str, season_number: int
) -> uuid.UUID:
    return await _run_ingestion(
        job_id,
        "season_ingestion",
        lambda session, provider: ingest_season(
            session, provider, series_external_id, season_number, ingestion_job_id=job_id
        ),
    )


async def get_series_season_numbers(series_id: uuid.UUID) -> list[int]:
    """All season numbers a just-ingested series has, for cascading into
    per-season ingestion jobs (see `app/workers/tasks.py:ingest_series_task`,
    `include_seasons`). A short-lived session of its own since this runs
    after `run_series_ingestion`'s own session has already closed."""

    async with session_scope() as session:
        seasons = await get_seasons_for_series(session, series_id)
        return [season.season_number for season in seasons]


async def get_series_season_ids(series_id: uuid.UUID) -> list[uuid.UUID]:
    async with session_scope() as session:
        seasons = await get_seasons_for_series(session, series_id)
        return [season.id for season in seasons]


async def get_primary_artwork_for_download(
    entity_type: str, entity_ids: list[uuid.UUID]
) -> list[uuid.UUID]:
    """Artwork ids worth auto-downloading for local hosting after an
    ingestion succeeds — see app/services/artwork/download.py and
    app/workers/tasks.py, which enqueue one download_artwork_task per id
    this returns. A short-lived session of its own, same reasoning as
    get_series_season_numbers above."""

    async with session_scope() as session:
        return await get_primary_artwork_ids(session, entity_type, entity_ids)