214 lines
7.7 KiB
Python
214 lines
7.7 KiB
Python
"""
|
|
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
|