diff --git a/backend/app/celery_app.py b/backend/app/celery_app.py index 613a78f..e5d91f2 100644 --- a/backend/app/celery_app.py +++ b/backend/app/celery_app.py @@ -28,6 +28,7 @@ def make_celery() -> Celery: "backend.app.tasks.import_file", "backend.app.tasks.thumbnail", "backend.app.tasks.maintenance", + "backend.app.tasks.migration", "backend.app.tasks.ml", "backend.app.tasks.download", ], @@ -41,6 +42,7 @@ def make_celery() -> Celery: "backend.app.tasks.download.*": {"queue": "download"}, "backend.app.tasks.scan.*": {"queue": "scan"}, "backend.app.tasks.maintenance.*": {"queue": "maintenance"}, + "backend.app.tasks.migration.*": {"queue": "maintenance"}, }, # Heavy ML tasks need fair dispatch — see ImageRepo's precedent. task_acks_late=True, diff --git a/backend/app/tasks/migration.py b/backend/app/tasks/migration.py new file mode 100644 index 0000000..dd200bd --- /dev/null +++ b/backend/app/tasks/migration.py @@ -0,0 +1,186 @@ +"""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: backup, gs_ingest, ir_ingest, tag_apply, ml_queue, verify, rollback +""" +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, async_sessionmaker, create_async_engine + +from ..celery_app import celery +from ..config import get_config +from ..models import MigrationRun +from ..services.credential_crypto import CredentialCrypto +from ..services.migrators import ( + backup as backup_mod, + gs_ingest, + ir_ingest, + ml_queue, + rollback as rollback_mod, + tag_apply, + verify, +) + +log = logging.getLogger(__name__) + +IMAGES_ROOT = Path("/images") +_KEY_PATH = IMAGES_ROOT / "secrets" / "credential_key.b64" + + +def _async_session_factory(): + cfg = get_config() + engine = create_async_engine(cfg.database_url, future=True, pool_pre_ping=True) + return async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False), engine + + +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 == "backup": + manifest = backup_mod.create_backup( + db_url=get_config().database_url_sync, + images_root=IMAGES_ROOT, + tag=params.get("tag", "manual"), + ) + await _update_run( + db, run_id, status="ok", + counts={"rows_processed": 0, "rows_inserted": 0, + "rows_skipped": 0, "files_copied": 0, + "bytes_copied": 0, "conflicts": 0}, + finished_at=datetime.now(UTC), + metadata_patch={"manifest": manifest}, + ) + return manifest + + 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 == "rollback": + result = rollback_mod.rollback_to_pre_migration( + db_url=get_config().database_url_sync, + images_root=IMAGES_ROOT, + ) + await _update_run( + db, run_id, status="ok", + counts={"rows_processed": 0, "rows_inserted": 0, + "rows_skipped": 0, "files_copied": 0, + "bytes_copied": 0, "conflicts": 0}, + finished_at=datetime.now(UTC), + metadata_patch={"rollback_result": result}, + ) + 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))