feat(scheduler): implement scheduler service for background tasks and credit refills; add endpoints for admin control
This commit is contained in:
JSC
@@ -7,6 +7,7 @@ from flask_jwt_extended import JWTManager
|
||||
|
||||
from app.database import init_db
|
||||
from app.services.auth_service import AuthService
|
||||
from app.services.scheduler_service import scheduler_service
|
||||
|
||||
# Global auth service instance
|
||||
auth_service = AuthService()
|
||||
@@ -60,10 +61,20 @@ def create_app():
|
||||
# Initialize authentication service with app
|
||||
auth_service.init_app(app)
|
||||
|
||||
# Start scheduler for background tasks
|
||||
scheduler_service.start()
|
||||
|
||||
# Register blueprints
|
||||
from app.routes import auth, main
|
||||
|
||||
app.register_blueprint(main.bp, url_prefix="/api")
|
||||
app.register_blueprint(auth.bp, url_prefix="/api/auth")
|
||||
|
||||
# Shutdown scheduler when app is torn down
|
||||
@app.teardown_appcontext
|
||||
def shutdown_scheduler(exception):
|
||||
"""Stop scheduler when app context is torn down."""
|
||||
if exception:
|
||||
scheduler_service.stop()
|
||||
|
||||
return app
|
||||
|
||||
@@ -8,6 +8,7 @@ from app.services.decorators import (
|
||||
require_credits,
|
||||
require_role,
|
||||
)
|
||||
from app.services.scheduler_service import scheduler_service
|
||||
|
||||
bp = Blueprint("main", __name__)
|
||||
|
||||
@@ -84,3 +85,19 @@ def expensive_operation() -> dict[str, str]:
|
||||
"user": user["email"],
|
||||
"operation_cost": 10,
|
||||
}
|
||||
|
||||
|
||||
@bp.route("/admin/scheduler/status")
|
||||
@require_auth
|
||||
@require_role("admin")
|
||||
def scheduler_status() -> dict:
|
||||
"""Get scheduler status (admin only)."""
|
||||
return scheduler_service.get_scheduler_status()
|
||||
|
||||
|
||||
@bp.route("/admin/credits/refill", methods=["POST"])
|
||||
@require_auth
|
||||
@require_role("admin")
|
||||
def manual_credit_refill() -> dict:
|
||||
"""Manually trigger credit refill for all users (admin only)."""
|
||||
return scheduler_service.trigger_credit_refill_now()
|
||||
|
||||
133
app/services/credit_service.py
Normal file
133
app/services/credit_service.py
Normal file
@@ -0,0 +1,133 @@
|
||||
"""Credit management service for handling daily credit refills."""
|
||||
|
||||
import logging
|
||||
from datetime import datetime
|
||||
|
||||
from app.database import db
|
||||
from app.models.user import User
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CreditService:
|
||||
"""Service for managing user credits and daily refills."""
|
||||
|
||||
@staticmethod
|
||||
def refill_all_users_credits() -> dict:
|
||||
"""
|
||||
Refill credits for all active users based on their plan.
|
||||
|
||||
This function:
|
||||
1. Gets all active users
|
||||
2. For each user, adds their plan's daily credit amount
|
||||
3. Ensures credits never exceed the plan's max_credits limit
|
||||
4. Updates all users in a single database transaction
|
||||
|
||||
Returns:
|
||||
dict: Summary of the refill operation
|
||||
"""
|
||||
try:
|
||||
# Get all active users with their plans
|
||||
users = User.query.filter_by(is_active=True).all()
|
||||
|
||||
if not users:
|
||||
logger.info("No active users found for credit refill")
|
||||
return {
|
||||
"success": True,
|
||||
"users_processed": 0,
|
||||
"credits_added": 0,
|
||||
"message": "No active users found"
|
||||
}
|
||||
|
||||
users_processed = 0
|
||||
total_credits_added = 0
|
||||
|
||||
for user in users:
|
||||
if not user.plan:
|
||||
logger.warning(f"User {user.email} has no plan assigned, skipping")
|
||||
continue
|
||||
|
||||
# Calculate new credit amount, capped at plan max
|
||||
current_credits = user.credits or 0
|
||||
plan_daily_credits = user.plan.credits
|
||||
max_credits = user.plan.max_credits
|
||||
|
||||
# Add daily credits but don't exceed maximum
|
||||
new_credits = min(current_credits + plan_daily_credits, max_credits)
|
||||
credits_added = new_credits - current_credits
|
||||
|
||||
if credits_added > 0:
|
||||
user.credits = new_credits
|
||||
user.updated_at = datetime.utcnow()
|
||||
total_credits_added += credits_added
|
||||
|
||||
logger.debug(
|
||||
f"User {user.email}: {current_credits} -> {new_credits} "
|
||||
f"(+{credits_added} credits, plan: {user.plan.code})"
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
f"User {user.email}: Already at max credits "
|
||||
f"({current_credits}/{max_credits})"
|
||||
)
|
||||
|
||||
users_processed += 1
|
||||
|
||||
# Commit all changes in a single transaction
|
||||
db.session.commit()
|
||||
|
||||
logger.info(
|
||||
f"Daily credit refill completed: {users_processed} users processed, "
|
||||
f"{total_credits_added} total credits added"
|
||||
)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"users_processed": users_processed,
|
||||
"credits_added": total_credits_added,
|
||||
"message": f"Successfully refilled credits for {users_processed} users"
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
# Rollback transaction on error
|
||||
db.session.rollback()
|
||||
logger.error(f"Error during daily credit refill: {str(e)}")
|
||||
|
||||
return {
|
||||
"success": False,
|
||||
"users_processed": 0,
|
||||
"credits_added": 0,
|
||||
"error": str(e),
|
||||
"message": "Credit refill failed"
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_user_credit_info(user_id: int) -> dict:
|
||||
"""
|
||||
Get detailed credit information for a specific user.
|
||||
|
||||
Args:
|
||||
user_id: The user's ID
|
||||
|
||||
Returns:
|
||||
dict: User's credit information
|
||||
"""
|
||||
user = User.query.get(user_id)
|
||||
if not user:
|
||||
return {"error": "User not found"}
|
||||
|
||||
if not user.plan:
|
||||
return {"error": "User has no plan assigned"}
|
||||
|
||||
return {
|
||||
"user_id": user.id,
|
||||
"email": user.email,
|
||||
"current_credits": user.credits,
|
||||
"plan": {
|
||||
"code": user.plan.code,
|
||||
"name": user.plan.name,
|
||||
"daily_credits": user.plan.credits,
|
||||
"max_credits": user.plan.max_credits
|
||||
},
|
||||
"is_active": user.is_active
|
||||
}
|
||||
107
app/services/scheduler_service.py
Normal file
107
app/services/scheduler_service.py
Normal file
@@ -0,0 +1,107 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user