393 lines
14 KiB
Python
393 lines
14 KiB
Python
"""
|
|
Scheduler service for managing scheduled download tasks.
|
|
Uses APScheduler with task queue for sequential execution.
|
|
"""
|
|
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
|
|
from datetime import datetime, time as time_type, timedelta
|
|
from sqlalchemy.orm import Session
|
|
import pytz
|
|
import logging
|
|
import asyncio
|
|
|
|
from .database import SessionLocal, DATA_DIR
|
|
from .models import ScheduledTask, TaskExecutionHistory
|
|
from .task_queue import task_queue
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Scheduler job store database (separate from main DB)
|
|
SCHEDULER_DB_PATH = DATA_DIR / "scheduler_jobs.db"
|
|
|
|
# Configure job stores
|
|
jobstores = {
|
|
'default': SQLAlchemyJobStore(url=f'sqlite:///{SCHEDULER_DB_PATH}')
|
|
}
|
|
|
|
# Configure scheduler for Docker container
|
|
scheduler = AsyncIOScheduler(
|
|
jobstores=jobstores,
|
|
timezone=pytz.UTC, # Container runs in UTC
|
|
job_defaults={
|
|
'coalesce': True, # Combine multiple missed executions into one
|
|
'max_instances': 1, # Only one instance of each job at a time
|
|
'misfire_grace_time': 3600 # Allow up to 1 hour late execution
|
|
}
|
|
)
|
|
|
|
|
|
async def queue_scheduled_task(task_id: str):
|
|
"""
|
|
Called by scheduler at the configured time.
|
|
Adds task to queue rather than executing immediately.
|
|
|
|
Args:
|
|
task_id: UUID of the scheduled task
|
|
"""
|
|
logger.info(f"Scheduler triggered for task {task_id}, adding to queue")
|
|
try:
|
|
await task_queue.add_task(task_id, priority=0) # Normal priority for scheduled tasks
|
|
except Exception as e:
|
|
logger.error(f"Failed to queue scheduled task {task_id}: {e}", exc_info=True)
|
|
|
|
|
|
def schedule_task(task: ScheduledTask):
|
|
"""
|
|
Schedule a task to be added to the queue at specified time.
|
|
|
|
Args:
|
|
task: ScheduledTask model instance
|
|
"""
|
|
if not task.enabled:
|
|
logger.info(f"Skipping scheduling for disabled task {task.id}")
|
|
return
|
|
|
|
try:
|
|
# Convert user's timezone to UTC for container execution
|
|
user_tz = pytz.timezone(task.timezone)
|
|
|
|
# Parse time
|
|
if isinstance(task.run_time, str):
|
|
hour, minute = map(int, task.run_time.split(':')[:2])
|
|
else:
|
|
hour = task.run_time.hour
|
|
minute = task.run_time.minute
|
|
|
|
# Create cron trigger with user's timezone
|
|
trigger = CronTrigger(
|
|
hour=hour,
|
|
minute=minute,
|
|
timezone=user_tz
|
|
)
|
|
|
|
# Add job to scheduler
|
|
scheduler.add_job(
|
|
func=queue_scheduled_task,
|
|
trigger=trigger,
|
|
args=[task.id],
|
|
id=str(task.id),
|
|
replace_existing=True,
|
|
name=f"{task.name} ({task.source_type}/{task.source_name})"
|
|
)
|
|
|
|
logger.info(f"Scheduled task {task.id} '{task.name}' for {hour:02d}:{minute:02d} {task.timezone}")
|
|
|
|
# Update next_run_at
|
|
next_run = trigger.get_next_fire_time(None, datetime.now(user_tz))
|
|
if next_run:
|
|
db = SessionLocal()
|
|
try:
|
|
db_task = db.query(ScheduledTask).filter(ScheduledTask.id == task.id).first()
|
|
if db_task:
|
|
db_task.next_run_at = next_run
|
|
db.commit()
|
|
finally:
|
|
db.close()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to schedule task {task.id}: {e}", exc_info=True)
|
|
|
|
|
|
def unschedule_task(task_id: str):
|
|
"""
|
|
Remove a task from the scheduler.
|
|
|
|
Args:
|
|
task_id: UUID of the scheduled task
|
|
"""
|
|
try:
|
|
if scheduler.get_job(task_id):
|
|
scheduler.remove_job(task_id)
|
|
logger.info(f"Unscheduled task {task_id}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to unschedule task {task_id}: {e}", exc_info=True)
|
|
|
|
|
|
def load_scheduled_tasks():
|
|
"""
|
|
Load all enabled scheduled tasks from database and schedule them.
|
|
Called on application startup.
|
|
"""
|
|
db = SessionLocal()
|
|
try:
|
|
# Check if the required columns exist before querying
|
|
try:
|
|
# First, check if we can access the table at all
|
|
db.query(ScheduledTask).first()
|
|
|
|
# Try to query with the new columns - this will fail if columns don't exist
|
|
tasks = db.query(ScheduledTask).filter(ScheduledTask.enabled == True).all()
|
|
except Exception as column_error:
|
|
if "no such column" in str(column_error):
|
|
logger.warning(f"Database schema is outdated. Required columns missing: {column_error}")
|
|
logger.info("Attempting to upgrade database schema...")
|
|
|
|
# Try to upgrade the database
|
|
try:
|
|
from .database import upgrade_database_schema
|
|
upgrade_database_schema()
|
|
|
|
# Now try again to load tasks
|
|
tasks = db.query(ScheduledTask).filter(ScheduledTask.enabled == True).all()
|
|
logger.info("Database upgraded successfully, continuing with task loading")
|
|
except Exception as upgrade_error:
|
|
logger.error(f"Failed to upgrade database: {upgrade_error}")
|
|
logger.info("Skipping scheduled task loading until database is manually upgraded")
|
|
return
|
|
else:
|
|
raise
|
|
|
|
logger.info(f"Loading {len(tasks)} enabled scheduled tasks")
|
|
|
|
for task in tasks:
|
|
schedule_task(task)
|
|
|
|
logger.info(f"Loaded {len(tasks)} scheduled tasks")
|
|
except Exception as e:
|
|
logger.error(f"Failed to load scheduled tasks: {e}", exc_info=True)
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
async def execute_scheduled_task(task_id: str):
|
|
"""
|
|
Execute a scheduled download task.
|
|
This function BLOCKS until the download is complete,
|
|
ensuring sequential execution.
|
|
|
|
Args:
|
|
task_id: UUID of the scheduled task
|
|
"""
|
|
db = SessionLocal()
|
|
try:
|
|
# Load task from database
|
|
task = db.query(ScheduledTask).filter(ScheduledTask.id == task_id).first()
|
|
|
|
if not task:
|
|
logger.error(f"Task {task_id} not found in database")
|
|
return
|
|
|
|
if not task.enabled:
|
|
logger.info(f"Skipping disabled task {task_id}")
|
|
# Still record in history that it was skipped
|
|
execution = TaskExecutionHistory(
|
|
task_id=task_id,
|
|
status='skipped',
|
|
completed_at=datetime.now(pytz.UTC)
|
|
)
|
|
db.add(execution)
|
|
db.commit()
|
|
return
|
|
|
|
logger.info(f"Executing scheduled task {task_id}: {task.source_type}/{task.source_name}")
|
|
|
|
# Create execution history record
|
|
execution = TaskExecutionHistory(
|
|
task_id=task_id,
|
|
status='running',
|
|
started_at=datetime.now(pytz.UTC)
|
|
)
|
|
db.add(execution)
|
|
db.commit()
|
|
execution_id = execution.id
|
|
|
|
try:
|
|
# Import here to avoid circular imports
|
|
from .main import create_download_with_bdfr_api, bdfr_manager
|
|
|
|
# Build download parameters
|
|
kwargs = {
|
|
'limit': task.limit,
|
|
'sort': task.sort,
|
|
'time_filter': 'day', # Always "day" for daily scheduled tasks
|
|
'no_dupes': True, # Always enabled for scheduled tasks
|
|
'simple_check': task.simple_check,
|
|
'auth_state': task.auth_state,
|
|
'upvoted': task.upvoted,
|
|
'saved': task.saved
|
|
}
|
|
|
|
# For user downloads, set submitted appropriately
|
|
if task.source_type == 'user':
|
|
# For likes or saved, don't download submitted posts
|
|
kwargs['submitted'] = not (task.upvoted or task.saved)
|
|
|
|
# Create download using existing API
|
|
download_id = await create_download_with_bdfr_api(
|
|
download_type=task.source_type,
|
|
name=task.source_name,
|
|
**kwargs
|
|
)
|
|
|
|
logger.info(f"Task {task_id} started as download {download_id}")
|
|
|
|
# Update execution record with download_id
|
|
execution = db.query(TaskExecutionHistory).filter(TaskExecutionHistory.id == execution_id).first()
|
|
if execution:
|
|
execution.download_id = download_id
|
|
db.commit()
|
|
|
|
# **CRITICAL: Wait for download to complete before returning**
|
|
await wait_for_download_completion(download_id, bdfr_manager)
|
|
|
|
# Check final status
|
|
download_status = bdfr_manager.get_download_status(download_id)
|
|
|
|
# Update execution history
|
|
execution = db.query(TaskExecutionHistory).filter(TaskExecutionHistory.id == execution_id).first()
|
|
if execution:
|
|
if download_status and download_status['status'] == 'completed':
|
|
execution.status = 'success'
|
|
execution.items_found = download_status.get('items_found', 0)
|
|
execution.items_downloaded = download_status.get('items_processed', 0)
|
|
logger.info(f"Task {task_id} completed successfully")
|
|
else:
|
|
execution.status = 'failed'
|
|
execution.error_message = download_status.get('error', 'Unknown error') if download_status else 'Download status not found'
|
|
logger.error(f"Task {task_id} failed: {execution.error_message}")
|
|
|
|
execution.completed_at = datetime.now(pytz.UTC)
|
|
db.commit()
|
|
|
|
# Update task timestamps
|
|
task.last_run_at = datetime.now(pytz.UTC)
|
|
task.next_run_at = calculate_next_run(task)
|
|
db.commit()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Task {task_id} execution error: {e}", exc_info=True)
|
|
|
|
# Update execution history with error
|
|
execution = db.query(TaskExecutionHistory).filter(TaskExecutionHistory.id == execution_id).first()
|
|
if execution:
|
|
execution.status = 'failed'
|
|
execution.error_message = str(e)
|
|
execution.completed_at = datetime.now(pytz.UTC)
|
|
db.commit()
|
|
|
|
raise
|
|
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
async def wait_for_download_completion(download_id: str, bdfr_manager, timeout: int = 3600):
|
|
"""
|
|
Wait for a download to complete.
|
|
Polls the download status until it's no longer running.
|
|
|
|
Args:
|
|
download_id: The download to wait for
|
|
bdfr_manager: BDFRManager instance
|
|
timeout: Maximum seconds to wait (default 1 hour)
|
|
"""
|
|
start_time = datetime.now()
|
|
check_interval = 5 # Check every 5 seconds
|
|
|
|
logger.info(f"Waiting for download {download_id} to complete...")
|
|
|
|
while True:
|
|
# Check if timeout exceeded
|
|
elapsed = (datetime.now() - start_time).total_seconds()
|
|
if elapsed > timeout:
|
|
logger.error(f"Download {download_id} timed out after {timeout}s")
|
|
raise TimeoutError(f"Download exceeded timeout of {timeout}s")
|
|
|
|
# Check download status
|
|
status = bdfr_manager.get_download_status(download_id)
|
|
|
|
if not status:
|
|
logger.warning(f"Download {download_id} status not found, assuming complete")
|
|
break
|
|
|
|
download_status = status.get('status', 'unknown')
|
|
|
|
# Check if download is finished (completed, failed, or cancelled)
|
|
if download_status in ['completed', 'failed', 'cancelled']:
|
|
logger.info(f"Download {download_id} finished with status: {download_status}")
|
|
break
|
|
|
|
# Still running, wait before checking again
|
|
await asyncio.sleep(check_interval)
|
|
|
|
|
|
def calculate_next_run(task: ScheduledTask) -> datetime:
|
|
"""
|
|
Calculate the next run time for a task based on its schedule.
|
|
|
|
Args:
|
|
task: ScheduledTask instance
|
|
|
|
Returns:
|
|
Next run datetime in UTC
|
|
"""
|
|
user_tz = pytz.timezone(task.timezone)
|
|
now = datetime.now(user_tz)
|
|
|
|
# Parse run time
|
|
if isinstance(task.run_time, str):
|
|
hour, minute = map(int, task.run_time.split(':')[:2])
|
|
else:
|
|
hour = task.run_time.hour
|
|
minute = task.run_time.minute
|
|
|
|
# Calculate next run
|
|
next_run = now.replace(hour=hour, minute=minute, second=0, microsecond=0)
|
|
|
|
# If time has passed today, schedule for tomorrow
|
|
if next_run <= now:
|
|
next_run += timedelta(days=1)
|
|
|
|
# Convert to UTC
|
|
return next_run.astimezone(pytz.UTC)
|
|
|
|
|
|
def start_scheduler():
|
|
"""Start the scheduler. Called on application startup."""
|
|
try:
|
|
if not scheduler.running:
|
|
scheduler.start()
|
|
logger.info("Scheduler started")
|
|
|
|
# Set execute callback for task queue
|
|
task_queue.set_execute_callback(execute_scheduled_task)
|
|
|
|
# Load existing tasks
|
|
load_scheduled_tasks()
|
|
else:
|
|
logger.info("Scheduler already running")
|
|
except Exception as e:
|
|
logger.error(f"Failed to start scheduler: {e}", exc_info=True)
|
|
raise
|
|
|
|
|
|
def stop_scheduler():
|
|
"""Stop the scheduler. Called on application shutdown."""
|
|
try:
|
|
if scheduler.running:
|
|
scheduler.shutdown(wait=True)
|
|
logger.info("Scheduler stopped")
|
|
except Exception as e:
|
|
logger.error(f"Failed to stop scheduler: {e}", exc_info=True) |