doris-sync/cron/scheduler.py

#!/usr/bin env python
# -*- coding: utf-8 -*-
"""
调度器
"""

from typing import Callable, Dict, Tuple, Any, Optional, List
from apscheduler.schedulers.blocking import BlockingScheduler
from apscheduler.triggers.cron import CronTrigger


# 任务注册表
JOBS: Dict[str, Dict[str, Any]] = {}


def register(
    name: str,
    func: Callable,
    cron: str,
    enabled: bool = True
) -> None:
    """
    注册任务
    
    Args:
        name: 任务名称
        func: 执行函数
        cron: cron 表达式
        enabled: 是否启用
    """
    JOBS[name] = {
        "func": func,
        "cron": cron,
        "enabled": enabled
    }


def run_job(
    name: str,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    group_id: Optional[str] = None
) -> None:
    """
    运行指定任务
    
    Args:
        name: 任务名称
        start_date: 开始日期
        end_date: 结束日期
        group_id: 集团ID
    """
    job = JOBS.get(name)
    if not job:
        raise ValueError(f"任务 {name} 不存在")
    
    func = job["func"]
    
    if group_id:
        func(start_date, end_date, group_id)
    elif start_date and end_date:
        func(start_date, end_date)
    else:
        func()


def list_jobs() -> List[Tuple[str, str, bool]]:
    """列出所有任务"""
    return [
        (name, job["cron"], job["enabled"])
        for name, job in JOBS.items()
    ]


def start_scheduler() -> None:
    """启动调度器"""
    scheduler = BlockingScheduler(timezone="Asia/Shanghai")
    
    for name, job in JOBS.items():
        if not job["enabled"]:
            print(f"[跳过] {name} (已禁用)")
            continue
        
        trigger = CronTrigger.from_crontab(
            job["cron"],
            timezone="Asia/Shanghai"
        )
        scheduler.add_job(job["func"], trigger=trigger, id=name, name=name)
        print(f"[注册] {name} - {job['cron']}")
    
    print("=" * 50)
    print("调度器启动中...")
    print("按 Ctrl+C 停止")
    print("=" * 50)
    
    try:
        scheduler.start()
    except KeyboardInterrupt:
        print("收到停止信号")
        scheduler.shutdown()