video-gen/video-gen-api/app/tasks/private_portrait_asset_tasks.py

from __future__ import annotations

import asyncio
import logging
from typing import Any

from app.enums.private_portrait import (
    PRIVATE_PORTRAIT_ASSET_POLL_BATCH_SIZE,
    PRIVATE_PORTRAIT_REMOTE_DELETE_RECOVERY_BATCH_SIZE,
    PrivatePortraitEventSource,
    PrivatePortraitEventType,
)
from app.models import async_session
from app.services.operation_log_service import log_operation_error
from app.services.private_portrait.asset_service import (
    DOMAIN,
    delete_asset_remote,
    delete_project_remote,
    poll_due_assets_once,
    recover_remote_deletes_once,
    sync_asset_status,
)
from app.tasks.celery_app import celery_app

logger = logging.getLogger(__name__)


def _run(coro):
    return asyncio.run(coro)


def _retry_countdown(retries: int) -> int:
    return min(300, 30 * (2 ** max(0, retries)))


async def _rollback_and_reraise(db, *, event_type: str, exc: BaseException, detail: dict[str, Any] | None = None, **kwargs: Any):
    await db.rollback()
    log_operation_error(
        domain=DOMAIN,
        event_type=event_type,
        source=PrivatePortraitEventSource.CELERY.value,
        exc=exc,
        detail=detail,
        **kwargs,
    )
    raise exc


@celery_app.task(name="private_portrait.poll_asset_status", queue="gen_private_portrait", bind=True, max_retries=5, default_retry_delay=30)
def poll_private_portrait_asset_status(self, asset_id: str) -> None:
    async def _inner():
        async with async_session() as db:
            try:
                await sync_asset_status(db, user_id=None, asset_id=asset_id)
                await db.commit()
            except Exception as exc:
                logger.exception("poll private portrait asset failed: %s", asset_id)
                await _rollback_and_reraise(
                    db,
                    event_type=PrivatePortraitEventType.ASSET_POLL_FAILED.value,
                    exc=exc,
                    asset_id=asset_id,
                    detail={"celery_task": "private_portrait.poll_asset_status"},
                )
    try:
        _run(_inner())
    except Exception as exc:
        raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))


@celery_app.task(name="private_portrait.sync_due_assets", queue="gen_private_portrait", bind=True, max_retries=3, default_retry_delay=60)
def sync_private_portrait_due_assets(self) -> int:
    async def _inner() -> int:
        async with async_session() as db:
            try:
                count = await poll_due_assets_once(db, limit=PRIVATE_PORTRAIT_ASSET_POLL_BATCH_SIZE)
                await db.commit()
                return count
            except Exception as exc:
                logger.exception("sync private portrait due assets failed")
                await _rollback_and_reraise(
                    db,
                    event_type=PrivatePortraitEventType.SYNC_DUE_ASSETS_FAILED.value,
                    exc=exc,
                    detail={"celery_task": "private_portrait.sync_due_assets"},
                )
    try:
        return _run(_inner())
    except Exception as exc:
        raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))


@celery_app.task(name="private_portrait.delete_asset_remote", queue="gen_private_portrait", bind=True, max_retries=3, default_retry_delay=60)
def delete_private_portrait_asset_remote(self, asset_id: str) -> None:
    async def _inner():
        async with async_session() as db:
            try:
                await delete_asset_remote(db, asset_id=asset_id)
                await db.commit()
            except Exception as exc:
                logger.exception("delete private portrait asset remote failed: %s", asset_id)
                await _rollback_and_reraise(
                    db,
                    event_type=PrivatePortraitEventType.ASSET_DELETE_REMOTE_FAILED.value,
                    exc=exc,
                    asset_id=asset_id,
                    detail={"celery_task": "private_portrait.delete_asset_remote"},
                )
    try:
        _run(_inner())
    except Exception as exc:
        raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))


@celery_app.task(name="private_portrait.delete_project_remote", queue="gen_private_portrait", bind=True, max_retries=3, default_retry_delay=60)
def delete_private_portrait_project_remote(self, project_id: str) -> None:
    async def _inner():
        async with async_session() as db:
            try:
                await delete_project_remote(db, project_id=project_id)
                await db.commit()
            except Exception as exc:
                logger.exception("delete private portrait project remote failed: %s", project_id)
                await _rollback_and_reraise(
                    db,
                    event_type=PrivatePortraitEventType.PROJECT_DELETE_REMOTE_FAILED.value,
                    exc=exc,
                    project_id=project_id,
                    detail={"celery_task": "private_portrait.delete_project_remote"},
                )
    try:
        _run(_inner())
    except Exception as exc:
        raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))


@celery_app.task(name="private_portrait.recover_remote_deletes", queue="gen_private_portrait", bind=True, max_retries=3, default_retry_delay=60)
def recover_private_portrait_remote_deletes(self) -> dict[str, int]:
    async def _inner() -> dict[str, int]:
        async with async_session() as db:
            try:
                result = await recover_remote_deletes_once(db, limit=PRIVATE_PORTRAIT_REMOTE_DELETE_RECOVERY_BATCH_SIZE)
                await db.commit()
                return result
            except Exception as exc:
                logger.exception("recover private portrait remote deletes failed")
                await _rollback_and_reraise(
                    db,
                    event_type=PrivatePortraitEventType.REMOTE_DELETE_RECOVERY_FAILED.value,
                    exc=exc,
                    detail={"celery_task": "private_portrait.recover_remote_deletes"},
                )
    try:
        return _run(_inner())
    except Exception as exc:
        raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))