company-ai-platform/app/core/task_queue.py

from collections.abc import Callable
from typing import Any

from app.core.config import get_settings
from app.modules.feishu.constants import FeishuReceiveIdType

TASK_PUSH_DAILY_BRIEF = "reports.push_daily_brief"
TASK_PUSH_PROJECT_WEEKLY = "reports.push_project_weekly"
TASK_GENERATE_RISK_EVENTS = "risks.generate_events"
TASK_SYNC_LEGACY_PROJECTS = "legacy.sync_projects"
TASK_SYNC_LEGACY_TASKS = "legacy.sync_tasks"
TASK_DISPATCH_PENDING_EVENTS = "events.dispatch_pending"


def dispatch_task(
    task_name: str,
    kwargs: dict[str, Any],
    inline: Callable[[], Any],
) -> dict[str, Any]:
    """Dispatch a task through Celery when enabled, otherwise run inline."""

    settings = get_settings()
    if settings.task_queue_enabled:
        from app.tasks import celery_app

        async_result = celery_app.signature(task_name, kwargs=kwargs).apply_async()
        return {
            "queued": True,
            "mode": "celery",
            "task_name": task_name,
            "task_id": async_result.id,
        }
    return {
        "queued": False,
        "mode": "inline",
        "task_name": task_name,
        "result": inline(),
    }


def enqueue_daily_brief_push(
    receive_id: str | None = None,
    receive_id_type: str = FeishuReceiveIdType.CHAT_ID,
    actor: str = "scheduler",
) -> dict[str, Any]:
    settings = get_settings()
    if settings.task_queue_enabled:
        from app.core.database import SessionLocal
        from app.modules.reports.constants import ReportPushStatus, ReportTitle, ReportType
        from app.modules.reports.service import ReportService
        from app.tasks import celery_app

        db = SessionLocal()
        try:
            push_run = ReportService(db).create_push_run(
                report_type=ReportType.DAILY,
                title=ReportTitle.DAILY_BRIEF,
                receive_id=receive_id,
                receive_id_type=receive_id_type,
                actor=actor,
                status=ReportPushStatus.QUEUED,
            )
            async_result = celery_app.signature(
                TASK_PUSH_DAILY_BRIEF,
                kwargs={
                    "receive_id": receive_id,
                    "receive_id_type": receive_id_type,
                    "actor": actor,
                    "push_run_code": push_run.code,
                },
            ).apply_async()
            ReportService(db).update_push_run(
                push_run.code,
                ReportPushStatus.QUEUED,
                task_id=async_result.id,
            )
        finally:
            db.close()
        return {
            "queued": True,
            "mode": "celery",
            "task_name": TASK_PUSH_DAILY_BRIEF,
            "task_id": async_result.id,
            "push_run_code": push_run.code,
        }

    def inline() -> Any:
        from app.core.database import SessionLocal
        from app.modules.reports.service import ReportService

        db = SessionLocal()
        try:
            report = ReportService(db).daily_brief()
            return ReportService(db).push_report(report, receive_id, receive_id_type, actor)
        finally:
            db.close()

    return dispatch_task(
        TASK_PUSH_DAILY_BRIEF,
        {
            "receive_id": receive_id,
            "receive_id_type": receive_id_type,
            "actor": actor,
        },
        inline,
    )


