Skip to content

Commit 5e81e19

Browse files
committed
fix: fix SQL channel selection algo (songquanpeng#1197)
1 parent 96d7a99 commit 5e81e19

File tree

2 files changed

+10
-4
lines changed

2 files changed

+10
-4
lines changed

model/ability.go

+9-3
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package model
22

33
import (
44
"github.com/songquanpeng/one-api/common"
5+
"gorm.io/gorm"
56
"strings"
67
)
78

@@ -13,7 +14,7 @@ type Ability struct {
1314
Priority *int64 `json:"priority" gorm:"bigint;default:0;index"`
1415
}
1516

16-
func GetRandomSatisfiedChannel(group string, model string) (*Channel, error) {
17+
func GetRandomSatisfiedChannel(group string, model string, ignoreFirstPriority bool) (*Channel, error) {
1718
ability := Ability{}
1819
groupCol := "`group`"
1920
trueVal := "1"
@@ -23,8 +24,13 @@ func GetRandomSatisfiedChannel(group string, model string) (*Channel, error) {
2324
}
2425

2526
var err error = nil
26-
maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where(groupCol+" = ? and model = ? and enabled = "+trueVal, group, model)
27-
channelQuery := DB.Where(groupCol+" = ? and model = ? and enabled = "+trueVal+" and priority = (?)", group, model, maxPrioritySubQuery)
27+
var channelQuery *gorm.DB
28+
if ignoreFirstPriority {
29+
channelQuery = DB.Where(groupCol+" = ? and model = ? and enabled = "+trueVal, group, model)
30+
} else {
31+
maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where(groupCol+" = ? and model = ? and enabled = "+trueVal, group, model)
32+
channelQuery = DB.Where(groupCol+" = ? and model = ? and enabled = "+trueVal+" and priority = (?)", group, model, maxPrioritySubQuery)
33+
}
2834
if common.UsingSQLite || common.UsingPostgreSQL {
2935
err = channelQuery.Order("RANDOM()").First(&ability).Error
3036
} else {

model/cache.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -205,7 +205,7 @@ func SyncChannelCache(frequency int) {
205205

206206
func CacheGetRandomSatisfiedChannel(group string, model string, ignoreFirstPriority bool) (*Channel, error) {
207207
if !config.MemoryCacheEnabled {
208-
return GetRandomSatisfiedChannel(group, model)
208+
return GetRandomSatisfiedChannel(group, model, ignoreFirstPriority)
209209
}
210210
channelSyncLock.RLock()
211211
defer channelSyncLock.RUnlock()

0 commit comments

Comments
 (0)