169 lines
5.5 KiB
Python
169 lines
5.5 KiB
Python
from apscheduler.schedulers.background import BackgroundScheduler
|
|
from apscheduler.executors.pool import ThreadPoolExecutor
|
|
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
|
|
from apscheduler.events import EVENT_JOB_EXECUTED, EVENT_JOB_ERROR
|
|
from ..core.config import settings
|
|
from ..core.database import SessionLocal
|
|
from ..core.logging_config import get_logger
|
|
from ..models.account import Account
|
|
|
|
# 获取日志记录器
|
|
logger = get_logger(__name__)
|
|
|
|
class TaskScheduler:
|
|
def __init__(self):
|
|
# 配置作业存储
|
|
jobstores = {
|
|
'default': SQLAlchemyJobStore(url=settings.DATABASE_URL)
|
|
}
|
|
|
|
# 配置执行器
|
|
executors = {
|
|
'default': ThreadPoolExecutor(20)
|
|
}
|
|
|
|
# 作业默认配置
|
|
job_defaults = {
|
|
'coalesce': False,
|
|
'max_instances': 3
|
|
}
|
|
|
|
# 创建调度器
|
|
self.scheduler = BackgroundScheduler(
|
|
jobstores=jobstores,
|
|
executors=executors,
|
|
job_defaults=job_defaults,
|
|
timezone='Asia/Shanghai'
|
|
)
|
|
|
|
# 添加事件监听器
|
|
self.scheduler.add_listener(self._job_listener, EVENT_JOB_EXECUTED | EVENT_JOB_ERROR)
|
|
|
|
def _job_listener(self, event):
|
|
"""作业执行监听器"""
|
|
if event.exception:
|
|
logger.error(f"Job {event.job_id} crashed: {event.exception}")
|
|
else:
|
|
logger.info(f"Job {event.job_id} executed successfully")
|
|
|
|
def start(self):
|
|
"""启动调度器"""
|
|
if not self.scheduler.running:
|
|
self.scheduler.start()
|
|
logger.info("定时任务调度器已启动")
|
|
|
|
# 启动时自动添加系统定时任务
|
|
self._setup_system_tasks()
|
|
|
|
def shutdown(self):
|
|
"""关闭调度器"""
|
|
if self.scheduler.running:
|
|
self.scheduler.shutdown()
|
|
logger.info("定时任务调度器已关闭")
|
|
|
|
def _setup_system_tasks(self):
|
|
"""设置系统定时任务"""
|
|
try:
|
|
# 检查是否已存在每日重置任务
|
|
existing_job = self.scheduler.get_job("daily_reset_today_updated")
|
|
if not existing_job:
|
|
# 添加每天午夜12点重置today_updated字段的任务
|
|
self.scheduler.add_job(
|
|
reset_today_updated_task,
|
|
'cron',
|
|
id='daily_reset_today_updated',
|
|
hour=0,
|
|
minute=0,
|
|
second=0,
|
|
name='每日重置账号更新状态'
|
|
)
|
|
logger.info("系统定时任务：每日重置账号更新状态已添加")
|
|
except Exception as e:
|
|
logger.error(f"设置系统定时任务失败: {e}")
|
|
|
|
def add_cron_job(self, func, job_id: str, **kwargs):
|
|
"""添加cron定时任务"""
|
|
return self.scheduler.add_job(func, 'cron', id=job_id, **kwargs)
|
|
|
|
def add_interval_job(self, func, job_id: str, **kwargs):
|
|
"""添加间隔执行任务"""
|
|
return self.scheduler.add_job(func, 'interval', id=job_id, **kwargs)
|
|
|
|
def add_date_job(self, func, job_id: str, **kwargs):
|
|
"""添加指定时间执行任务"""
|
|
return self.scheduler.add_job(func, 'date', id=job_id, **kwargs)
|
|
|
|
def remove_job(self, job_id: str):
|
|
"""移除任务"""
|
|
try:
|
|
self.scheduler.remove_job(job_id)
|
|
logger.info(f"任务 {job_id} 已移除")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"移除任务失败: {e}")
|
|
return False
|
|
|
|
def get_job(self, job_id: str):
|
|
"""获取任务"""
|
|
return self.scheduler.get_job(job_id)
|
|
|
|
def get_jobs(self):
|
|
"""获取所有任务"""
|
|
return self.scheduler.get_jobs()
|
|
|
|
def pause_job(self, job_id: str):
|
|
"""暂停任务"""
|
|
try:
|
|
self.scheduler.pause_job(job_id)
|
|
logger.info(f"任务 {job_id} 已暂停")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"暂停任务失败: {e}")
|
|
return False
|
|
|
|
def resume_job(self, job_id: str):
|
|
"""恢复任务"""
|
|
try:
|
|
self.scheduler.resume_job(job_id)
|
|
logger.info(f"任务 {job_id} 已恢复")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"恢复任务失败: {e}")
|
|
return False
|
|
|
|
# 全局调度器实例
|
|
task_scheduler = TaskScheduler()
|
|
|
|
# 系统定时任务函数
|
|
def reset_today_updated_task():
|
|
"""每日重置账号today_updated字段为0的任务"""
|
|
db = SessionLocal()
|
|
try:
|
|
logger.info("开始执行每日重置账号更新状态任务")
|
|
|
|
# 更新所有账号的today_updated字段为0
|
|
updated_count = db.query(Account).update({Account.today_updated: 0})
|
|
db.commit()
|
|
|
|
logger.info(f"每日重置任务完成，已重置 {updated_count} 个账号的today_updated字段")
|
|
return f"成功重置 {updated_count} 个账号的更新状态"
|
|
|
|
except Exception as e:
|
|
db.rollback()
|
|
logger.error(f"每日重置任务执行失败: {e}")
|
|
raise e
|
|
finally:
|
|
db.close()
|
|
|
|
# 示例定时任务函数
|
|
def example_task():
|
|
"""示例定时任务"""
|
|
logger.info("执行示例定时任务")
|
|
# 这里可以添加具体的业务逻辑
|
|
return "任务执行完成"
|
|
|
|
def database_cleanup_task():
|
|
"""数据库清理任务示例"""
|
|
logger.info("执行数据库清理任务")
|
|
# 这里可以添加数据库清理逻辑
|
|
return "数据库清理完成" |