package model import ( "encoding/json" "fmt" "one-api/common" "one-api/common/logger" "strconv" "time" ) var ( TokenCacheSeconds = 0 ) func CacheGetTokenByKey(key string) (*Token, error) { keyCol := "`key`" if common.UsingPostgreSQL { keyCol = `"key"` } var token Token if !common.RedisEnabled { err := DB.Where(keyCol+" = ?", key).First(&token).Error return &token, err } tokenObjectString, err := common.RedisGet(fmt.Sprintf("token:%s", key)) if err != nil { err := DB.Where(keyCol+" = ?", key).First(&token).Error if err != nil { return nil, err } jsonBytes, err := json.Marshal(token) if err != nil { return nil, err } err = common.RedisSet(fmt.Sprintf("token:%s", key), string(jsonBytes), time.Duration(TokenCacheSeconds)*time.Second) if err != nil { logger.SysError("Redis set token error: " + err.Error()) } return &token, nil } err = json.Unmarshal([]byte(tokenObjectString), &token) return &token, err } func CacheGetUserGroup(id int) (group string, err error) { if !common.RedisEnabled { return GetUserGroup(id) } group, err = common.RedisGet(fmt.Sprintf("user_group:%d", id)) if err != nil { group, err = GetUserGroup(id) if err != nil { return "", err } err = common.RedisSet(fmt.Sprintf("user_group:%d", id), group, time.Duration(TokenCacheSeconds)*time.Second) if err != nil { logger.SysError("Redis set user group error: " + err.Error()) } } return group, err } func CacheGetUserQuota(id int) (quota int, err error) { if !common.RedisEnabled { return GetUserQuota(id) } quotaString, err := common.RedisGet(fmt.Sprintf("user_quota:%d", id)) if err != nil { quota, err = GetUserQuota(id) if err != nil { return 0, err } err = common.RedisSet(fmt.Sprintf("user_quota:%d", id), fmt.Sprintf("%d", quota), time.Duration(TokenCacheSeconds)*time.Second) if err != nil { logger.SysError("Redis set user quota error: " + err.Error()) } return quota, err } quota, err = strconv.Atoi(quotaString) return quota, err } func CacheUpdateUserQuota(id int) error { if !common.RedisEnabled { return nil } quota, err := GetUserQuota(id) if err != nil { return err } err = common.RedisSet(fmt.Sprintf("user_quota:%d", id), fmt.Sprintf("%d", quota), time.Duration(TokenCacheSeconds)*time.Second) return err } func CacheDecreaseUserQuota(id int, quota int) error { if !common.RedisEnabled { return nil } err := common.RedisDecrease(fmt.Sprintf("user_quota:%d", id), int64(quota)) return err } func CacheIsUserEnabled(userId int) (bool, error) { if !common.RedisEnabled { return IsUserEnabled(userId) } enabled, err := common.RedisGet(fmt.Sprintf("user_enabled:%d", userId)) if err == nil { return enabled == "1", nil } userEnabled, err := IsUserEnabled(userId) if err != nil { return false, err } enabled = "0" if userEnabled { enabled = "1" } err = common.RedisSet(fmt.Sprintf("user_enabled:%d", userId), enabled, time.Duration(TokenCacheSeconds)*time.Second) if err != nil { logger.SysError("Redis set user enabled error: " + err.Error()) } return userEnabled, err }