weibo_signin/backend/task_scheduler/app/celery_app.py

"""
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