def enqueue_project_weekly_push(
    receive_id: str | None = None,
    receive_id_type: str = FeishuReceiveIdType.CHAT_ID,
    actor: str = "scheduler",
) -> dict[str, Any]:
    settings = get_settings()
    if settings.task_queue_enabled:
        from app.core.database import SessionLocal
        from app.modules.reports.constants import ReportPushStatus, ReportTitle, ReportType
        from app.modules.reports.service import ReportService
        from app.tasks import celery_app

        db = SessionLocal()
        try:
            push_run = ReportService(db).create_push_run(
                report_type=ReportType.WEEKLY,
                title=ReportTitle.PROJECT_WEEKLY,
                receive_id=receive_id,
                receive_id_type=receive_id_type,
                actor=actor,
                status=ReportPushStatus.QUEUED,
            )
            async_result = celery_app.signature(
                TASK_PUSH_PROJECT_WEEKLY,
                kwargs={
                    "receive_id": receive_id,
                    "receive_id_type": receive_id_type,
                    "actor": actor,
                    "push_run_code": push_run.code,
                },
            ).apply_async()
            ReportService(db).update_push_run(
                push_run.code,
                ReportPushStatus.QUEUED,
                task_id=async_result.id,
            )
        finally:
            db.close()
        return {
            "queued": True,
            "mode": "celery",
            "task_name": TASK_PUSH_PROJECT_WEEKLY,
            "task_id": async_result.id,
            "push_run_code": push_run.code,
        }

    def inline() -> Any:
        from app.core.database import SessionLocal
        from app.modules.reports.service import ReportService

        db = SessionLocal()
        try:
            report = ReportService(db).project_weekly()
            return ReportService(db).push_report(report, receive_id, receive_id_type, actor)
        finally:
            db.close()

    return dispatch_task(
        TASK_PUSH_PROJECT_WEEKLY,
        {
            "receive_id": receive_id,
            "receive_id_type": receive_id_type,
            "actor": actor,
        },
        inline,
    )


def enqueue_risk_event_generation(actor: str = "scheduler") -> dict[str, Any]:
    def inline() -> Any:
        from app.core.database import SessionLocal
        from app.modules.risk.service import RiskService

        db = SessionLocal()
        try:
            return RiskService(db).generate_events(actor=actor)
        finally:
            db.close()

    return dispatch_task(
        TASK_GENERATE_RISK_EVENTS,
        {"actor": actor},
        inline,
    )


def enqueue_event_dispatch(
    limit: int | None = None,
    actor: str = "scheduler",
) -> dict[str, Any]:
    def inline() -> Any:
        from app.core.config import get_settings
        from app.core.database import SessionLocal
        from app.modules.events.service import EventService

        db = SessionLocal()
        try:
            return EventService(db).dispatch_pending(
                limit=limit or get_settings().event_dispatch_batch_size,
                worker_id=actor,
            )
        finally:
            db.close()

    return dispatch_task(
        TASK_DISPATCH_PENDING_EVENTS,
        {"limit": limit, "actor": actor},
        inline,
    )


def enqueue_legacy_project_sync(
    source_query: str | None = None,
    source_query_name: str | None = None,
    field_map: dict[str, str] | None = None,
    limit: int = 100,
    dry_run: bool = False,
    actor: str = "scheduler",
) -> dict[str, Any]:
    def inline() -> Any:
        from app.core.database import SessionLocal
        from app.modules.legacy_mysql.service import LegacyMySQLService

        db = SessionLocal()
        try:
            return LegacyMySQLService(db).sync_projects(
                source_query=source_query,
                source_query_name=source_query_name,
                field_map=field_map or {},
                limit=limit,
                dry_run=dry_run,
                actor=actor,
            )
        finally:
            db.close()

    return dispatch_task(
        TASK_SYNC_LEGACY_PROJECTS,
        {
            "source_query": source_query,
            "source_query_name": source_query_name,
            "field_map": field_map or {},
            "limit": limit,
            "dry_run": dry_run,
            "actor": actor,
        },
        inline,
    )


def enqueue_legacy_task_sync(
    source_query: str | None = None,
    source_query_name: str | None = None,
    field_map: dict[str, str] | None = None,
    limit: int = 100,
    dry_run: bool = False,
    actor: str = "scheduler",
) -> dict[str, Any]:
    def inline() -> Any:
        from app.core.database import SessionLocal
        from app.modules.legacy_mysql.service import LegacyMySQLService

        db = SessionLocal()
        try:
            return LegacyMySQLService(db).sync_tasks(
                source_query=source_query,
                source_query_name=source_query_name,
                field_map=field_map or {},
                limit=limit,
                dry_run=dry_run,
                actor=actor,
            )
        finally:
            db.close()

    return dispatch_task(
        TASK_SYNC_LEGACY_TASKS,
        {
            "source_query": source_query,
            "source_query_name": source_query_name,
            "field_map": field_map or {},
            "limit": limit,
            "dry_run": dry_run,
            "actor": actor,
        },
        inline,
    )