89 lines
3.3 KiB
Python
89 lines
3.3 KiB
Python
from typing import Callable, Dict, Any
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from apscheduler.triggers.interval import IntervalTrigger
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
import logging
|
|
from apscheduler.schedulers.base import SchedulerNotRunningError
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class BaseScheduler:
|
|
def __init__(self):
|
|
self.scheduler = AsyncIOScheduler()
|
|
self.jobs: Dict[str, Any] = {}
|
|
self._is_running = False
|
|
|
|
async def schedule_task(
|
|
self,
|
|
task_name: str,
|
|
func: Callable,
|
|
interval_seconds: int = None,
|
|
cron_expression: str = None,
|
|
*args,
|
|
**kwargs
|
|
) -> None:
|
|
"""Schedule a task to run at regular intervals or at specific times using APScheduler
|
|
|
|
Args:
|
|
task_name: Name of the task
|
|
func: Function to execute
|
|
interval_seconds: Interval in seconds for periodic execution (mutually exclusive with cron_expression)
|
|
cron_expression: Cron expression for time-based scheduling (mutually exclusive with interval_seconds)
|
|
*args: Additional positional arguments for the function
|
|
**kwargs: Additional keyword arguments for the function
|
|
"""
|
|
if task_name in self.jobs:
|
|
logger.warning(f"Task {task_name} already exists. Removing existing job.")
|
|
self.jobs[task_name].remove()
|
|
|
|
if interval_seconds and cron_expression:
|
|
raise ValueError("Cannot specify both interval_seconds and cron_expression")
|
|
elif not interval_seconds and not cron_expression:
|
|
raise ValueError("Must specify either interval_seconds or cron_expression")
|
|
|
|
if interval_seconds:
|
|
trigger = IntervalTrigger(seconds=interval_seconds)
|
|
else:
|
|
trigger = CronTrigger.from_crontab(cron_expression)
|
|
|
|
job = self.scheduler.add_job(
|
|
func,
|
|
trigger=trigger,
|
|
args=args,
|
|
kwargs=kwargs,
|
|
id=task_name,
|
|
replace_existing=True
|
|
)
|
|
|
|
self.jobs[task_name] = job
|
|
if interval_seconds:
|
|
logger.info(f"Scheduled task {task_name} to run every {interval_seconds} seconds")
|
|
else:
|
|
logger.info(f"Scheduled task {task_name} with cron expression: {cron_expression}")
|
|
|
|
def remove_task(self, task_name: str) -> None:
|
|
"""Remove a scheduled task"""
|
|
if task_name in self.jobs:
|
|
self.jobs[task_name].remove()
|
|
del self.jobs[task_name]
|
|
logger.info(f"Removed task {task_name}")
|
|
|
|
def start(self) -> None:
|
|
"""Start the scheduler"""
|
|
if not self._is_running:
|
|
self.scheduler.start()
|
|
self._is_running = True
|
|
logger.info("Scheduler started")
|
|
|
|
async def shutdown(self) -> None:
|
|
"""Shutdown the scheduler"""
|
|
if self._is_running:
|
|
try:
|
|
self.scheduler.shutdown()
|
|
self._is_running = False
|
|
logger.info("Scheduler stopped")
|
|
except SchedulerNotRunningError:
|
|
logger.warning("Scheduler was already stopped")
|
|
except Exception as e:
|
|
logger.error(f"Error shutting down scheduler: {str(e)}")
|
|
raise |