82 lines
2.8 KiB
Python
82 lines
2.8 KiB
Python
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
from sqlalchemy import select
|
|
from typing import Optional
|
|
import logging
|
|
|
|
from app.database import AsyncSessionLocal
|
|
from app.models import Schedule
|
|
from backup.engine import BackupEngine
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BackupScheduler:
|
|
def __init__(self):
|
|
self.scheduler = AsyncIOScheduler()
|
|
self._job_map = {}
|
|
|
|
def start(self):
|
|
"""Start the scheduler."""
|
|
self.scheduler.start()
|
|
logger.info("Backup scheduler started")
|
|
|
|
def shutdown(self):
|
|
"""Shutdown the scheduler."""
|
|
self.scheduler.shutdown()
|
|
logger.info("Backup scheduler shutdown")
|
|
|
|
async def sync_schedules(self):
|
|
"""Sync all enabled schedules from database."""
|
|
async with AsyncSessionLocal() as db:
|
|
result = await db.execute(
|
|
select(Schedule).where(Schedule.enabled == True)
|
|
)
|
|
schedules = result.scalars().all()
|
|
|
|
# Clear existing jobs
|
|
for schedule_id, job_id in list(self._job_map.items()):
|
|
self.scheduler.remove_job(job_id)
|
|
del self._job_map[schedule_id]
|
|
|
|
# Add new jobs
|
|
for schedule in schedules:
|
|
await self._add_schedule_job(schedule)
|
|
|
|
async def _add_schedule_job(self, schedule: Schedule):
|
|
"""Add a single schedule job to the scheduler."""
|
|
try:
|
|
trigger = CronTrigger.from_crontab(schedule.cron_expression)
|
|
job = self.scheduler.add_job(
|
|
self._run_backup_job,
|
|
trigger=trigger,
|
|
args=[schedule.job_id],
|
|
id=f"backup_job_{schedule.job_id}",
|
|
replace_existing=True
|
|
)
|
|
self._job_map[schedule.id] = job.id
|
|
logger.info(f"Scheduled backup job {schedule.job_id} with cron: {schedule.cron_expression}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to schedule job {schedule.job_id}: {e}")
|
|
|
|
async def _run_backup_job(self, job_id: int):
|
|
"""Execute a backup job."""
|
|
logger.info(f"Running scheduled backup job {job_id}")
|
|
async with AsyncSessionLocal() as db:
|
|
engine = BackupEngine(db)
|
|
await engine.execute_job(job_id, triggered_by="schedule")
|
|
|
|
async def add_schedule(self, schedule: Schedule):
|
|
"""Add a new schedule to the scheduler."""
|
|
await self._add_schedule_job(schedule)
|
|
|
|
def remove_schedule(self, schedule_id: int):
|
|
"""Remove a schedule from the scheduler."""
|
|
if schedule_id in self._job_map:
|
|
self.scheduler.remove_job(self._job_map[schedule_id])
|
|
del self._job_map[schedule_id]
|
|
|
|
|
|
# Global scheduler instance
|
|
backup_scheduler = BackupScheduler()
|