fix: return quota to user when delete token (close #37)

songquanpeng · May 4, 2023 · 331177d · 331177d
1 parent 4fed003
commit 331177d
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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