diff --git a/controller/relay.go b/controller/relay.go index af214e7768..595cc29871 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -137,7 +137,7 @@ func relayHelper(c *gin.Context) error { ratio = common.RatioGPT3dot5 } quota = int(float64(quota) * ratio) - err := model.ConsumeTokenQuota(tokenId, quota) + err := model.DecreaseTokenQuota(tokenId, quota) if err != nil { common.SysError("Error consuming token remain quota: " + err.Error()) } diff --git a/model/redemption.go b/model/redemption.go index 838d7ecd0a..d499108a21 100644 --- a/model/redemption.go +++ b/model/redemption.go @@ -55,7 +55,7 @@ func Redeem(key string, tokenId int) (quota int, err error) { if redemption.Status != common.RedemptionCodeStatusEnabled { return 0, errors.New("该兑换码已被使用") } - err = TopUpTokenQuota(tokenId, redemption.Quota) + err = IncreaseTokenQuota(tokenId, redemption.Quota) if err != nil { return 0, err } diff --git a/model/token.go b/model/token.go index ff1806afc6..64e5c4e1cf 100644 --- a/model/token.go +++ b/model/token.go @@ -116,15 +116,26 @@ func DeleteTokenById(id int, userId int) (err error) { if err != nil { return err } + quota := token.RemainQuota + if quota != 0 { + if quota > 0 { + err = IncreaseUserQuota(userId, quota) + } else { + err = DecreaseUserQuota(userId, quota) + } + } + if err != nil { + return err + } return token.Delete() } -func ConsumeTokenQuota(id int, quota int) (err error) { - err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota - ?", quota)).Error +func IncreaseTokenQuota(id int, quota int) (err error) { + err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota + ?", quota)).Error return err } -func TopUpTokenQuota(id int, quota int) (err error) { - err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota + ?", quota)).Error +func DecreaseTokenQuota(id int, quota int) (err error) { + err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota - ?", quota)).Error return err } diff --git a/model/user.go b/model/user.go index bc2348b938..3a8f2313d0 100644 --- a/model/user.go +++ b/model/user.go @@ -225,6 +225,11 @@ func GetUserQuota(id int) (quota int, err error) { return quota, err } +func IncreaseUserQuota(id int, quota int) (err error) { + err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error + return err +} + func DecreaseUserQuota(id int, quota int) (err error) { err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error return err