Parcourir la source

chore: cache username

CaIon il y a 2 ans
Parent
commit
f07b9f8ab2
4 fichiers modifiés avec 40 ajouts et 8 suppressions
  1. 18 0
      model/cache.go
  2. 3 2
      model/log.go
  3. 16 3
      model/usedata.go
  4. 3 3
      model/user.go

+ 18 - 0
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)

+ 3 - 2
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,

+ 16 - 3
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表中查询数据

+ 3 - 3
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
 }