109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
"""Credit management service for handling daily credit refills."""
|
|
|
|
import logging
|
|
from datetime import datetime
|
|
from zoneinfo import ZoneInfo
|
|
|
|
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.now(tz=ZoneInfo("UTC"))
|
|
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: {e!s}")
|
|
|
|
return {
|
|
"success": False,
|
|
"users_processed": 0,
|
|
"credits_added": 0,
|
|
"error": str(e),
|
|
"message": "Credit refill failed",
|
|
}
|