feat: publish v1.0.0-rc.36-custom.1

model/user_quota_adjustment.go

@@ -0,0 +1,87 @@
+package model
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/shopspring/decimal"
+	"gorm.io/gorm"
+)
+
+var (
+	ErrInvalidUserQuotaAdjustment = errors.New("invalid user quota adjustment")
+	ErrUserQuotaPermission        = errors.New("cannot adjust quota for this user role")
+)
+
+// UserQuotaAdjustment is the immutable database snapshot of a committed manual
+// adjustment. Pending relay deductions in the quota cache are not part of it.
+type UserQuotaAdjustment struct {
+	UserID   int
+	Username string
+	Before   int
+	After    int
+}
+
+func AdjustUserQuota(userID, operatorRole int, mode string, value int) (*UserQuotaAdjustment, error) {
+	if userID <= 0 || (mode != "add" && mode != "subtract" && mode != "override") {
+		return nil, ErrInvalidUserQuotaAdjustment
+	}
+	if mode != "override" && value <= 0 {
+		return nil, ErrInvalidUserQuotaAdjustment
+	}
+	if value > common.MaxWalletQuota || value < -common.MaxWalletQuota {
+		return nil, ErrWalletQuotaLimitExceeded
+	}
+
+	var adjustment UserQuotaAdjustment
+	err := DB.Transaction(func(tx *gorm.DB) error {
+		var user User
+		if err := lockForUpdate(tx).First(&user, userID).Error; err != nil {
+			return err
+		}
+		if operatorRole != common.RoleRootUser && operatorRole <= user.Role {
+			return ErrUserQuotaPermission
+		}
+		if user.Quota > common.MaxWalletQuota || user.Quota < -common.MaxWalletQuota {
+			return ErrWalletQuotaLimitExceeded
+		}
+		quota := decimal.NewFromInt(int64(value))
+		switch mode {
+		case "add":
+			quota = decimal.NewFromInt(int64(user.Quota)).Add(quota)
+		case "subtract":
+			quota = decimal.NewFromInt(int64(user.Quota)).Sub(quota)
+		}
+		after, err := common.WalletQuotaFromDecimalStrict(quota)
+		if err != nil {
+			return ErrWalletQuotaLimitExceeded
+		}
+		// An unchanged override is a successful operation, including on MySQL
+		// configurations that count only changed rows in RowsAffected.
+		if after != user.Quota {
+			result := tx.Model(&User{}).Where("id = ?", userID).Update("quota", after)
+			if result.Error != nil {
+				return result.Error
+			}
+			if result.RowsAffected != 1 {
+				return gorm.ErrRecordNotFound
+			}
+		}
+		adjustment = UserQuotaAdjustment{UserID: user.Id, Username: user.Username, Before: user.Quota, After: after}
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	// Apply only the committed difference, preserving outstanding reservations.
+	// Both balances are bounded above, so their difference fits in int64.
+	delta := int64(adjustment.After) - int64(adjustment.Before)
+	if delta != 0 {
+		if err := cacheIncrUserQuota(userID, delta); err != nil {
+			common.SysError(fmt.Sprintf("failed to sync manual quota adjustment for user %d: %s", userID, err))
+		}
+	}
+	return &adjustment, nil
+}