88 lines
2.8 KiB
Go
88 lines
2.8 KiB
Go
package model
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"math"
|
|
|
|
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
|
"github.com/QuantumNous/new-api/types"
|
|
"github.com/gin-gonic/gin"
|
|
"gorm.io/gorm/clause"
|
|
)
|
|
|
|
// Stored independently from User and UserSetting to prevent profile mass assignment.
|
|
type PersonalGroupRatio struct {
|
|
UserID int `gorm:"primaryKey;autoIncrement:false"`
|
|
Ratios string `gorm:"type:text;not null"`
|
|
}
|
|
|
|
const personalRatioContextKey = "personal_group_ratios"
|
|
|
|
func GetPersonalGroupRatios(ctx context.Context, userID int) (map[string]float64, error) {
|
|
var rows []PersonalGroupRatio
|
|
if err := DB.WithContext(ctx).Where("user_id = ?", userID).Find(&rows).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
result := map[string]float64{}
|
|
if len(rows) != 0 {
|
|
if err := json.Unmarshal([]byte(rows[0].Ratios), &result); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func SavePersonalGroupRatios(ctx context.Context, userID int, ratios map[string]float64) error {
|
|
for group, ratio := range ratios {
|
|
if !ratio_setting.ContainsGroupRatio(group) || group == "auto" || ratio < 0 || math.IsNaN(ratio) || math.IsInf(ratio, 0) {
|
|
return fmt.Errorf("invalid group or ratio: %s", group)
|
|
}
|
|
}
|
|
if ratios == nil {
|
|
ratios = map[string]float64{}
|
|
}
|
|
body, err := json.Marshal(ratios)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
row := PersonalGroupRatio{UserID: userID, Ratios: string(body)}
|
|
if err = DB.WithContext(ctx).Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "user_id"}}, DoUpdates: clause.AssignmentColumns([]string{"ratios"})}).Create(&row).Error; err != nil {
|
|
return err
|
|
}
|
|
// Pricing is loaded from the database at request admission, never from a stale
|
|
// process/Redis cache. The request snapshot survives retries and price updates.
|
|
return invalidateUserCache(userID)
|
|
}
|
|
|
|
func LoadPersonalGroupRatios(c *gin.Context, userID int) error {
|
|
ratios, err := GetPersonalGroupRatios(c.Request.Context(), userID)
|
|
if err == nil {
|
|
c.Set(personalRatioContextKey, ratios)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func EffectiveGroupRatio(c *gin.Context, userGroup, group string) types.GroupRatioInfo {
|
|
result := types.GroupRatioInfo{GroupRatio: ratio_setting.GetGroupRatio(group), GroupSpecialRatio: -1}
|
|
if ratio, ok := ratio_setting.GetGroupGroupRatio(userGroup, group); ok {
|
|
result.GroupRatio, result.GroupSpecialRatio, result.HasSpecialRatio = ratio, ratio, true
|
|
}
|
|
if value, ok := c.Get(personalRatioContextKey); ok {
|
|
if ratios, ok := value.(map[string]float64); ok {
|
|
if ratio, ok := ratios[group]; ok {
|
|
result.GroupRatio, result.GroupSpecialRatio, result.HasSpecialRatio = ratio, ratio, true
|
|
}
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
func HasPersonalGroupRatio(c *gin.Context, group string) bool {
|
|
value, _ := c.Get(personalRatioContextKey)
|
|
ratios, _ := value.(map[string]float64)
|
|
_, ok := ratios[group]
|
|
return ok
|
|
}
|