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

from socket import gethostname
from typing import Any

from fastapi import FastAPI

from app.core.config import get_settings
from app.core.constants import ActorValue
from app.modules.feishu.constants import FeishuReceiveIdType
from app.modules.observability.constants import HeartbeatComponent


def attach_scheduler(app: FastAPI) -> None:
    """Attach optional APScheduler jobs to the FastAPI application."""

    settings = get_settings()
    if not settings.scheduler_enabled:
        return

    scheduler = create_scheduler(app)

    @app.on_event("startup")
    def start_scheduler() -> None:
        scheduler.start()

    @app.on_event("shutdown")
    def stop_scheduler() -> None:
        scheduler.shutdown(wait=False)


def create_scheduler(app: FastAPI | None = None) -> Any:
    """Create the V2/V3 scheduler without requiring a FastAPI process."""

    from apscheduler.schedulers.background import BackgroundScheduler

    from app.core.database import SessionLocal
    from app.core.task_queue import (
        enqueue_daily_brief_push,
        enqueue_event_dispatch,
        enqueue_legacy_project_sync,
        enqueue_legacy_task_sync,
        enqueue_project_weekly_push,
    )
    from app.modules.observability.service import ObservabilityService
    from app.modules.reports.service import ReportService

    settings = get_settings()
    scheduler = BackgroundScheduler(timezone="Asia/Shanghai")

    def run_daily_brief() -> None:
        db = SessionLocal()
        try:
            report = ReportService(db).daily_brief()
            _set_state(app, "last_daily_brief", report)
            if (
                settings.feishu_app_id
                and settings.feishu_app_secret
                and settings.feishu_default_chat_id
            ):
                if settings.task_queue_enabled:
                    dispatch = enqueue_daily_brief_push(
                        receive_id=settings.feishu_default_chat_id,
                        receive_id_type=FeishuReceiveIdType.CHAT_ID,
                        actor=ActorValue.SCHEDULER,
                    )
                    _set_state(app, "last_daily_brief_dispatch", dispatch)
                    return
                ReportService(db).push_report(
                    report,
                    receive_id=settings.feishu_default_chat_id,
                    receive_id_type=FeishuReceiveIdType.CHAT_ID,
                    actor=ActorValue.SCHEDULER,
                )
        finally:
            db.close()

    def run_project_weekly() -> None:
        db = SessionLocal()
        try:
            report = ReportService(db).project_weekly()
            _set_state(app, "last_project_weekly", report)
            if (
                settings.feishu_app_id
                and settings.feishu_app_secret
                and settings.feishu_default_chat_id
            ):
                if settings.task_queue_enabled:
                    dispatch = enqueue_project_weekly_push(
                        receive_id=settings.feishu_default_chat_id,
                        receive_id_type=FeishuReceiveIdType.CHAT_ID,
                        actor=ActorValue.SCHEDULER,
                    )
                    _set_state(app, "last_project_weekly_dispatch", dispatch)
                    return
                ReportService(db).push_report(
                    report,
                    receive_id=settings.feishu_default_chat_id,
                    receive_id_type=FeishuReceiveIdType.CHAT_ID,
                    actor=ActorValue.SCHEDULER,
                )
        finally:
            db.close()

    def run_legacy_project_sync() -> None:
        dispatch = enqueue_legacy_project_sync(actor=ActorValue.SCHEDULER)
        _set_state(app, "last_legacy_project_sync_dispatch", dispatch)

    def run_legacy_task_sync() -> None:
        dispatch = enqueue_legacy_task_sync(actor=ActorValue.SCHEDULER)
        _set_state(app, "last_legacy_task_sync_dispatch", dispatch)

    def run_event_dispatch() -> None:
        dispatch = enqueue_event_dispatch(
            limit=settings.event_dispatch_batch_size,
            actor=ActorValue.SCHEDULER,
        )
        _set_state(app, "last_event_dispatch", dispatch)

    def record_scheduler_heartbeat() -> None:
        db = SessionLocal()
        try:
            heartbeat = ObservabilityService(db).record_heartbeat(
                component=HeartbeatComponent.SCHEDULER,
                instance_id=gethostname(),
                actor=ActorValue.SCHEDULER,
            )
            _set_state(app, "last_scheduler_heartbeat", heartbeat)
        finally:
            db.close()

    scheduler.add_job(
        run_daily_brief,
        trigger="cron",
        hour=settings.daily_brief_cron_hour,
        minute=settings.daily_brief_cron_minute,
        id="daily_brief_push",
        replace_existing=True,
    )
    scheduler.add_job(
        run_project_weekly,
        trigger="cron",
        day_of_week=settings.weekly_project_report_day_of_week,
        hour=settings.weekly_project_report_cron_hour,
        minute=settings.weekly_project_report_cron_minute,
        id="project_weekly_push",
        replace_existing=True,
    )
    if settings.event_dispatch_enabled:
        scheduler.add_job(
            run_event_dispatch,
            trigger="cron",
            minute=settings.event_dispatch_cron_minute,
            id="event_dispatch",
            replace_existing=True,
        )
    scheduler.add_job(
        record_scheduler_heartbeat,
        trigger="interval",
        seconds=settings.heartbeat_interval_seconds,
        id="scheduler_heartbeat",
        replace_existing=True,
    )
    if settings.legacy_sync_enabled and settings.legacy_project_query:
        scheduler.add_job(
            run_legacy_project_sync,
            trigger="cron",
            hour=settings.legacy_project_sync_cron_hour,
            minute=settings.legacy_project_sync_cron_minute,
            id="legacy_project_sync",
            replace_existing=True,
        )
    if settings.legacy_sync_enabled and settings.legacy_task_query:
        scheduler.add_job(
            run_legacy_task_sync,
            trigger="cron",
            hour=settings.legacy_task_sync_cron_hour,
            minute=settings.legacy_task_sync_cron_minute,
            id="legacy_task_sync",
            replace_existing=True,
        )
    return scheduler


def _set_state(app: FastAPI | None, key: str, value: Any) -> None:
    if app is not None:
        setattr(app.state, key, value)