From a398f35968fe1d567b3a1c9b0fb387251ce11120 Mon Sep 17 00:00:00 2001 From: Bryan Date: Sun, 22 Oct 2023 18:38:29 +0800 Subject: [PATCH] fix: fix postgresql support (#606) * fix postgresql support fixes #517 * fix: fix pg support * chore: delete useless code --------- Co-authored-by: JustSong --- common/constants.go | 3 --- common/database.go | 6 ++++++ common/utils.go | 8 ++++++++ model/ability.go | 13 ++++++++++--- model/cache.go | 8 ++++++-- model/channel.go | 17 +++++------------ model/main.go | 1 + model/redemption.go | 7 ++++++- model/user.go | 7 ++++++- 9 files changed, 48 insertions(+), 22 deletions(-) create mode 100644 common/database.go diff --git a/common/constants.go b/common/constants.go index a0361c35..c25785c7 100644 --- a/common/constants.go +++ b/common/constants.go @@ -21,12 +21,9 @@ var QuotaPerUnit = 500 * 1000.0 // $0.002 / 1K tokens var DisplayInCurrencyEnabled = true var DisplayTokenStatEnabled = true -var UsingSQLite = false - // Any options with "Secret", "Token" in its key won't be return by GetOptions var SessionSecret = uuid.New().String() -var SQLitePath = "one-api.db" var OptionMap map[string]string var OptionMapRWMutex sync.RWMutex diff --git a/common/database.go b/common/database.go new file mode 100644 index 00000000..c7e9fd52 --- /dev/null +++ b/common/database.go @@ -0,0 +1,6 @@ +package common + +var UsingSQLite = false +var UsingPostgreSQL = false + +var SQLitePath = "one-api.db" diff --git a/common/utils.go b/common/utils.go index ab901b77..21bec8f5 100644 --- a/common/utils.go +++ b/common/utils.go @@ -199,3 +199,11 @@ func GetOrDefault(env string, defaultValue int) int { func MessageWithRequestId(message string, id string) string { return fmt.Sprintf("%s (request id: %s)", message, id) } + +func String2Int(str string) int { + num, err := strconv.Atoi(str) + if err != nil { + return 0 + } + return num +} diff --git a/model/ability.go b/model/ability.go index 50972a26..3da83be8 100644 --- a/model/ability.go +++ b/model/ability.go @@ -15,10 +15,17 @@ type Ability struct { func GetRandomSatisfiedChannel(group string, model string) (*Channel, error) { ability := Ability{} + groupCol := "`group`" + trueVal := "1" + if common.UsingPostgreSQL { + groupCol = `"group"` + trueVal = "true" + } + var err error = nil - maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where("`group` = ? and model = ? and enabled = 1", group, model) - channelQuery := DB.Where("`group` = ? and model = ? and enabled = 1 and priority = (?)", group, model, maxPrioritySubQuery) - if common.UsingSQLite { + maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where(groupCol+" = ? and model = ? and enabled = "+trueVal, group, model) + channelQuery := DB.Where(groupCol+" = ? and model = ? and enabled = "+trueVal+" and priority = (?)", group, model, maxPrioritySubQuery) + if common.UsingSQLite || common.UsingPostgreSQL { err = channelQuery.Order("RANDOM()").First(&ability).Error } else { err = channelQuery.Order("RAND()").First(&ability).Error diff --git a/model/cache.go b/model/cache.go index a7f5c06f..c6d0c70a 100644 --- a/model/cache.go +++ b/model/cache.go @@ -21,14 +21,18 @@ var ( ) func CacheGetTokenByKey(key string) (*Token, error) { + keyCol := "`key`" + if common.UsingPostgreSQL { + keyCol = `"key"` + } var token Token if !common.RedisEnabled { - err := DB.Where("`key` = ?", key).First(&token).Error + err := DB.Where(keyCol+" = ?", key).First(&token).Error return &token, err } tokenObjectString, err := common.RedisGet(fmt.Sprintf("token:%s", key)) if err != nil { - err := DB.Where("`key` = ?", key).First(&token).Error + err := DB.Where(keyCol+" = ?", key).First(&token).Error if err != nil { return nil, err } diff --git a/model/channel.go b/model/channel.go index 091a0d71..7e7b42e6 100644 --- a/model/channel.go +++ b/model/channel.go @@ -38,7 +38,11 @@ func GetAllChannels(startIdx int, num int, selectAll bool) ([]*Channel, error) { } func SearchChannels(keyword string) (channels []*Channel, err error) { - err = DB.Omit("key").Where("id = ? or name LIKE ? or `key` = ?", keyword, keyword+"%", keyword).Find(&channels).Error + keyCol := "`key`" + if common.UsingPostgreSQL { + keyCol = `"key"` + } + err = DB.Omit("key").Where("id = ? or name LIKE ? or "+keyCol+" = ?", common.String2Int(keyword), keyword+"%", keyword).Find(&channels).Error return channels, err } @@ -53,17 +57,6 @@ func GetChannelById(id int, selectAll bool) (*Channel, error) { return &channel, err } -func GetRandomChannel() (*Channel, error) { - channel := Channel{} - var err error = nil - if common.UsingSQLite { - err = DB.Where("status = ? and `group` = ?", common.ChannelStatusEnabled, "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 - } - return &channel, err -} - func BatchInsertChannels(channels []Channel) error { var err error err = DB.Create(&channels).Error diff --git a/model/main.go b/model/main.go index 0e962049..08182634 100644 --- a/model/main.go +++ b/model/main.go @@ -42,6 +42,7 @@ func chooseDB() (*gorm.DB, error) { if strings.HasPrefix(dsn, "postgres://") { // Use PostgreSQL common.SysLog("using PostgreSQL as database") + common.UsingPostgreSQL = true return gorm.Open(postgres.New(postgres.Config{ DSN: dsn, PreferSimpleProtocol: true, // disables implicit prepared statement usage diff --git a/model/redemption.go b/model/redemption.go index fafb2145..f16412b5 100644 --- a/model/redemption.go +++ b/model/redemption.go @@ -50,8 +50,13 @@ func Redeem(key string, userId int) (quota int, err error) { } redemption := &Redemption{} + keyCol := "`key`" + if common.UsingPostgreSQL { + keyCol = `"key"` + } + err = DB.Transaction(func(tx *gorm.DB) error { - err := tx.Set("gorm:query_option", "FOR UPDATE").Where("`key` = ?", key).First(redemption).Error + err := tx.Set("gorm:query_option", "FOR UPDATE").Where(keyCol+" = ?", key).First(redemption).Error if err != nil { return errors.New("无效的兑换码") } diff --git a/model/user.go b/model/user.go index 1b2ec7e6..7844eb6a 100644 --- a/model/user.go +++ b/model/user.go @@ -266,7 +266,12 @@ func GetUserEmail(id int) (email string, err error) { } func GetUserGroup(id int) (group string, err error) { - err = DB.Model(&User{}).Where("id = ?", id).Select("`group`").Find(&group).Error + groupCol := "`group`" + if common.UsingPostgreSQL { + groupCol = `"group"` + } + + err = DB.Model(&User{}).Where("id = ?", id).Select(groupCol).Find(&group).Error return group, err }