diff --git a/model/cache.go b/model/cache.go index 975cedd..995d2cd 100644 --- a/model/cache.go +++ b/model/cache.go @@ -68,6 +68,24 @@ func CacheGetUserGroup(id int) (group string, err error) { return group, err } +func CacheGetUsername(id int) (username string, err error) { + if !common.RedisEnabled { + return GetUsernameById(id) + } + username, err = common.RedisGet(fmt.Sprintf("user_name:%d", id)) + if err != nil { + username, err = GetUserGroup(id) + if err != nil { + return "", err + } + err = common.RedisSet(fmt.Sprintf("user_name:%d", id), username, time.Duration(UserId2GroupCacheSeconds)*time.Second) + if err != nil { + common.SysError("Redis set user group error: " + err.Error()) + } + } + return username, err +} + func CacheGetUserQuota(id int) (quota int, err error) { if !common.RedisEnabled { return GetUserQuota(id) diff --git a/model/log.go b/model/log.go index 85de9d3..f96191f 100644 --- a/model/log.go +++ b/model/log.go @@ -41,9 +41,10 @@ func RecordLog(userId int, logType int, content string) { if logType == LogTypeConsume && !common.LogConsumeEnabled { return } + username, _ := CacheGetUsername(userId) log := &Log{ UserId: userId, - Username: GetUsernameById(userId), + Username: username, CreatedAt: common.GetTimestamp(), Type: logType, Content: content, @@ -59,7 +60,7 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke if !common.LogConsumeEnabled { return } - username := GetUsernameById(userId) + username, _ := CacheGetUsername(userId) log := &Log{ UserId: userId, Username: username, diff --git a/model/usedata.go b/model/usedata.go index aca3a36..9ca623d 100644 --- a/model/usedata.go +++ b/model/usedata.go @@ -2,6 +2,7 @@ package model import ( "fmt" + "gorm.io/gorm" "one-api/common" "sync" "time" @@ -78,9 +79,10 @@ func SaveQuotaDataCache() { DB.Table("quota_data").Where("user_id = ? and username = ? and model_name = ? and created_at = ?", quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.CreatedAt).First(quotaDataDB) if quotaDataDB.Id > 0 { - quotaDataDB.Count += quotaData.Count - quotaDataDB.Quota += quotaData.Quota - DB.Table("quota_data").Save(quotaDataDB) + //quotaDataDB.Count += quotaData.Count + //quotaDataDB.Quota += quotaData.Quota + //DB.Table("quota_data").Save(quotaDataDB) + increaseQuotaData(quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.Count, quotaData.Quota, quotaData.CreatedAt) } else { DB.Table("quota_data").Create(quotaData) } @@ -89,6 +91,17 @@ func SaveQuotaDataCache() { common.SysLog(fmt.Sprintf("保存数据看板数据成功,共保存%d条数据", size)) } +func increaseQuotaData(userId int, username string, modelName string, count int, quota int, createdAt int64) { + err := DB.Table("quota_data").Where("user_id = ? and username = ? and model_name = ? and created_at = ?", + userId, username, modelName, createdAt).Updates(map[string]interface{}{ + "count": gorm.Expr("count + ?", count), + "quota": gorm.Expr("quota + ?", quota), + }).Error + if err != nil { + common.SysLog(fmt.Sprintf("increaseQuotaData error: %s", err)) + } +} + func GetQuotaDataByUsername(username string, startTime int64, endTime int64) (quotaData []*QuotaData, err error) { var quotaDatas []*QuotaData // 从quota_data表中查询数据 diff --git a/model/user.go b/model/user.go index 30d09e5..360afc3 100644 --- a/model/user.go +++ b/model/user.go @@ -452,7 +452,7 @@ func updateUserRequestCount(id int, count int) { } } -func GetUsernameById(id int) (username string) { - DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username) - return username +func GetUsernameById(id int) (username string, err error) { + err = DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username).Error + return username, err }