224 lines
8.6 KiB
Python
224 lines
8.6 KiB
Python
"""
|
|
Scheduler Service
|
|
Manages APScheduler and dynamic task execution
|
|
"""
|
|
import logging
|
|
import asyncio
|
|
from datetime import datetime
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.database import AsyncSessionLocal
|
|
from app.models.db.scheduled_job import ScheduledJob, JobType
|
|
from app.models.db.task import Task
|
|
from app.services.task_service import task_service
|
|
from app.jobs.registry import task_registry
|
|
# Import predefined tasks to register them
|
|
import app.jobs.predefined # noqa: F401
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SchedulerService:
|
|
def __init__(self):
|
|
self.scheduler = AsyncIOScheduler()
|
|
self.jobs = {}
|
|
|
|
def start(self):
|
|
"""Start the scheduler"""
|
|
if not self.scheduler.running:
|
|
self.scheduler.start()
|
|
logger.info("Scheduler started")
|
|
# Load jobs from DB
|
|
asyncio.create_task(self.load_jobs())
|
|
|
|
def shutdown(self):
|
|
"""Shutdown the scheduler"""
|
|
if self.scheduler.running:
|
|
self.scheduler.shutdown()
|
|
logger.info("Scheduler stopped")
|
|
|
|
async def load_jobs(self):
|
|
"""Load active jobs from database and schedule them"""
|
|
logger.info("Loading scheduled jobs from database...")
|
|
async with AsyncSessionLocal() as session:
|
|
result = await session.execute(select(ScheduledJob).where(ScheduledJob.is_active == True))
|
|
jobs = result.scalars().all()
|
|
|
|
for job in jobs:
|
|
self.add_job_to_scheduler(job)
|
|
|
|
logger.info(f"Loaded {len(jobs)} scheduled jobs")
|
|
|
|
def add_job_to_scheduler(self, job: ScheduledJob):
|
|
"""Add a single job to APScheduler"""
|
|
try:
|
|
# Remove existing job if any (to update)
|
|
if str(job.id) in self.jobs:
|
|
self.scheduler.remove_job(str(job.id))
|
|
|
|
# Create trigger from cron expression
|
|
# Cron format: "minute hour day month day_of_week"
|
|
# APScheduler expects kwargs, so we need to parse or use from_crontab if strictly standard
|
|
# But CronTrigger.from_crontab is standard.
|
|
trigger = CronTrigger.from_crontab(job.cron_expression)
|
|
|
|
self.scheduler.add_job(
|
|
self.execute_job,
|
|
trigger,
|
|
args=[job.id],
|
|
id=str(job.id),
|
|
name=job.name,
|
|
replace_existing=True
|
|
)
|
|
self.jobs[str(job.id)] = job
|
|
logger.info(f"Scheduled job '{job.name}' (ID: {job.id}) with cron: {job.cron_expression}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to schedule job '{job.name}': {e}")
|
|
|
|
async def execute_job(self, job_id: int):
|
|
"""
|
|
Execute either a predefined task or dynamic python code for a job.
|
|
This runs in the scheduler's event loop.
|
|
"""
|
|
logger.info(f"Executing job ID: {job_id}")
|
|
|
|
async with AsyncSessionLocal() as session:
|
|
# Fetch job details again to get latest configuration
|
|
result = await session.execute(select(ScheduledJob).where(ScheduledJob.id == job_id))
|
|
job = result.scalar_one_or_none()
|
|
|
|
if not job:
|
|
logger.error(f"Job {job_id} not found")
|
|
return
|
|
|
|
# Validate job configuration
|
|
if job.job_type == JobType.PREDEFINED and not job.predefined_function:
|
|
logger.error(f"Job {job_id} is predefined type but has no function name")
|
|
return
|
|
elif job.job_type == JobType.CUSTOM_CODE and not job.python_code:
|
|
logger.error(f"Job {job_id} is custom_code type but has no code")
|
|
return
|
|
|
|
# Create a Task record for this execution history
|
|
task_record = await task_service.create_task(
|
|
session,
|
|
task_type="scheduled_job",
|
|
description=f"Scheduled execution of '{job.name}'",
|
|
params={"job_id": job.id, "job_type": job.job_type.value},
|
|
created_by=None # System
|
|
)
|
|
|
|
# Update Task to running
|
|
await task_service.update_task(session, task_record.id, status="running", started_at=datetime.utcnow(), progress=0)
|
|
|
|
# Update Job last run time
|
|
job.last_run_at = datetime.utcnow()
|
|
await session.commit()
|
|
|
|
try:
|
|
# Execute based on job type
|
|
if job.job_type == JobType.PREDEFINED:
|
|
# Execute predefined task from registry
|
|
logger.debug(f"Executing predefined task: {job.predefined_function}")
|
|
result_val = await task_registry.execute_task(
|
|
name=job.predefined_function,
|
|
db=session,
|
|
logger=logger,
|
|
params=job.function_params or {}
|
|
)
|
|
else:
|
|
# Execute custom Python code (legacy support)
|
|
logger.debug(f"Executing custom code for job: {job.name}")
|
|
# Prepare execution context
|
|
# We inject useful services and variables
|
|
context = {
|
|
"db": session,
|
|
"logger": logger,
|
|
"task_id": task_record.id,
|
|
"asyncio": asyncio,
|
|
# Import commonly used services here if needed, or let code import them
|
|
}
|
|
|
|
# Wrap code in an async function to allow await
|
|
# Indent code to fit inside the wrapper
|
|
indented_code = "\n".join([" " + line for line in job.python_code.split("\n")])
|
|
wrapper_code = f"async def _dynamic_func():\n{indented_code}"
|
|
|
|
# Execute definition
|
|
exec(wrapper_code, context)
|
|
|
|
# Execute the function
|
|
_func = context["_dynamic_func"]
|
|
result_val = await _func()
|
|
|
|
# Success
|
|
await task_service.update_task(
|
|
session,
|
|
task_record.id,
|
|
status="completed",
|
|
progress=100,
|
|
completed_at=datetime.utcnow(),
|
|
result={"output": str(result_val) if result_val else "Success"}
|
|
)
|
|
job.last_run_status = "success"
|
|
logger.info(f"Job '{job.name}' completed successfully")
|
|
|
|
except Exception as e:
|
|
# Failure
|
|
import traceback
|
|
error_msg = f"{str(e)}\n{traceback.format_exc()}"
|
|
logger.error(f"Job '{job.name}' failed: {e}")
|
|
|
|
# Rollback the current transaction
|
|
await session.rollback()
|
|
|
|
# Start a new transaction to update task status
|
|
try:
|
|
await task_service.update_task(
|
|
session,
|
|
task_record.id,
|
|
status="failed",
|
|
error_message=error_msg,
|
|
completed_at=datetime.utcnow()
|
|
)
|
|
job.last_run_status = "failed"
|
|
# Commit the failed task update in new transaction
|
|
await session.commit()
|
|
except Exception as update_error:
|
|
logger.error(f"Failed to update task status: {update_error}")
|
|
await session.rollback()
|
|
|
|
else:
|
|
# Success - commit only if no exception
|
|
await session.commit()
|
|
|
|
async def reload_job(self, job_id: int):
|
|
"""Reload a specific job from DB (after update)"""
|
|
async with AsyncSessionLocal() as session:
|
|
result = await session.execute(select(ScheduledJob).where(ScheduledJob.id == job_id))
|
|
job = result.scalar_one_or_none()
|
|
if job:
|
|
if job.is_active:
|
|
self.add_job_to_scheduler(job)
|
|
else:
|
|
self.remove_job(job_id)
|
|
|
|
def remove_job(self, job_id: int):
|
|
"""Remove job from scheduler"""
|
|
if str(job_id) in self.jobs:
|
|
if self.scheduler.get_job(str(job_id)):
|
|
self.scheduler.remove_job(str(job_id))
|
|
del self.jobs[str(job_id)]
|
|
logger.info(f"Removed job ID: {job_id}")
|
|
|
|
async def run_job_now(self, job_id: int):
|
|
"""Manually trigger a job immediately"""
|
|
return await self.execute_job(job_id)
|
|
|
|
|
|
# Singleton
|
|
scheduler_service = SchedulerService()
|