fastapi使用apscheduler和redis任务存储例子

from contextlib import asynccontextmanager
from datetime import datetime
from typing import Optional, Dict, Any
from urllib.parse import quote
from fastapi import FastAPI, HTTPException, Query, Body
from pydantic import BaseModel
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.interval import IntervalTrigger
from apscheduler.triggers.date import DateTrigger
from apscheduler.jobstores.redis import RedisJobStore
from pytz import timezone
from uvicorn import run


# 配置Redis作为任务存储，使用JSON序列化
jobstores = {
    'default': RedisJobStore(
        jobs_key='apscheduler.jobs', 
        run_times_key='apscheduler.run_times',
        host='172.16.xx.xx',
        port=6379,
        db=10,
        password='xxxx@123',
        username='default'
    )
}

# 创建调度器并配置Redis作为任务存储
scheduler = AsyncIOScheduler(
    jobstores=jobstores,
    timezone=timezone('Asia/Shanghai')
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动调度器
    scheduler.start()
    yield
    # 关闭调度器
    scheduler.shutdown()


app = FastAPI(lifespan=lifespan)


class ScheduleRequest(BaseModel):
    """用于创建或更新任务调度的请求模型"""
    # 触发器类型: 'date', 'interval', 或 'cron'
    trigger_type: str
    
    # 触发器参数
    # 对于 'date': {'run_date': '2023-01-01 12:00:00'}
    # 对于 'interval': {'seconds': 30, 'minutes': 0, ...}
    # 对于 'cron': {'hour': '12', 'minute': '30', ...}
    trigger_args: Dict[str, Any]
    
    # 可选的任务ID，用于添加新任务时
    job_id: Optional[str] = None
    
    # 可选的函数名称，用于添加新任务时
    func_name: Optional[str] = None


@app.get('/')
async def home():
    """首页接口，用于检查API是否正常运行"""
    return {'message': 'API is up and running'}


@app.get('/say_hello')
async def say_hello(name: str = 'World'):
    """打招呼接口，返回问候消息"""
    return {'message': f'Hello, {name}!'}


@app.get('/tasks')
async def get_all_tasks():
    """获取所有调度任务的列表
    
    返回所有任务的ID、名称、下次运行时间、触发器类型和活动状态
    """
    # 强制从Redis重新加载任务
    scheduler.get_jobs()
    
    # 获取所有任务
    jobs = scheduler.get_jobs()
    tasks = []
    for job in jobs:
        tasks.append({
            'id': job.id,
            'name': job.name,
            'next_run_time': str(job.next_run_time) if job.next_run_time else None,
            'trigger': str(job.trigger),
            'active': job.next_run_time is not None
        })
    return {'tasks': tasks}


@app.post('/tasks/{job_id}/pause')
async def pause_task(job_id: str):
    """暂停指定的任务
    
    参数:
        job_id: 要暂停的任务ID
    """
    job = scheduler.get_job(job_id)
    if not job:
        raise HTTPException(status_code=404, detail=f"Job with ID {job_id} not found")
    
    scheduler.pause_job(job_id)
    return {'message': f'Job {job_id} paused successfully'}


@app.post('/tasks/{job_id}/resume')
async def resume_task(job_id: str):
    """恢复指定的任务
    
    参数:
        job_id: 要恢复的任务ID
    """
    job = scheduler.get_job(job_id)
    if not job:
        raise HTTPException(status_code=404, detail=f"Job with ID {job_id} not found")
    
    scheduler.resume_job(job_id)
    return {'message': f'Job {job_id} resumed successfully'}


@app.post('/tasks/pause-all')
async def pause_all_tasks():
    """暂停所有任务"""
    scheduler.pause()
    return {'message': 'All jobs paused successfully'}


@app.post('/tasks/resume-all')
async def resume_all_tasks():
    """恢复所有任务"""
    scheduler.resume()
    return {'message': 'All jobs resumed successfully'}


@app.delete('/tasks/{job_id}')
async def remove_task(job_id: str):
    """删除指定的任务
    
    参数:
        job_id: 要删除的任务ID
    """
    job = scheduler.get_job(job_id)
    if not job:
        raise HTTPException(status_code=404, detail=f"Job with ID {job_id} not found")
    
    scheduler.remove_job(job_id)
    return {'message': f'Job {job_id} removed successfully'}


@app.put('/tasks/{job_id}/schedule')
async def update_task_schedule(job_id: str, schedule_data: ScheduleRequest):
    """更新任务的调度计划
    
    参数:
        job_id: 要更新的任务ID
        schedule_data: 包含触发器类型和参数的请求体
    """
    job = scheduler.get_job(job_id)
    if not job:
        raise HTTPException(status_code=404, detail=f"Job with ID {job_id} not found")
    
    try:
        trigger = create_trigger(schedule_data.trigger_type, schedule_data.trigger_args)
        scheduler.reschedule_job(job_id, trigger=trigger)
        return {
            'message': f'Job {job_id} schedule updated successfully',
            'next_run_time': str(scheduler.get_job(job_id).next_run_time)
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.post('/tasks/schedule')
async def add_scheduled_task(schedule_data: ScheduleRequest):
    """添加一个新的调度任务
    
    参数:
        schedule_data: 包含触发器类型、参数、函数名和可选任务ID的请求体
    """
    if not schedule_data.func_name:
        raise HTTPException(status_code=400, detail="Function name is required")
    
    available_functions = {
        'add_info_file': add_info_file,
        'add_info_file_1': add_info_file_1
    }
    
    if schedule_data.func_name not in available_functions:
        raise HTTPException(status_code=400, detail=f"Function {schedule_data.func_name} not found")
    
    try:
        trigger = create_trigger(schedule_data.trigger_type, schedule_data.trigger_args)
        job = scheduler.add_job(
            available_functions[schedule_data.func_name],
            trigger=trigger,
            id=schedule_data.job_id
        )
        
        return {
            'message': 'Job added successfully',
            'job_id': job.id,
            'next_run_time': str(job.next_run_time)
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


def create_trigger(trigger_type: str, trigger_args: Dict[str, Any]):
    """创建触发器对象
    
    参数:
        trigger_type: 触发器类型，可以是 'date', 'interval' 或 'cron'
        trigger_args: 触发器参数
        
    返回:
        对应类型的触发器对象
    """
    if trigger_type == 'date':
        # 日期触发器 - 在特定时间点执行一次
        return DateTrigger(**trigger_args, timezone=timezone('Asia/Shanghai'))
    elif trigger_type == 'interval':
        # 间隔触发器 - 每隔一段时间执行一次
        return IntervalTrigger(**trigger_args, timezone=timezone('Asia/Shanghai'))
    elif trigger_type == 'cron':
        # Cron触发器 - 按cron表达式执行
        return CronTrigger(**trigger_args, timezone=timezone('Asia/Shanghai'))
    else:
        raise ValueError(f"Unsupported trigger type: {trigger_type}")


@app.put('/tasks/{job_id}/function')
async def update_task_function(job_id: str, func_name: str = Query(...)):
    """更新任务的执行函数
    
    参数:
        job_id: 要更新的任务ID
        func_name: 新的函数名称
    """
    job = scheduler.get_job(job_id)
    if not job:
        raise HTTPException(status_code=404, detail=f"Job with ID {job_id} not found")
    
    available_functions = {
        'add_info_file': add_info_file,
        'add_info_file_1': add_info_file_1
    }
    
    if func_name not in available_functions:
        raise HTTPException(status_code=400, detail=f"Function {func_name} not found")
    
    try:
        scheduler.modify_job(job_id, func=available_functions[func_name])
        return {'message': f'Job {job_id} function updated to {func_name} successfully'}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.post('/recreate-default-job')
async def recreate_default_job():
    """重新创建默认任务
    
    如果默认任务丢失或无法正常工作，可以使用此接口重新创建
    """
    try:
        # 先尝试删除现有的默认任务
        try:
            scheduler.remove_job('defaultjob')  # 注意：使用正确的任务ID
        except Exception:
            # 如果任务不存在，忽略错误
            pass
        
        # 创建新的默认任务
        job = scheduler.add_job(
            add_info_file,
            'interval',
            seconds=10,
            id='defaultjob',  # 注意：使用正确的任务ID
            replace_existing=True
        )
        
        return {
            'message': 'Default job recreated successfully',
            'job_id': job.id,
            'next_run_time': str(job.next_run_time)
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error recreating default job: {str(e)}")



# 使用Redis存储的调度器添加任务
@scheduler.scheduled_job('interval', seconds=10, id='defaultjob')
async def add_info_file():
    """默认的定时任务，每10秒执行一次，将当前时间写入文件"""
    print('add_info_file')
    file_path = r"D:\1fkl\all_test\for_work\dir_2025\month5\aps\storage\1.txt"
    with open(file_path, 'a', encoding='utf-8') as f:
        f.write(f'写入时间-{datetime.now().strftime("%Y-%m-%d %H:%M:%S")} \n')
        
async def add_info_file_1():
    """备用的定时任务，将当前时间以不同格式写入文件"""
    print('add_info_file_1')
    file_path = r"D:\1fkl\all_test\for_work\dir_2025\month5\aps\storage\1.txt"
    with open(file_path, 'a', encoding='utf-8') as f:
        f.write(f'插入时间-{datetime.now().strftime("%Y-%m-%d %H:%M:%S")} \n')


if __name__ == '__main__':
    run(app, host="127.0.0.1", port=8144)