""" 服务层 (Service) 定义 SchedulerManager 类作为定时任务服务的公共 API 入口。 它负责编排业务逻辑,并调用 Repository 和 Adapter 层来完成具体工作。 """ from collections.abc import Callable, Coroutine from datetime import datetime from typing import Any, ClassVar import nonebot from pydantic import BaseModel from zhenxun.configs.config import Config from zhenxun.models.schedule_info import ScheduleInfo from zhenxun.services.log import logger from .adapter import APSchedulerAdapter from .job import _execute_job from .repository import ScheduleRepository from .targeter import ScheduleTargeter class SchedulerManager: ALL_GROUPS: ClassVar[str] = "__ALL_GROUPS__" _registered_tasks: ClassVar[ dict[str, dict[str, Callable | type[BaseModel] | None]] ] = {} _declared_tasks: ClassVar[list[dict[str, Any]]] = [] _running_tasks: ClassVar[set] = set() def target(self, **filters: Any) -> ScheduleTargeter: """ 创建目标选择器以执行批量操作 参数: **filters: 过滤条件,支持plugin_name、group_id、bot_id等字段。 返回: ScheduleTargeter: 目标选择器对象,可用于批量操作。 """ return ScheduleTargeter(self, **filters) def task( self, trigger: str, group_id: str | None = None, bot_id: str | None = None, **trigger_kwargs, ): """ 声明式定时任务装饰器 参数: trigger: 触发器类型,如'cron'、'interval'等。 group_id: 目标群组ID,None表示全局任务。 bot_id: 目标Bot ID,None表示使用默认Bot。 **trigger_kwargs: 触发器配置参数。 返回: Callable: 装饰器函数。 """ def decorator(func: Callable[..., Coroutine]) -> Callable[..., Coroutine]: try: plugin = nonebot.get_plugin_by_module_name(func.__module__) if not plugin: raise ValueError(f"函数 {func.__name__} 不在任何已加载的插件中。") plugin_name = plugin.name task_declaration = { "plugin_name": plugin_name, "func": func, "group_id": group_id, "bot_id": bot_id, "trigger_type": trigger, "trigger_config": trigger_kwargs, "job_kwargs": {}, } self._declared_tasks.append(task_declaration) logger.debug( f"发现声明式定时任务 '{plugin_name}',将在启动时进行注册。" ) except Exception as e: logger.error(f"注册声明式定时任务失败: {func.__name__}, 错误: {e}") return func return decorator def register( self, plugin_name: str, params_model: type[BaseModel] | None = None ) -> Callable: """ 注册可调度的任务函数 参数: plugin_name: 插件名称,用于标识任务。 params_model: 参数验证模型,继承自BaseModel的类。 返回: Callable: 装饰器函数。 """ def decorator(func: Callable[..., Coroutine]) -> Callable[..., Coroutine]: if plugin_name in self._registered_tasks: logger.warning(f"插件 '{plugin_name}' 的定时任务已被重复注册。") self._registered_tasks[plugin_name] = { "func": func, "model": params_model, } model_name = params_model.__name__ if params_model else "无" logger.debug( f"插件 '{plugin_name}' 的定时任务已注册,参数模型: {model_name}" ) return func return decorator def get_registered_plugins(self) -> list[str]: """ 获取已注册插件列表 返回: list[str]: 已注册的插件名称列表。 """ return list(self._registered_tasks.keys()) async def add_daily_task( self, plugin_name: str, group_id: str | None, hour: int, minute: int, second: int = 0, job_kwargs: dict | None = None, bot_id: str | None = None, ) -> "ScheduleInfo | None": """ 添加每日定时任务 参数: plugin_name: 插件名称。 group_id: 目标群组ID,None表示全局任务。 hour: 执行小时(0-23)。 minute: 执行分钟(0-59)。 second: 执行秒数(0-59),默认为0。 job_kwargs: 任务参数字典。 bot_id: 目标Bot ID,None表示使用默认Bot。 返回: ScheduleInfo | None: 创建的任务信息,失败时返回None。 """ trigger_config = { "hour": hour, "minute": minute, "second": second, "timezone": Config.get_config("SchedulerManager", "SCHEDULER_TIMEZONE"), } return await self.add_schedule( plugin_name, group_id, "cron", trigger_config, job_kwargs=job_kwargs, bot_id=bot_id, ) async def add_interval_task( self, plugin_name: str, group_id: str | None, *, weeks: int = 0, days: int = 0, hours: int = 0, minutes: int = 0, seconds: int = 0, start_date: str | datetime | None = None, job_kwargs: dict | None = None, bot_id: str | None = None, ) -> "ScheduleInfo | None": """添加间隔性定时任务""" trigger_config = { "weeks": weeks, "days": days, "hours": hours, "minutes": minutes, "seconds": seconds, "start_date": start_date, } trigger_config = {k: v for k, v in trigger_config.items() if v} return await self.add_schedule( plugin_name, group_id, "interval", trigger_config, job_kwargs=job_kwargs, bot_id=bot_id, ) def _validate_and_prepare_kwargs( self, plugin_name: str, job_kwargs: dict | None ) -> tuple[bool, str | dict]: """验证并准备任务参数,应用默认值""" from pydantic import ValidationError task_meta = self._registered_tasks.get(plugin_name) if not task_meta: return False, f"插件 '{plugin_name}' 未注册。" params_model = task_meta.get("model") job_kwargs = job_kwargs if job_kwargs is not None else {} if not params_model: if job_kwargs: logger.warning( f"插件 '{plugin_name}' 未定义参数模型,但收到了参数: {job_kwargs}" ) return True, job_kwargs if not (isinstance(params_model, type) and issubclass(params_model, BaseModel)): logger.error(f"插件 '{plugin_name}' 的参数模型不是有效的 BaseModel 类") return False, f"插件 '{plugin_name}' 的参数模型配置错误" try: model_validate = getattr(params_model, "model_validate", None) if not model_validate: return False, f"插件 '{plugin_name}' 的参数模型不支持验证" validated_model = model_validate(job_kwargs) model_dump = getattr(validated_model, "model_dump", None) if not model_dump: return False, f"插件 '{plugin_name}' 的参数模型不支持导出" return True, model_dump() except ValidationError as e: errors = [f" - {err['loc'][0]}: {err['msg']}" for err in e.errors()] error_str = "\n".join(errors) msg = f"插件 '{plugin_name}' 的任务参数验证失败:\n{error_str}" return False, msg async def add_schedule( self, plugin_name: str, group_id: str | None, trigger_type: str, trigger_config: dict, job_kwargs: dict | None = None, bot_id: str | None = None, ) -> "ScheduleInfo | None": """ 添加定时任务(通用方法) 参数: plugin_name: 插件名称。 group_id: 目标群组ID,None表示全局任务。 trigger_type: 触发器类型,如'cron'、'interval'等。 trigger_config: 触发器配置字典。 job_kwargs: 任务参数字典。 bot_id: 目标Bot ID,None表示使用默认Bot。 返回: ScheduleInfo | None: 创建的任务信息,失败时返回None。 """ if plugin_name not in self._registered_tasks: logger.error(f"插件 '{plugin_name}' 没有注册可用的定时任务。") return None is_valid, result = self._validate_and_prepare_kwargs(plugin_name, job_kwargs) if not is_valid: logger.error(f"任务参数校验失败: {result}") return None search_kwargs = {"plugin_name": plugin_name, "group_id": group_id} if bot_id and group_id == self.ALL_GROUPS: search_kwargs["bot_id"] = bot_id else: search_kwargs["bot_id__isnull"] = True defaults = { "trigger_type": trigger_type, "trigger_config": trigger_config, "job_kwargs": result, "is_enabled": True, } schedule, created = await ScheduleRepository.update_or_create( defaults, **search_kwargs ) APSchedulerAdapter.add_or_reschedule_job(schedule) action = "设置" if created else "更新" logger.info( f"已成功{action}插件 '{plugin_name}' 的定时任务 (ID: {schedule.id})。" ) return schedule async def get_all_schedules(self) -> list[ScheduleInfo]: """ 获取所有定时任务信息 """ return await self.get_schedules() async def get_schedules( self, plugin_name: str | None = None, group_id: str | None = None, bot_id: str | None = None, ) -> list[ScheduleInfo]: """ 根据条件获取定时任务列表 参数: plugin_name: 插件名称,None表示不限制。 group_id: 群组ID,None表示不限制。 bot_id: Bot ID,None表示不限制。 返回: list[ScheduleInfo]: 符合条件的任务信息列表。 """ return await ScheduleRepository.query_schedules( plugin_name=plugin_name, group_id=group_id, bot_id=bot_id ) async def update_schedule( self, schedule_id: int, trigger_type: str | None = None, trigger_config: dict | None = None, job_kwargs: dict | None = None, ) -> tuple[bool, str]: """ 更新定时任务配置 参数: schedule_id: 任务ID。 trigger_type: 新的触发器类型,None表示不更新。 trigger_config: 新的触发器配置,None表示不更新。 job_kwargs: 新的任务参数,None表示不更新。 返回: tuple[bool, str]: (是否成功, 结果消息)。 """ schedule = await ScheduleRepository.get_by_id(schedule_id) if not schedule: return False, f"未找到 ID 为 {schedule_id} 的任务。" updated_fields = [] if trigger_config is not None: schedule.trigger_config = trigger_config updated_fields.append("trigger_config") if trigger_type is not None and schedule.trigger_type != trigger_type: schedule.trigger_type = trigger_type updated_fields.append("trigger_type") if job_kwargs is not None: existing_kwargs = ( schedule.job_kwargs.copy() if isinstance(schedule.job_kwargs, dict) else {} ) existing_kwargs.update(job_kwargs) is_valid, result = self._validate_and_prepare_kwargs( schedule.plugin_name, existing_kwargs ) if not is_valid: return False, str(result) assert isinstance(result, dict), "验证成功时 result 应该是字典类型" schedule.job_kwargs = result updated_fields.append("job_kwargs") if not updated_fields: return True, "没有任何需要更新的配置。" await ScheduleRepository.save(schedule, update_fields=updated_fields) APSchedulerAdapter.add_or_reschedule_job(schedule) return True, f"成功更新了任务 ID: {schedule_id} 的配置。" async def get_schedule_status(self, schedule_id: int) -> dict | None: """获取定时任务的详细状态信息""" schedule = await ScheduleRepository.get_by_id(schedule_id) if not schedule: return None status_from_scheduler = APSchedulerAdapter.get_job_status(schedule.id) status_text = ( "运行中" if schedule_id in self._running_tasks else ("启用" if schedule.is_enabled else "暂停") ) return { "id": schedule.id, "bot_id": schedule.bot_id, "plugin_name": schedule.plugin_name, "group_id": schedule.group_id, "is_enabled": status_text, "trigger_type": schedule.trigger_type, "trigger_config": schedule.trigger_config, "job_kwargs": schedule.job_kwargs, **status_from_scheduler, } async def pause_schedule(self, schedule_id: int) -> tuple[bool, str]: """暂停指定的定时任务""" schedule = await ScheduleRepository.get_by_id(schedule_id) if not schedule or not schedule.is_enabled: return False, "任务不存在或已暂停。" schedule.is_enabled = False await ScheduleRepository.save(schedule, update_fields=["is_enabled"]) APSchedulerAdapter.pause_job(schedule_id) return True, f"已暂停任务 (ID: {schedule.id})。" async def resume_schedule(self, schedule_id: int) -> tuple[bool, str]: """恢复指定的定时任务""" schedule = await ScheduleRepository.get_by_id(schedule_id) if not schedule or schedule.is_enabled: return False, "任务不存在或已启用。" schedule.is_enabled = True await ScheduleRepository.save(schedule, update_fields=["is_enabled"]) APSchedulerAdapter.resume_job(schedule_id) return True, f"已恢复任务 (ID: {schedule.id})。" async def trigger_now(self, schedule_id: int) -> tuple[bool, str]: """立即手动触发指定的定时任务""" schedule = await ScheduleRepository.get_by_id(schedule_id) if not schedule: return False, f"未找到 ID 为 {schedule_id} 的定时任务。" if schedule.plugin_name not in self._registered_tasks: return False, f"插件 '{schedule.plugin_name}' 没有注册可用的定时任务。" try: await _execute_job(schedule.id) return True, f"已手动触发任务 (ID: {schedule.id})。" except Exception as e: logger.error(f"手动触发任务失败: {e}") return False, f"手动触发任务失败: {e}" scheduler_manager = SchedulerManager()