Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: pgsql new

Browse Source
This commit is contained in:
ckt1031 2023-10-06 22:34:28 +08:00
parent c034cebc79
commit c78fba83eb
2 changed files with 28 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
model/ability.go
View File

@ -1,7 +1,6 @@
package model
import (
"fmt"
"one-api/common"
"strings"
)
@ -20,21 +19,33 @@ func GetRandomSatisfiedChannel(group string, model string, stream bool) (*Channe
ability := Ability{}
var err error = nil
cmd := "`group` = ? and model = ? and enabled = 1"
if common.UsingPostgreSQL {
// Make cmd compatible with PostgreSQL
cmd = "\"group\" = ? and model = ? and enabled = true"
}
var cmdWhere *Ability
if stream {
cmd += fmt.Sprintf(" and allow_streaming = %d", common.ChannelAllowStreamEnabled)
cmdWhere = &Ability{
Group: group,
Model: model,
Enabled: true,
AllowStreaming: common.ChannelAllowStreamEnabled,
}
} else {
cmd += fmt.Sprintf(" and allow_non_streaming = %d", common.ChannelAllowNonStreamEnabled)
cmdWhere = &Ability{
Group: group,
Model: model,
Enabled: true,
AllowNonStreaming: common.ChannelAllowNonStreamEnabled,
}
}
maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where(cmd, group, model)
channelQuery := DB.Where("`group` = ? and model = ? and enabled = 1 and priority = (?)", group, model, maxPrioritySubQuery)
maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where(cmdWhere)
cmd1 := "`group` = ? and model = ? and enabled = 1 and priority = (?)"
if common.UsingPostgreSQL {
cmd1 = "\"group\" = ? and model = ? and enabled = 1 and priority = (?)"
}
channelQuery := DB.Where(cmd1, group, model, maxPrioritySubQuery)
if common.UsingSQLite || common.UsingPostgreSQL {
err = channelQuery.Order("RANDOM()").First(&ability).Error
} else {

15
model/channel.go
View File

@ -2,6 +2,7 @@ package model
import (
"one-api/common"
"strconv"
"gorm.io/gorm"
)
@ -41,13 +42,9 @@ func GetAllChannels(startIdx int, num int, selectAll bool) ([]*Channel, error) {
}
func SearchChannels(keyword string) (channels []*Channel, err error) {
whereItem := "id = ? or name LIKE ? or `key` = ?"
idKeyword, err := strconv.Atoi(keyword)
if common.UsingPostgreSQL {
whereItem = "id = ? or name LIKE ? or \"key\" = ?"
}
err = DB.Omit("key").Where(whereItem, keyword, keyword+"%", keyword).Find(&channels).Error
err = DB.Omit("key").Where("name LIKE ?", keyword+"%").Or(&Channel{Id: idKeyword}).Or(&Channel{Key: keyword}).Find(&channels).Error
return channels, err
}
@ -66,11 +63,11 @@ func GetRandomChannel() (*Channel, error) {
channel := Channel{}
var err error = nil
if common.UsingPostgreSQL {
err = DB.Where("status = ? and \"group\" = ?", common.ChannelStatusEnabled, "default").Order("RANDOM()").Limit(1).First(&channel).Error
err = DB.Where(&Channel{Status: common.ChannelStatusEnabled, Group: "default"}).Order("RANDOM()").Limit(1).First(&channel).Error
} else if common.UsingSQLite {
err = DB.Where("status = ? and `group` = ?", common.ChannelStatusEnabled, "default").Order("RANDOM()").Limit(1).First(&channel).Error
err = DB.Where(&Channel{Status: common.ChannelStatusEnabled, Group: "default"}).Order("RANDOM()").Limit(1).First(&channel).Error
} else {
err = DB.Where("status = ? and `group` = ?", common.ChannelStatusEnabled, "default").Order("RAND()").Limit(1).First(&channel).Error
err = DB.Where(&Channel{Status: common.ChannelStatusEnabled, Group: "default"}).Order("RAND()").Limit(1).First(&channel).Error
}
return &channel, err
}