FabledCurator/backend/app/tasks/migration.py

"""FC-5 run_migration Celery task.

Dispatches to the right migrator based on `kind`. Updates MigrationRun
row's status/counts/finished_at as it runs. Failures set status='error'
with the error message preserved.

kinds: gs_ingest, ir_ingest, tag_apply, ml_queue, verify, cleanup
(backup + rollback retired 2026-05-24 → see /api/system/backup/*)
"""
from __future__ import annotations

import asyncio
import logging
from datetime import UTC, datetime
from pathlib import Path

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from ..celery_app import celery
from ..models import MigrationRun
from ..services.credential_crypto import CredentialCrypto
from ..services.migrators import cleanup as cleanup_mod
from ..services.migrators import gs_ingest, ir_ingest, ml_queue, tag_apply, verify
from ._async_session import async_session_factory

log = logging.getLogger(__name__)

IMAGES_ROOT = Path("/images")
_KEY_PATH = IMAGES_ROOT / "secrets" / "credential_key.b64"


async def _update_run(
    db: AsyncSession, run_id: int, *,
    status: str | None = None, counts: dict | None = None,
    error: str | None = None, finished_at: datetime | None = None,
    metadata_patch: dict | None = None,
) -> None:
    run = (await db.execute(
        select(MigrationRun).where(MigrationRun.id == run_id)
    )).scalar_one()
    if status is not None:
        run.status = status
    if counts is not None:
        run.counts = counts
    if error is not None:
        run.error = error
    if finished_at is not None:
        run.finished_at = finished_at
    if metadata_patch:
        run.metadata_ = {**(run.metadata_ or {}), **metadata_patch}
    await db.commit()


async def _run_async(run_id: int, kind: str, params: dict) -> dict:
    factory, engine = async_session_factory()
    try:
        async with factory() as db:
            await _update_run(db, run_id, status="running")
            try:
                if kind in ("backup", "rollback"):
                    raise ValueError(
                        f"kind {kind!r} retired in FC-3h; "
                        "use /api/system/backup/* instead"
                    )

                elif kind == "gs_ingest":
                    fc_crypto = CredentialCrypto(_KEY_PATH)
                    counts = await gs_ingest.migrate_async(
                        db, data=params["data"],
                        fc_crypto=fc_crypto,
                        dry_run=params.get("dry_run", False),
                    )
                    await _update_run(
                        db, run_id, status="ok", counts=counts,
                        finished_at=datetime.now(UTC),
                    )
                    return counts

                elif kind == "ir_ingest":
                    counts = await ir_ingest.migrate_async(
                        db, data=params["data"],
                        images_root=IMAGES_ROOT,
                        dry_run=params.get("dry_run", False),
                    )
                    await _update_run(
                        db, run_id, status="ok", counts=counts,
                        finished_at=datetime.now(UTC),
                    )
                    return counts

                elif kind == "tag_apply":
                    result = await tag_apply.apply_async(
                        db, images_root=IMAGES_ROOT,
                        dry_run=params.get("dry_run", False),
                    )
                    await _update_run(
                        db, run_id, status="ok",
                        counts=result["counts"],
                        finished_at=datetime.now(UTC),
                        metadata_patch={"unmatched": result["unmatched"]},
                    )
                    return result

                elif kind == "ml_queue":
                    count = await ml_queue.queue_all_unprocessed_async(db)
                    await _update_run(
                        db, run_id, status="ok",
                        counts={"rows_processed": count, "rows_inserted": 0,
                                "rows_skipped": 0, "files_copied": 0,
                                "bytes_copied": 0, "conflicts": 0},
                        finished_at=datetime.now(UTC),
                    )
                    return {"queued": count}

                elif kind == "verify":
                    checks = await verify.verify_async(db, expected=params.get("expected"))
                    sample = await verify.verify_sha256_sample(
                        db, sample_size=params.get("sample_size", 20),
                    )
                    await _update_run(
                        db, run_id, status="ok",
                        counts={"rows_processed": sample["sample_size"],
                                "rows_inserted": 0, "rows_skipped": 0,
                                "files_copied": 0, "bytes_copied": 0,
                                "conflicts": sample["mismatched"] + sample["missing"]},
                        finished_at=datetime.now(UTC),
                        metadata_patch={"checks": checks, "sample": sample},
                    )
                    return {"checks": checks, "sample": sample}

                elif kind == "cleanup":
                    slug = params.get("slug")
                    if not slug:
                        raise ValueError("cleanup requires params.slug")
                    result = await cleanup_mod.cleanup_artist_async(
                        db, slug=slug, images_root=IMAGES_ROOT,
                        dry_run=params.get("dry_run", False),
                        source_path_prefix=params.get("source_path_prefix"),
                    )
                    await _update_run(
                        db, run_id, status="ok",
                        counts=result["counts"],
                        finished_at=datetime.now(UTC),
                        metadata_patch={
                            "artist": result["artist"],
                            "summary": result["summary"],
                        },
                    )
                    return result

                else:
                    raise ValueError(f"unknown kind: {kind}")

            except Exception as exc:
                log.exception("migration kind=%s failed", kind)
                await _update_run(
                    db, run_id, status="error", error=str(exc),
                    finished_at=datetime.now(UTC),
                )
                raise
    finally:
        await engine.dispose()


@celery.task(name="backend.app.tasks.migration.run_migration", bind=True, acks_late=True)
def run_migration(self, run_id: int, kind: str, params: dict) -> dict:
    """FC-5: dispatch a migration kind. Updates MigrationRun row as it goes."""
    return asyncio.run(_run_async(run_id, kind, params))