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 sqlalchemy.ext.asyncio import AsyncSession

from app.core.database import session_scope
from app.core.logging import get_logger
from app.models.provider import IngestionJob
from app.providers.base import MetadataProvider
from app.providers.tmdb import TMDbProvider
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.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))
        raise
    finally:
        await provider.aclose()


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
        ),
    )