| package model
|
|
|
| import (
|
| "database/sql/driver"
|
| "encoding/json"
|
| "errors"
|
| "fmt"
|
| "math/rand"
|
| "strings"
|
| "sync"
|
|
|
| "github.com/QuantumNous/new-api/common"
|
| "github.com/QuantumNous/new-api/constant"
|
| "github.com/QuantumNous/new-api/dto"
|
| "github.com/QuantumNous/new-api/types"
|
|
|
| "github.com/samber/lo"
|
| "gorm.io/gorm"
|
| )
|
|
|
| type Channel struct {
|
| Id int `json:"id"`
|
| Type int `json:"type" gorm:"default:0"`
|
| Key string `json:"key" gorm:"not null"`
|
| OpenAIOrganization *string `json:"openai_organization"`
|
| TestModel *string `json:"test_model"`
|
| Status int `json:"status" gorm:"default:1"`
|
| Name string `json:"name" gorm:"index"`
|
| Weight *uint `json:"weight" gorm:"default:0"`
|
| CreatedTime int64 `json:"created_time" gorm:"bigint"`
|
| TestTime int64 `json:"test_time" gorm:"bigint"`
|
| ResponseTime int `json:"response_time"`
|
| BaseURL *string `json:"base_url" gorm:"column:base_url;default:''"`
|
| Other string `json:"other"`
|
| Balance float64 `json:"balance"`
|
| BalanceUpdatedTime int64 `json:"balance_updated_time" gorm:"bigint"`
|
| Models string `json:"models"`
|
| Group string `json:"group" gorm:"type:varchar(64);default:'default'"`
|
| UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"`
|
| ModelMapping *string `json:"model_mapping" gorm:"type:text"`
|
|
|
| StatusCodeMapping *string `json:"status_code_mapping" gorm:"type:varchar(1024);default:''"`
|
| Priority *int64 `json:"priority" gorm:"bigint;default:0"`
|
| AutoBan *int `json:"auto_ban" gorm:"default:1"`
|
| OtherInfo string `json:"other_info"`
|
| Tag *string `json:"tag" gorm:"index"`
|
| Setting *string `json:"setting" gorm:"type:text"`
|
| ParamOverride *string `json:"param_override" gorm:"type:text"`
|
| HeaderOverride *string `json:"header_override" gorm:"type:text"`
|
| Remark *string `json:"remark" gorm:"type:varchar(255)" validate:"max=255"`
|
|
|
| ChannelInfo ChannelInfo `json:"channel_info" gorm:"type:json"`
|
|
|
| OtherSettings string `json:"settings" gorm:"column:settings"`
|
|
|
|
|
| Keys []string `json:"-" gorm:"-"`
|
| }
|
|
|
| type ChannelInfo struct {
|
| IsMultiKey bool `json:"is_multi_key"`
|
| MultiKeySize int `json:"multi_key_size"`
|
| MultiKeyStatusList map[int]int `json:"multi_key_status_list"`
|
| MultiKeyDisabledReason map[int]string `json:"multi_key_disabled_reason,omitempty"`
|
| MultiKeyDisabledTime map[int]int64 `json:"multi_key_disabled_time,omitempty"`
|
| MultiKeyPollingIndex int `json:"multi_key_polling_index"`
|
| MultiKeyMode constant.MultiKeyMode `json:"multi_key_mode"`
|
| }
|
|
|
|
|
| func (c ChannelInfo) Value() (driver.Value, error) {
|
| return common.Marshal(&c)
|
| }
|
|
|
|
|
| func (c *ChannelInfo) Scan(value interface{}) error {
|
| bytesValue, _ := value.([]byte)
|
| return common.Unmarshal(bytesValue, c)
|
| }
|
|
|
| func (channel *Channel) GetKeys() []string {
|
| if channel.Key == "" {
|
| return []string{}
|
| }
|
| if len(channel.Keys) > 0 {
|
| return channel.Keys
|
| }
|
| trimmed := strings.TrimSpace(channel.Key)
|
|
|
| if strings.HasPrefix(trimmed, "[") {
|
| var arr []json.RawMessage
|
| if err := common.Unmarshal([]byte(trimmed), &arr); err == nil {
|
| res := make([]string, len(arr))
|
| for i, v := range arr {
|
| res[i] = string(v)
|
| }
|
| return res
|
| }
|
| }
|
|
|
| keys := strings.Split(strings.Trim(channel.Key, "\n"), "\n")
|
| return keys
|
| }
|
|
|
| func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) {
|
|
|
| if !channel.ChannelInfo.IsMultiKey {
|
| return channel.Key, 0, nil
|
| }
|
|
|
|
|
| keys := channel.GetKeys()
|
| if len(keys) == 0 {
|
|
|
| return "", 0, types.NewError(errors.New("no keys available"), types.ErrorCodeChannelNoAvailableKey)
|
| }
|
|
|
| lock := GetChannelPollingLock(channel.Id)
|
| lock.Lock()
|
| defer lock.Unlock()
|
|
|
| statusList := channel.ChannelInfo.MultiKeyStatusList
|
|
|
| getStatus := func(idx int) int {
|
| if statusList == nil {
|
| return common.ChannelStatusEnabled
|
| }
|
| if status, ok := statusList[idx]; ok {
|
| return status
|
| }
|
| return common.ChannelStatusEnabled
|
| }
|
|
|
|
|
| enabledIdx := make([]int, 0, len(keys))
|
| for i := range keys {
|
| if getStatus(i) == common.ChannelStatusEnabled {
|
| enabledIdx = append(enabledIdx, i)
|
| }
|
| }
|
|
|
|
|
|
|
| if len(enabledIdx) == 0 {
|
| return "", 0, types.NewError(errors.New("no enabled keys"), types.ErrorCodeChannelNoAvailableKey)
|
| }
|
|
|
| switch channel.ChannelInfo.MultiKeyMode {
|
| case constant.MultiKeyModeRandom:
|
|
|
| selectedIdx := enabledIdx[rand.Intn(len(enabledIdx))]
|
| return keys[selectedIdx], selectedIdx, nil
|
| case constant.MultiKeyModePolling:
|
|
|
|
|
| channelInfo, err := CacheGetChannelInfo(channel.Id)
|
| if err != nil {
|
| return "", 0, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
|
| }
|
|
|
| defer func() {
|
| if common.DebugEnabled {
|
| println(fmt.Sprintf("channel %d polling index: %d", channel.Id, channel.ChannelInfo.MultiKeyPollingIndex))
|
| }
|
| if !common.MemoryCacheEnabled {
|
| _ = channel.SaveChannelInfo()
|
| } else {
|
|
|
| }
|
| }()
|
|
|
| start := channelInfo.MultiKeyPollingIndex
|
| if start < 0 || start >= len(keys) {
|
| start = 0
|
| }
|
| for i := 0; i < len(keys); i++ {
|
| idx := (start + i) % len(keys)
|
| if getStatus(idx) == common.ChannelStatusEnabled {
|
|
|
| channel.ChannelInfo.MultiKeyPollingIndex = (idx + 1) % len(keys)
|
| return keys[idx], idx, nil
|
| }
|
| }
|
|
|
| return keys[enabledIdx[0]], enabledIdx[0], nil
|
| default:
|
|
|
| return keys[enabledIdx[0]], enabledIdx[0], nil
|
| }
|
| }
|
|
|
| func (channel *Channel) SaveChannelInfo() error {
|
| return DB.Model(channel).Update("channel_info", channel.ChannelInfo).Error
|
| }
|
|
|
| func (channel *Channel) GetModels() []string {
|
| if channel.Models == "" {
|
| return []string{}
|
| }
|
| return strings.Split(strings.Trim(channel.Models, ","), ",")
|
| }
|
|
|
| func (channel *Channel) GetGroups() []string {
|
| if channel.Group == "" {
|
| return []string{}
|
| }
|
| groups := strings.Split(strings.Trim(channel.Group, ","), ",")
|
| for i, group := range groups {
|
| groups[i] = strings.TrimSpace(group)
|
| }
|
| return groups
|
| }
|
|
|
| func (channel *Channel) GetOtherInfo() map[string]interface{} {
|
| otherInfo := make(map[string]interface{})
|
| if channel.OtherInfo != "" {
|
| err := common.Unmarshal([]byte(channel.OtherInfo), &otherInfo)
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to unmarshal other info: channel_id=%d, tag=%s, name=%s, error=%v", channel.Id, channel.GetTag(), channel.Name, err))
|
| }
|
| }
|
| return otherInfo
|
| }
|
|
|
| func (channel *Channel) SetOtherInfo(otherInfo map[string]interface{}) {
|
| otherInfoBytes, err := json.Marshal(otherInfo)
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to marshal other info: channel_id=%d, tag=%s, name=%s, error=%v", channel.Id, channel.GetTag(), channel.Name, err))
|
| return
|
| }
|
| channel.OtherInfo = string(otherInfoBytes)
|
| }
|
|
|
| func (channel *Channel) GetTag() string {
|
| if channel.Tag == nil {
|
| return ""
|
| }
|
| return *channel.Tag
|
| }
|
|
|
| func (channel *Channel) SetTag(tag string) {
|
| channel.Tag = &tag
|
| }
|
|
|
| func (channel *Channel) GetAutoBan() bool {
|
| if channel.AutoBan == nil {
|
| return false
|
| }
|
| return *channel.AutoBan == 1
|
| }
|
|
|
| func (channel *Channel) Save() error {
|
| return DB.Save(channel).Error
|
| }
|
|
|
| func (channel *Channel) SaveWithoutKey() error {
|
| if channel.Id == 0 {
|
| return errors.New("channel ID is 0")
|
| }
|
| return DB.Omit("key").Save(channel).Error
|
| }
|
|
|
| func GetAllChannels(startIdx int, num int, selectAll bool, idSort bool) ([]*Channel, error) {
|
| var channels []*Channel
|
| var err error
|
| order := "priority desc"
|
| if idSort {
|
| order = "id desc"
|
| }
|
| if selectAll {
|
| err = DB.Order(order).Find(&channels).Error
|
| } else {
|
| err = DB.Order(order).Limit(num).Offset(startIdx).Omit("key").Find(&channels).Error
|
| }
|
| return channels, err
|
| }
|
|
|
| func GetChannelsByTag(tag string, idSort bool, selectAll bool) ([]*Channel, error) {
|
| var channels []*Channel
|
| order := "priority desc"
|
| if idSort {
|
| order = "id desc"
|
| }
|
| query := DB.Where("tag = ?", tag).Order(order)
|
| if !selectAll {
|
| query = query.Omit("key")
|
| }
|
| err := query.Find(&channels).Error
|
| return channels, err
|
| }
|
|
|
| func SearchChannels(keyword string, group string, model string, idSort bool) ([]*Channel, error) {
|
| var channels []*Channel
|
| modelsCol := "`models`"
|
|
|
|
|
| if common.UsingPostgreSQL {
|
| modelsCol = `"models"`
|
| }
|
|
|
| baseURLCol := "`base_url`"
|
|
|
| if common.UsingPostgreSQL {
|
| baseURLCol = `"base_url"`
|
| }
|
|
|
| order := "priority desc"
|
| if idSort {
|
| order = "id desc"
|
| }
|
|
|
|
|
| baseQuery := DB.Model(&Channel{}).Omit("key")
|
|
|
|
|
| var whereClause string
|
| var args []interface{}
|
| if group != "" && group != "null" {
|
| var groupCondition string
|
| if common.UsingMySQL {
|
| groupCondition = `CONCAT(',', ` + commonGroupCol + `, ',') LIKE ?`
|
| } else {
|
|
|
| groupCondition = `(',' || ` + commonGroupCol + ` || ',') LIKE ?`
|
| }
|
| whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + ` LIKE ? AND ` + groupCondition
|
| args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%", "%,"+group+",%")
|
| } else {
|
| whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + " LIKE ?"
|
| args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%")
|
| }
|
|
|
|
|
| err := baseQuery.Where(whereClause, args...).Order(order).Find(&channels).Error
|
| if err != nil {
|
| return nil, err
|
| }
|
| return channels, nil
|
| }
|
|
|
| func GetChannelById(id int, selectAll bool) (*Channel, error) {
|
| channel := &Channel{Id: id}
|
| var err error = nil
|
| if selectAll {
|
| err = DB.First(channel, "id = ?", id).Error
|
| } else {
|
| err = DB.Omit("key").First(channel, "id = ?", id).Error
|
| }
|
| if err != nil {
|
| return nil, err
|
| }
|
| if channel == nil {
|
| return nil, errors.New("channel not found")
|
| }
|
| return channel, nil
|
| }
|
|
|
| func BatchInsertChannels(channels []Channel) error {
|
| if len(channels) == 0 {
|
| return nil
|
| }
|
| tx := DB.Begin()
|
| if tx.Error != nil {
|
| return tx.Error
|
| }
|
| defer func() {
|
| if r := recover(); r != nil {
|
| tx.Rollback()
|
| }
|
| }()
|
|
|
| for _, chunk := range lo.Chunk(channels, 50) {
|
| if err := tx.Create(&chunk).Error; err != nil {
|
| tx.Rollback()
|
| return err
|
| }
|
| for _, channel_ := range chunk {
|
| if err := channel_.AddAbilities(tx); err != nil {
|
| tx.Rollback()
|
| return err
|
| }
|
| }
|
| }
|
| return tx.Commit().Error
|
| }
|
|
|
| func BatchDeleteChannels(ids []int) error {
|
| if len(ids) == 0 {
|
| return nil
|
| }
|
|
|
| tx := DB.Begin()
|
| if tx.Error != nil {
|
| return tx.Error
|
| }
|
| for _, chunk := range lo.Chunk(ids, 200) {
|
| if err := tx.Where("id in (?)", chunk).Delete(&Channel{}).Error; err != nil {
|
| tx.Rollback()
|
| return err
|
| }
|
| if err := tx.Where("channel_id in (?)", chunk).Delete(&Ability{}).Error; err != nil {
|
| tx.Rollback()
|
| return err
|
| }
|
| }
|
| return tx.Commit().Error
|
| }
|
|
|
| func (channel *Channel) GetPriority() int64 {
|
| if channel.Priority == nil {
|
| return 0
|
| }
|
| return *channel.Priority
|
| }
|
|
|
| func (channel *Channel) GetWeight() int {
|
| if channel.Weight == nil {
|
| return 0
|
| }
|
| return int(*channel.Weight)
|
| }
|
|
|
| func (channel *Channel) GetBaseURL() string {
|
| if channel.BaseURL == nil {
|
| return ""
|
| }
|
| url := *channel.BaseURL
|
| if url == "" {
|
| url = constant.ChannelBaseURLs[channel.Type]
|
| }
|
| return url
|
| }
|
|
|
| func (channel *Channel) GetModelMapping() string {
|
| if channel.ModelMapping == nil {
|
| return ""
|
| }
|
| return *channel.ModelMapping
|
| }
|
|
|
| func (channel *Channel) GetStatusCodeMapping() string {
|
| if channel.StatusCodeMapping == nil {
|
| return ""
|
| }
|
| return *channel.StatusCodeMapping
|
| }
|
|
|
| func (channel *Channel) Insert() error {
|
| var err error
|
| err = DB.Create(channel).Error
|
| if err != nil {
|
| return err
|
| }
|
| err = channel.AddAbilities(nil)
|
| return err
|
| }
|
|
|
| func (channel *Channel) Update() error {
|
|
|
| if channel.ChannelInfo.IsMultiKey {
|
| var keyStr string
|
| if channel.Key != "" {
|
| keyStr = channel.Key
|
| } else {
|
|
|
| if existing, err := GetChannelById(channel.Id, true); err == nil {
|
| keyStr = existing.Key
|
| }
|
| }
|
|
|
| keys := []string{}
|
| if keyStr != "" {
|
| trimmed := strings.TrimSpace(keyStr)
|
| if strings.HasPrefix(trimmed, "[") {
|
| var arr []json.RawMessage
|
| if err := common.Unmarshal([]byte(trimmed), &arr); err == nil {
|
| keys = make([]string, len(arr))
|
| for i, v := range arr {
|
| keys[i] = string(v)
|
| }
|
| }
|
| }
|
| if len(keys) == 0 {
|
| keys = strings.Split(strings.Trim(keyStr, "\n"), "\n")
|
| }
|
| }
|
| channel.ChannelInfo.MultiKeySize = len(keys)
|
|
|
| if channel.ChannelInfo.MultiKeyStatusList != nil {
|
| for idx := range channel.ChannelInfo.MultiKeyStatusList {
|
| if idx >= channel.ChannelInfo.MultiKeySize {
|
| delete(channel.ChannelInfo.MultiKeyStatusList, idx)
|
| }
|
| }
|
| }
|
| }
|
| var err error
|
| err = DB.Model(channel).Updates(channel).Error
|
| if err != nil {
|
| return err
|
| }
|
| DB.Model(channel).First(channel, "id = ?", channel.Id)
|
| err = channel.UpdateAbilities(nil)
|
| return err
|
| }
|
|
|
| func (channel *Channel) UpdateResponseTime(responseTime int64) {
|
| err := DB.Model(channel).Select("response_time", "test_time").Updates(Channel{
|
| TestTime: common.GetTimestamp(),
|
| ResponseTime: int(responseTime),
|
| }).Error
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to update response time: channel_id=%d, error=%v", channel.Id, err))
|
| }
|
| }
|
|
|
| func (channel *Channel) UpdateBalance(balance float64) {
|
| err := DB.Model(channel).Select("balance_updated_time", "balance").Updates(Channel{
|
| BalanceUpdatedTime: common.GetTimestamp(),
|
| Balance: balance,
|
| }).Error
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to update balance: channel_id=%d, error=%v", channel.Id, err))
|
| }
|
| }
|
|
|
| func (channel *Channel) Delete() error {
|
| var err error
|
| err = DB.Delete(channel).Error
|
| if err != nil {
|
| return err
|
| }
|
| err = channel.DeleteAbilities()
|
| return err
|
| }
|
|
|
| var channelStatusLock sync.Mutex
|
|
|
|
|
| var channelPollingLocks sync.Map
|
|
|
|
|
| func GetChannelPollingLock(channelId int) *sync.Mutex {
|
| if lock, exists := channelPollingLocks.Load(channelId); exists {
|
| return lock.(*sync.Mutex)
|
| }
|
|
|
| newLock := &sync.Mutex{}
|
| actual, _ := channelPollingLocks.LoadOrStore(channelId, newLock)
|
| return actual.(*sync.Mutex)
|
| }
|
|
|
|
|
|
|
| func CleanupChannelPollingLocks() {
|
| var activeChannelIds []int
|
| DB.Model(&Channel{}).Pluck("id", &activeChannelIds)
|
|
|
| activeChannelSet := make(map[int]bool)
|
| for _, id := range activeChannelIds {
|
| activeChannelSet[id] = true
|
| }
|
|
|
| channelPollingLocks.Range(func(key, value interface{}) bool {
|
| channelId := key.(int)
|
| if !activeChannelSet[channelId] {
|
| channelPollingLocks.Delete(channelId)
|
| }
|
| return true
|
| })
|
| }
|
|
|
| func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int, reason string) {
|
| keys := channel.GetKeys()
|
| if len(keys) == 0 {
|
| channel.Status = status
|
| } else {
|
| var keyIndex int
|
| for i, key := range keys {
|
| if key == usingKey {
|
| keyIndex = i
|
| break
|
| }
|
| }
|
| if channel.ChannelInfo.MultiKeyStatusList == nil {
|
| channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
|
| }
|
| if status == common.ChannelStatusEnabled {
|
| delete(channel.ChannelInfo.MultiKeyStatusList, keyIndex)
|
| } else {
|
| channel.ChannelInfo.MultiKeyStatusList[keyIndex] = status
|
| if channel.ChannelInfo.MultiKeyDisabledReason == nil {
|
| channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string)
|
| }
|
| if channel.ChannelInfo.MultiKeyDisabledTime == nil {
|
| channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64)
|
| }
|
| channel.ChannelInfo.MultiKeyDisabledReason[keyIndex] = reason
|
| channel.ChannelInfo.MultiKeyDisabledTime[keyIndex] = common.GetTimestamp()
|
| }
|
| if len(channel.ChannelInfo.MultiKeyStatusList) >= channel.ChannelInfo.MultiKeySize {
|
| channel.Status = common.ChannelStatusAutoDisabled
|
| info := channel.GetOtherInfo()
|
| info["status_reason"] = "All keys are disabled"
|
| info["status_time"] = common.GetTimestamp()
|
| channel.SetOtherInfo(info)
|
| }
|
| }
|
| }
|
|
|
| func UpdateChannelStatus(channelId int, usingKey string, status int, reason string) bool {
|
| if common.MemoryCacheEnabled {
|
| channelStatusLock.Lock()
|
| defer channelStatusLock.Unlock()
|
|
|
| channelCache, _ := CacheGetChannel(channelId)
|
| if channelCache == nil {
|
| return false
|
| }
|
| if channelCache.ChannelInfo.IsMultiKey {
|
|
|
| pollingLock := GetChannelPollingLock(channelId)
|
| pollingLock.Lock()
|
|
|
| handlerMultiKeyUpdate(channelCache, usingKey, status, reason)
|
| pollingLock.Unlock()
|
|
|
|
|
| } else {
|
|
|
| if channelCache.Status == status {
|
| return false
|
| }
|
| CacheUpdateChannelStatus(channelId, status)
|
| }
|
| }
|
|
|
| shouldUpdateAbilities := false
|
| defer func() {
|
| if shouldUpdateAbilities {
|
| err := UpdateAbilityStatus(channelId, status == common.ChannelStatusEnabled)
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to update ability status: channel_id=%d, error=%v", channelId, err))
|
| }
|
| }
|
| }()
|
| channel, err := GetChannelById(channelId, true)
|
| if err != nil {
|
| return false
|
| } else {
|
| if channel.Status == status {
|
| return false
|
| }
|
|
|
| if channel.ChannelInfo.IsMultiKey {
|
| beforeStatus := channel.Status
|
|
|
| pollingLock := GetChannelPollingLock(channelId)
|
| pollingLock.Lock()
|
| handlerMultiKeyUpdate(channel, usingKey, status, reason)
|
| pollingLock.Unlock()
|
| if beforeStatus != channel.Status {
|
| shouldUpdateAbilities = true
|
| }
|
| } else {
|
| info := channel.GetOtherInfo()
|
| info["status_reason"] = reason
|
| info["status_time"] = common.GetTimestamp()
|
| channel.SetOtherInfo(info)
|
| channel.Status = status
|
| shouldUpdateAbilities = true
|
| }
|
| err = channel.SaveWithoutKey()
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to update channel status: channel_id=%d, status=%d, error=%v", channel.Id, status, err))
|
| return false
|
| }
|
| }
|
| return true
|
| }
|
|
|
| func EnableChannelByTag(tag string) error {
|
| err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusEnabled).Error
|
| if err != nil {
|
| return err
|
| }
|
| err = UpdateAbilityStatusByTag(tag, true)
|
| return err
|
| }
|
|
|
| func DisableChannelByTag(tag string) error {
|
| err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusManuallyDisabled).Error
|
| if err != nil {
|
| return err
|
| }
|
| err = UpdateAbilityStatusByTag(tag, false)
|
| return err
|
| }
|
|
|
| func EditChannelByTag(tag string, newTag *string, modelMapping *string, models *string, group *string, priority *int64, weight *uint, paramOverride *string, headerOverride *string) error {
|
| updateData := Channel{}
|
| shouldReCreateAbilities := false
|
| updatedTag := tag
|
|
|
| if newTag != nil && *newTag != tag {
|
| updateData.Tag = newTag
|
| updatedTag = *newTag
|
| }
|
| if modelMapping != nil && *modelMapping != "" {
|
| updateData.ModelMapping = modelMapping
|
| }
|
| if models != nil && *models != "" {
|
| shouldReCreateAbilities = true
|
| updateData.Models = *models
|
| }
|
| if group != nil && *group != "" {
|
| shouldReCreateAbilities = true
|
| updateData.Group = *group
|
| }
|
| if priority != nil {
|
| updateData.Priority = priority
|
| }
|
| if weight != nil {
|
| updateData.Weight = weight
|
| }
|
| if paramOverride != nil {
|
| updateData.ParamOverride = paramOverride
|
| }
|
| if headerOverride != nil {
|
| updateData.HeaderOverride = headerOverride
|
| }
|
|
|
| err := DB.Model(&Channel{}).Where("tag = ?", tag).Updates(updateData).Error
|
| if err != nil {
|
| return err
|
| }
|
| if shouldReCreateAbilities {
|
| channels, err := GetChannelsByTag(updatedTag, false, false)
|
| if err == nil {
|
| for _, channel := range channels {
|
| err = channel.UpdateAbilities(nil)
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to update abilities: channel_id=%d, tag=%s, error=%v", channel.Id, channel.GetTag(), err))
|
| }
|
| }
|
| }
|
| } else {
|
| err := UpdateAbilityByTag(tag, newTag, priority, weight)
|
| if err != nil {
|
| return err
|
| }
|
| }
|
| return nil
|
| }
|
|
|
| func UpdateChannelUsedQuota(id int, quota int) {
|
| if common.BatchUpdateEnabled {
|
| addNewRecord(BatchUpdateTypeChannelUsedQuota, id, quota)
|
| return
|
| }
|
| updateChannelUsedQuota(id, quota)
|
| }
|
|
|
| func updateChannelUsedQuota(id int, quota int) {
|
| err := DB.Model(&Channel{}).Where("id = ?", id).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to update channel used quota: channel_id=%d, delta_quota=%d, error=%v", id, quota, err))
|
| }
|
| }
|
|
|
| func DeleteChannelByStatus(status int64) (int64, error) {
|
| result := DB.Where("status = ?", status).Delete(&Channel{})
|
| return result.RowsAffected, result.Error
|
| }
|
|
|
| func DeleteDisabledChannel() (int64, error) {
|
| result := DB.Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).Delete(&Channel{})
|
| return result.RowsAffected, result.Error
|
| }
|
|
|
| func GetPaginatedTags(offset int, limit int) ([]*string, error) {
|
| var tags []*string
|
| err := DB.Model(&Channel{}).Select("DISTINCT tag").Where("tag != ''").Offset(offset).Limit(limit).Find(&tags).Error
|
| return tags, err
|
| }
|
|
|
| func SearchTags(keyword string, group string, model string, idSort bool) ([]*string, error) {
|
| var tags []*string
|
| modelsCol := "`models`"
|
|
|
|
|
| if common.UsingPostgreSQL {
|
| modelsCol = `"models"`
|
| }
|
|
|
| baseURLCol := "`base_url`"
|
|
|
| if common.UsingPostgreSQL {
|
| baseURLCol = `"base_url"`
|
| }
|
|
|
| order := "priority desc"
|
| if idSort {
|
| order = "id desc"
|
| }
|
|
|
|
|
| baseQuery := DB.Model(&Channel{}).Omit("key")
|
|
|
|
|
| var whereClause string
|
| var args []interface{}
|
| if group != "" && group != "null" {
|
| var groupCondition string
|
| if common.UsingMySQL {
|
| groupCondition = `CONCAT(',', ` + commonGroupCol + `, ',') LIKE ?`
|
| } else {
|
|
|
| groupCondition = `(',' || ` + commonGroupCol + ` || ',') LIKE ?`
|
| }
|
| whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + ` LIKE ? AND ` + groupCondition
|
| args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%", "%,"+group+",%")
|
| } else {
|
| whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + " LIKE ?"
|
| args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%")
|
| }
|
|
|
| subQuery := baseQuery.Where(whereClause, args...).
|
| Select("tag").
|
| Where("tag != ''").
|
| Order(order)
|
|
|
| err := DB.Table("(?) as sub", subQuery).
|
| Select("DISTINCT tag").
|
| Find(&tags).Error
|
|
|
| if err != nil {
|
| return nil, err
|
| }
|
|
|
| return tags, nil
|
| }
|
|
|
| func (channel *Channel) ValidateSettings() error {
|
| channelParams := &dto.ChannelSettings{}
|
| if channel.Setting != nil && *channel.Setting != "" {
|
| err := common.Unmarshal([]byte(*channel.Setting), channelParams)
|
| if err != nil {
|
| return err
|
| }
|
| }
|
| return nil
|
| }
|
|
|
| func (channel *Channel) GetSetting() dto.ChannelSettings {
|
| setting := dto.ChannelSettings{}
|
| if channel.Setting != nil && *channel.Setting != "" {
|
| err := common.Unmarshal([]byte(*channel.Setting), &setting)
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to unmarshal setting: channel_id=%d, error=%v", channel.Id, err))
|
| channel.Setting = nil
|
| _ = channel.Save()
|
| }
|
| }
|
| return setting
|
| }
|
|
|
| func (channel *Channel) SetSetting(setting dto.ChannelSettings) {
|
| settingBytes, err := common.Marshal(setting)
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to marshal setting: channel_id=%d, error=%v", channel.Id, err))
|
| return
|
| }
|
| channel.Setting = common.GetPointer[string](string(settingBytes))
|
| }
|
|
|
| func (channel *Channel) GetOtherSettings() dto.ChannelOtherSettings {
|
| setting := dto.ChannelOtherSettings{}
|
| if channel.OtherSettings != "" {
|
| err := common.UnmarshalJsonStr(channel.OtherSettings, &setting)
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to unmarshal setting: channel_id=%d, error=%v", channel.Id, err))
|
| channel.OtherSettings = "{}"
|
| _ = channel.Save()
|
| }
|
| }
|
| return setting
|
| }
|
|
|
| func (channel *Channel) SetOtherSettings(setting dto.ChannelOtherSettings) {
|
| settingBytes, err := common.Marshal(setting)
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to marshal setting: channel_id=%d, error=%v", channel.Id, err))
|
| return
|
| }
|
| channel.OtherSettings = string(settingBytes)
|
| }
|
|
|
| func (channel *Channel) GetParamOverride() map[string]interface{} {
|
| paramOverride := make(map[string]interface{})
|
| if channel.ParamOverride != nil && *channel.ParamOverride != "" {
|
| err := common.Unmarshal([]byte(*channel.ParamOverride), ¶mOverride)
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to unmarshal param override: channel_id=%d, error=%v", channel.Id, err))
|
| }
|
| }
|
| return paramOverride
|
| }
|
|
|
| func (channel *Channel) GetHeaderOverride() map[string]interface{} {
|
| headerOverride := make(map[string]interface{})
|
| if channel.HeaderOverride != nil && *channel.HeaderOverride != "" {
|
| err := common.Unmarshal([]byte(*channel.HeaderOverride), &headerOverride)
|
| if err != nil {
|
| common.SysLog(fmt.Sprintf("failed to unmarshal header override: channel_id=%d, error=%v", channel.Id, err))
|
| }
|
| }
|
| return headerOverride
|
| }
|
|
|
| func GetChannelsByIds(ids []int) ([]*Channel, error) {
|
| var channels []*Channel
|
| err := DB.Where("id in (?)", ids).Find(&channels).Error
|
| return channels, err
|
| }
|
|
|
| func BatchSetChannelTag(ids []int, tag *string) error {
|
|
|
| tx := DB.Begin()
|
| if tx.Error != nil {
|
| return tx.Error
|
| }
|
|
|
|
|
| err := tx.Model(&Channel{}).Where("id in (?)", ids).Update("tag", tag).Error
|
| if err != nil {
|
| tx.Rollback()
|
| return err
|
| }
|
|
|
|
|
| channels, err := GetChannelsByIds(ids)
|
| if err != nil {
|
| tx.Rollback()
|
| return err
|
| }
|
|
|
| for _, channel := range channels {
|
| err = channel.UpdateAbilities(tx)
|
| if err != nil {
|
| tx.Rollback()
|
| return err
|
| }
|
| }
|
|
|
|
|
| return tx.Commit().Error
|
| }
|
|
|
|
|
| func CountAllChannels() (int64, error) {
|
| var total int64
|
| err := DB.Model(&Channel{}).Count(&total).Error
|
| return total, err
|
| }
|
|
|
|
|
| func CountAllTags() (int64, error) {
|
| var total int64
|
| err := DB.Model(&Channel{}).Where("tag is not null AND tag != ''").Distinct("tag").Count(&total).Error
|
| return total, err
|
| }
|
|
|
|
|
| func GetChannelsByType(startIdx int, num int, idSort bool, channelType int) ([]*Channel, error) {
|
| var channels []*Channel
|
| order := "priority desc"
|
| if idSort {
|
| order = "id desc"
|
| }
|
| err := DB.Where("type = ?", channelType).Order(order).Limit(num).Offset(startIdx).Omit("key").Find(&channels).Error
|
| return channels, err
|
| }
|
|
|
|
|
| func CountChannelsByType(channelType int) (int64, error) {
|
| var count int64
|
| err := DB.Model(&Channel{}).Where("type = ?", channelType).Count(&count).Error
|
| return count, err
|
| }
|
|
|
|
|
| func CountChannelsGroupByType() (map[int64]int64, error) {
|
| type result struct {
|
| Type int64 `gorm:"column:type"`
|
| Count int64 `gorm:"column:count"`
|
| }
|
| var results []result
|
| err := DB.Model(&Channel{}).Select("type, count(*) as count").Group("type").Find(&results).Error
|
| if err != nil {
|
| return nil, err
|
| }
|
| counts := make(map[int64]int64)
|
| for _, r := range results {
|
| counts[r.Type] = r.Count
|
| }
|
| return counts, nil
|
| }
|
|
|