107 lines
3.4 KiB
Python
107 lines
3.4 KiB
Python
"""Scheduler service for managing background tasks with APScheduler."""
|
|
|
|
import logging
|
|
from typing import Optional
|
|
|
|
from apscheduler.schedulers.background import BackgroundScheduler
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
|
|
from app.services.credit_service import CreditService
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SchedulerService:
|
|
"""Service for managing scheduled background tasks."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize the scheduler service."""
|
|
self.scheduler: Optional[BackgroundScheduler] = None
|
|
|
|
def start(self) -> None:
|
|
"""Start the scheduler and add all scheduled jobs."""
|
|
if self.scheduler is not None:
|
|
logger.warning("Scheduler is already running")
|
|
return
|
|
|
|
self.scheduler = BackgroundScheduler()
|
|
|
|
# Add daily credit refill job
|
|
self._add_daily_credit_refill_job()
|
|
|
|
# Start the scheduler
|
|
self.scheduler.start()
|
|
logger.info("Scheduler started successfully")
|
|
|
|
def stop(self) -> None:
|
|
"""Stop the scheduler."""
|
|
if self.scheduler is not None:
|
|
self.scheduler.shutdown()
|
|
self.scheduler = None
|
|
logger.info("Scheduler stopped")
|
|
|
|
def _add_daily_credit_refill_job(self) -> None:
|
|
"""Add the daily credit refill job."""
|
|
if self.scheduler is None:
|
|
raise RuntimeError("Scheduler not initialized")
|
|
|
|
# Schedule daily at 00:00 UTC
|
|
trigger = CronTrigger(hour=0, minute=0)
|
|
|
|
self.scheduler.add_job(
|
|
func=self._run_daily_credit_refill,
|
|
trigger=trigger,
|
|
id="daily_credit_refill",
|
|
name="Daily Credit Refill",
|
|
replace_existing=True,
|
|
)
|
|
|
|
logger.info("Daily credit refill job scheduled for 00:00 UTC")
|
|
|
|
def _run_daily_credit_refill(self) -> None:
|
|
"""Execute the daily credit refill task."""
|
|
logger.info("Starting daily credit refill task")
|
|
|
|
try:
|
|
result = CreditService.refill_all_users_credits()
|
|
|
|
if result["success"]:
|
|
logger.info(
|
|
f"Daily credit refill completed successfully: "
|
|
f"{result['users_processed']} users processed, "
|
|
f"{result['credits_added']} credits added"
|
|
)
|
|
else:
|
|
logger.error(f"Daily credit refill failed: {result['message']}")
|
|
|
|
except Exception as e:
|
|
logger.exception(f"Error during daily credit refill: {e}")
|
|
|
|
def trigger_credit_refill_now(self) -> dict:
|
|
"""Manually trigger credit refill for testing purposes."""
|
|
logger.info("Manually triggering credit refill")
|
|
return CreditService.refill_all_users_credits()
|
|
|
|
def get_scheduler_status(self) -> dict:
|
|
"""Get the current status of the scheduler."""
|
|
if self.scheduler is None:
|
|
return {"running": False, "jobs": []}
|
|
|
|
jobs = []
|
|
for job in self.scheduler.get_jobs():
|
|
jobs.append({
|
|
"id": job.id,
|
|
"name": job.name,
|
|
"next_run": job.next_run_time.isoformat()
|
|
if job.next_run_time else None,
|
|
"trigger": str(job.trigger),
|
|
})
|
|
|
|
return {
|
|
"running": self.scheduler.running,
|
|
"jobs": jobs,
|
|
}
|
|
|
|
|
|
# Global scheduler instance
|
|
scheduler_service = SchedulerService() |