Files
new-api-custom/model/group_metadata.go
T

75 lines
2.0 KiB
Go

package model
import (
"context"
"encoding/json"
"fmt"
"strings"
"unicode/utf8"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
type GroupMetadata struct {
GroupName string `gorm:"primaryKey;size:191"`
Value string `gorm:"type:text;not null"`
}
type GroupMetadataValue struct {
Tags []string `json:"tags"`
HideRecentDegraded bool `json:"hide_recent_degraded"`
}
func GetGroupMetadata(ctx context.Context) (map[string]GroupMetadataValue, error) {
var rows []GroupMetadata
if err := DB.WithContext(ctx).Find(&rows).Error; err != nil {
return nil, err
}
result := map[string]GroupMetadataValue{}
for _, row := range rows {
var value GroupMetadataValue
if err := json.Unmarshal([]byte(row.Value), &value); err != nil {
return nil, err
}
result[row.GroupName] = value
}
return result, nil
}
// Updates only submitted groups; absent groups keep their existing metadata.
func SaveGroupMetadata(ctx context.Context, groups map[string]GroupMetadataValue) error {
rows := []GroupMetadata{}
for group, value := range groups {
if !ratio_setting.ContainsGroupRatio(group) || group == "auto" {
return fmt.Errorf("invalid group: %s", group)
}
tags, seen := []string{}, map[string]bool{}
for _, tag := range value.Tags {
tag = strings.TrimSpace(tag)
if tag == "" || seen[tag] {
continue
}
if utf8.RuneCountInString(tag) > 40 || len(tags) >= 20 {
return fmt.Errorf("too many or too long group tags")
}
seen[tag] = true
tags = append(tags, tag)
}
value.Tags = tags
body, err := json.Marshal(value)
if err != nil {
return err
}
rows = append(rows, GroupMetadata{GroupName: group, Value: string(body)})
}
return DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
for _, row := range rows {
if err := tx.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "group_name"}}, DoUpdates: clause.AssignmentColumns([]string{"value"})}).Create(&row).Error; err != nil {
return err
}
}
return nil
})
}