扫码登录，获取cookies
This commit is contained in:
@@ -4,22 +4,35 @@ Celery Beat configuration for scheduled sign-in tasks
|
||||
"""
|
||||
|
||||
import os
|
||||
from celery import Celery
|
||||
from celery.schedules import crontab
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from sqlalchemy import select
|
||||
import sys
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, List
|
||||
from datetime import datetime
|
||||
|
||||
from ..config import settings
|
||||
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=["app.tasks.signin_tasks"]
|
||||
include=["task_scheduler.app.tasks.signin_tasks"]
|
||||
)
|
||||
|
||||
# Celery configuration
|
||||
@@ -33,65 +46,168 @@ celery_app.conf.update(
|
||||
beat_max_loop_interval=5,
|
||||
)
|
||||
|
||||
# Database configuration for task scheduler
|
||||
engine = create_async_engine(
|
||||
settings.DATABASE_URL,
|
||||
echo=settings.DEBUG,
|
||||
pool_size=10,
|
||||
max_overflow=20
|
||||
)
|
||||
|
||||
AsyncSessionLocal = sessionmaker(
|
||||
engine,
|
||||
class_=AsyncSession,
|
||||
expire_on_commit=False
|
||||
)
|
||||
|
||||
async def get_db():
|
||||
"""Get database session for task scheduler"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
try:
|
||||
yield session
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
class TaskSchedulerService:
|
||||
"""Service to manage scheduled tasks from database"""
|
||||
|
||||
def __init__(self):
|
||||
self.engine = engine
|
||||
self.scheduled_tasks: Dict[str, dict] = {}
|
||||
|
||||
async def load_scheduled_tasks(self):
|
||||
"""Load enabled tasks from database and schedule them"""
|
||||
from app.models.task_models import Task
|
||||
|
||||
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
|
||||
stmt = select(Task).where(Task.is_enabled == True)
|
||||
# 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)
|
||||
tasks = result.scalars().all()
|
||||
task_account_pairs = result.all()
|
||||
|
||||
print(f"📅 Loaded {len(tasks)} enabled tasks from database")
|
||||
logger.info(f"📅 Loaded {len(task_account_pairs)} enabled tasks from database")
|
||||
|
||||
# Here we would dynamically add tasks to Celery Beat
|
||||
# For now, we'll use static configuration in celery_config.py
|
||||
return tasks
|
||||
# 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:
|
||||
print(f"❌ Error loading tasks from database: {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"""
|
||||
service = TaskSchedulerService()
|
||||
"""Synchronous wrapper to load tasks on startup"""
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(service.load_scheduled_tasks())
|
||||
return loop.run_until_complete(scheduler_service.load_scheduled_tasks())
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
|
||||
# Import task modules to register them
|
||||
from app.tasks import signin_tasks
|
||||
from .tasks import signin_tasks
|
||||
|
||||
Reference in New Issue
Block a user