140 lines
4.4 KiB
Python
140 lines
4.4 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",
|
|
}
|
|
|
|
@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,
|
|
}
|