75 lines
2.0 KiB
Go
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
|
|
})
|
|
}
|