From 0cdab80a6e05441d265646d297f22d22da4a4aeb Mon Sep 17 00:00:00 2001
From: JustSong <songquanpeng@foxmail.com>
Date: Fri, 16 Jun 2023 16:02:00 +0800
Subject: [PATCH] feat: record channel's used quota (close #137)

---
 controller/relay.go | 2 ++
 model/channel.go    | 9 +++++++++
 2 files changed, 11 insertions(+)

diff --git a/controller/relay.go b/controller/relay.go
index bb519258..8e7073f7 100644
--- a/controller/relay.go
+++ b/controller/relay.go
@@ -262,6 +262,8 @@ func relayHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 			userId := c.GetInt("id")
 			model.RecordLog(userId, model.LogTypeConsume, fmt.Sprintf("使用模型 %s 消耗 %d 点额度（模型倍率 %.2f，分组倍率 %.2f）", textRequest.Model, quota, modelRatio, groupRatio))
 			model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
+			channelId := c.GetInt("channel_id")
+			model.UpdateChannelUsedQuota(channelId, quota)
 		}
 	}()
 
diff --git a/model/channel.go b/model/channel.go
index fdc89eb9..e53efc20 100644
--- a/model/channel.go
+++ b/model/channel.go
@@ -1,6 +1,7 @@
 package model
 
 import (
+	"gorm.io/gorm"
 	"one-api/common"
 )
 
@@ -20,6 +21,7 @@ type Channel struct {
 	BalanceUpdatedTime int64   `json:"balance_updated_time" gorm:"bigint"`
 	Models             string  `json:"models"`
 	Group              string  `json:"group" gorm:"type:varchar(32);default:'default'"`
+	UsedQuota          int64   `json:"used_quota" gorm:"bigint;default:0"`
 }
 
 func GetAllChannels(startIdx int, num int, selectAll bool) ([]*Channel, error) {
@@ -136,3 +138,10 @@ func UpdateChannelStatusById(id int, status int) {
 		common.SysError("failed to update channel status: " + err.Error())
 	}
 }
+
+func UpdateChannelUsedQuota(id int, quota int) {
+	err := DB.Model(&Channel{}).Where("id = ?", id).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error
+	if err != nil {
+		common.SysError("failed to update channel used quota: " + err.Error())
+	}
+}