64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
"""Scheduler service for periodic tasks."""
|
|
|
|
from collections.abc import Callable
|
|
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
|
|
|
from app.core.logging import get_logger
|
|
from app.services.credit import CreditService
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class SchedulerService:
|
|
"""Service for managing scheduled tasks."""
|
|
|
|
def __init__(self, db_session_factory: Callable[[], AsyncSession]) -> None:
|
|
"""Initialize the scheduler service.
|
|
|
|
Args:
|
|
db_session_factory: Factory function to create database sessions
|
|
|
|
"""
|
|
self.db_session_factory = db_session_factory
|
|
self.scheduler = AsyncIOScheduler()
|
|
self.credit_service = CreditService(db_session_factory)
|
|
|
|
async def start(self) -> None:
|
|
"""Start the scheduler and register all tasks."""
|
|
logger.info("Starting scheduler service...")
|
|
|
|
# Add daily credit recharge job (runs at midnight UTC)
|
|
self.scheduler.add_job(
|
|
self._daily_credit_recharge,
|
|
"cron",
|
|
hour=0,
|
|
minute=0,
|
|
id="daily_credit_recharge",
|
|
name="Daily Credit Recharge",
|
|
replace_existing=True,
|
|
)
|
|
|
|
self.scheduler.start()
|
|
logger.info("Scheduler service started successfully")
|
|
|
|
async def stop(self) -> None:
|
|
"""Stop the scheduler."""
|
|
logger.info("Stopping scheduler service...")
|
|
self.scheduler.shutdown()
|
|
logger.info("Scheduler service stopped")
|
|
|
|
async def _daily_credit_recharge(self) -> None:
|
|
"""Execute daily credit recharge for all users."""
|
|
logger.info("Starting daily credit recharge task...")
|
|
|
|
try:
|
|
stats = await self.credit_service.recharge_all_users_credits()
|
|
logger.info(
|
|
"Daily credit recharge completed successfully: %s",
|
|
stats,
|
|
)
|
|
except Exception:
|
|
logger.exception("Daily credit recharge task failed")
|