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