""" Weibo-HotSign Task Scheduler Service Celery Beat configuration for scheduled sign-in tasks """ import os import sys import asyncio import logging from typing import Dict, List from datetime import datetime from celery import Celery from celery.schedules import crontab from croniter import croniter from sqlalchemy import select # Add parent directory to path for imports sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../..")) from shared.models.base import AsyncSessionLocal from shared.models.task import Task from shared.models.account import Account from shared.config import shared_settings from .config import settings logger = logging.getLogger(__name__) # Create Celery app celery_app = Celery( "weibo_hot_sign_scheduler", broker=settings.CELERY_BROKER_URL, backend=settings.CELERY_RESULT_BACKEND, include=["task_scheduler.app.tasks.signin_tasks"] ) # Celery configuration celery_app.conf.update( task_serializer="json", accept_content=["json"], result_serializer="json", timezone="Asia/Shanghai", enable_utc=True, beat_schedule_filename="celerybeat-schedule", beat_max_loop_interval=5, ) class TaskSchedulerService: """Service to manage scheduled tasks from database""" def __init__(self): self.scheduled_tasks: Dict[str, dict] = {} async def load_scheduled_tasks(self) -> List[Task]: """ Load enabled tasks from database and register them to Celery Beat. Returns list of loaded tasks. """ try: async with AsyncSessionLocal() as session: # Query all enabled tasks with their accounts stmt = ( select(Task, Account) .join(Account, Task.account_id == Account.id) .where(Task.is_enabled == True) ) result = await session.execute(stmt) task_account_pairs = result.all() logger.info(f"📅 Loaded {len(task_account_pairs)} enabled tasks from database") # Register tasks to Celery Beat dynamically beat_schedule = {} for task, account in task_account_pairs: try: # Validate cron expression if not croniter.is_valid(task.cron_expression): logger.warning(f"Invalid cron expression for task {task.id}: {task.cron_expression}") continue # Create schedule entry schedule_name = f"task_{task.id}" beat_schedule[schedule_name] = { "task": "task_scheduler.app.tasks.signin_tasks.execute_signin_task", "schedule": self._parse_cron_to_celery(task.cron_expression), "args": (task.id, task.account_id, task.cron_expression), } self.scheduled_tasks[task.id] = { "task_id": task.id, "account_id": task.account_id, "cron_expression": task.cron_expression, "account_status": account.status, } logger.info(f"✅ Registered task {task.id} for account {account.weibo_user_id} with cron: {task.cron_expression}") except Exception as e: logger.error(f"Failed to register task {task.id}: {e}") continue # Update Celery Beat schedule celery_app.conf.beat_schedule.update(beat_schedule) return [task for task, _ in task_account_pairs] except Exception as e: logger.error(f"❌ Error loading tasks from database: {e}") return [] def _parse_cron_to_celery(self, cron_expression: str) -> crontab: """ Parse cron expression string to Celery crontab schedule. Format: minute hour day month day_of_week """ parts = cron_expression.split() if len(parts) != 5: raise ValueError(f"Invalid cron expression: {cron_expression}") return crontab( minute=parts[0], hour=parts[1], day_of_month=parts[2], month_of_year=parts[3], day_of_week=parts[4], ) async def add_task(self, task_id: str, account_id: str, cron_expression: str): """Dynamically add a new task to the schedule""" try: if not croniter.is_valid(cron_expression): raise ValueError(f"Invalid cron expression: {cron_expression}") schedule_name = f"task_{task_id}" celery_app.conf.beat_schedule[schedule_name] = { "task": "task_scheduler.app.tasks.signin_tasks.execute_signin_task", "schedule": self._parse_cron_to_celery(cron_expression), "args": (task_id, account_id, cron_expression), } self.scheduled_tasks[task_id] = { "task_id": task_id, "account_id": account_id, "cron_expression": cron_expression, } logger.info(f"✅ Added task {task_id} to schedule") except Exception as e: logger.error(f"Failed to add task {task_id}: {e}") raise async def remove_task(self, task_id: str): """Dynamically remove a task from the schedule""" try: schedule_name = f"task_{task_id}" if schedule_name in celery_app.conf.beat_schedule: del celery_app.conf.beat_schedule[schedule_name] logger.info(f"✅ Removed task {task_id} from schedule") if task_id in self.scheduled_tasks: del self.scheduled_tasks[task_id] except Exception as e: logger.error(f"Failed to remove task {task_id}: {e}") raise async def update_task(self, task_id: str, is_enabled: bool, cron_expression: str = None): """Update an existing task in the schedule""" try: if is_enabled: # Re-add or update the task async with AsyncSessionLocal() as session: stmt = select(Task).where(Task.id == task_id) result = await session.execute(stmt) task = result.scalar_one_or_none() if task: await self.add_task( task_id, task.account_id, cron_expression or task.cron_expression ) else: # Remove the task await self.remove_task(task_id) logger.info(f"✅ Updated task {task_id}, enabled={is_enabled}") except Exception as e: logger.error(f"Failed to update task {task_id}: {e}") raise # Global scheduler service instance scheduler_service = TaskSchedulerService() # Synchronous wrapper for async function def sync_load_tasks(): """Synchronous wrapper to load tasks on startup""" loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: return loop.run_until_complete(scheduler_service.load_scheduled_tasks()) finally: loop.close() # Import task modules to register them from .tasks import signin_tasks