|
|
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 { |
|
|
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 |
|
|
} |
|
|
|