channel.go 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858
  1. package model
  2. import (
  3. "database/sql/driver"
  4. "encoding/json"
  5. "errors"
  6. "fmt"
  7. "math/rand"
  8. "one-api/common"
  9. "one-api/constant"
  10. "one-api/dto"
  11. "one-api/types"
  12. "strings"
  13. "sync"
  14. "gorm.io/gorm"
  15. )
  16. type Channel struct {
  17. Id int `json:"id"`
  18. Type int `json:"type" gorm:"default:0"`
  19. Key string `json:"key" gorm:"not null"`
  20. OpenAIOrganization *string `json:"openai_organization"`
  21. TestModel *string `json:"test_model"`
  22. Status int `json:"status" gorm:"default:1"`
  23. Name string `json:"name" gorm:"index"`
  24. Weight *uint `json:"weight" gorm:"default:0"`
  25. CreatedTime int64 `json:"created_time" gorm:"bigint"`
  26. TestTime int64 `json:"test_time" gorm:"bigint"`
  27. ResponseTime int `json:"response_time"` // in milliseconds
  28. BaseURL *string `json:"base_url" gorm:"column:base_url;default:''"`
  29. Other string `json:"other"`
  30. Balance float64 `json:"balance"` // in USD
  31. BalanceUpdatedTime int64 `json:"balance_updated_time" gorm:"bigint"`
  32. Models string `json:"models"`
  33. Group string `json:"group" gorm:"type:varchar(64);default:'default'"`
  34. UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"`
  35. ModelMapping *string `json:"model_mapping" gorm:"type:text"`
  36. //MaxInputTokens *int `json:"max_input_tokens" gorm:"default:0"`
  37. StatusCodeMapping *string `json:"status_code_mapping" gorm:"type:varchar(1024);default:''"`
  38. Priority *int64 `json:"priority" gorm:"bigint;default:0"`
  39. AutoBan *int `json:"auto_ban" gorm:"default:1"`
  40. OtherInfo string `json:"other_info"`
  41. Tag *string `json:"tag" gorm:"index"`
  42. Setting *string `json:"setting" gorm:"type:text"` // 渠道额外设置
  43. ParamOverride *string `json:"param_override" gorm:"type:text"`
  44. // add after v0.8.5
  45. ChannelInfo ChannelInfo `json:"channel_info" gorm:"type:json"`
  46. }
  47. type ChannelInfo struct {
  48. IsMultiKey bool `json:"is_multi_key"` // 是否多Key模式
  49. MultiKeySize int `json:"multi_key_size"` // 多Key模式下的Key数量
  50. MultiKeyStatusList map[int]int `json:"multi_key_status_list"` // key状态列表,key index -> status
  51. MultiKeyPollingIndex int `json:"multi_key_polling_index"` // 多Key模式下轮询的key索引
  52. MultiKeyMode constant.MultiKeyMode `json:"multi_key_mode"`
  53. }
  54. // Value implements driver.Valuer interface
  55. func (c ChannelInfo) Value() (driver.Value, error) {
  56. return common.Marshal(&c)
  57. }
  58. // Scan implements sql.Scanner interface
  59. func (c *ChannelInfo) Scan(value interface{}) error {
  60. bytesValue, _ := value.([]byte)
  61. return common.Unmarshal(bytesValue, c)
  62. }
  63. func (channel *Channel) getKeys() []string {
  64. if channel.Key == "" {
  65. return []string{}
  66. }
  67. // use \n to split keys
  68. keys := strings.Split(strings.Trim(channel.Key, "\n"), "\n")
  69. return keys
  70. }
  71. func (channel *Channel) GetNextEnabledKey() (string, *types.NewAPIError) {
  72. // If not in multi-key mode, return the original key string directly.
  73. if !channel.ChannelInfo.IsMultiKey {
  74. return channel.Key, nil
  75. }
  76. // Obtain all keys (split by \n)
  77. keys := channel.getKeys()
  78. if len(keys) == 0 {
  79. // No keys available, return error, should disable the channel
  80. return "", types.NewError(errors.New("no keys available"), types.ErrorCodeChannelNoAvailableKey)
  81. }
  82. statusList := channel.ChannelInfo.MultiKeyStatusList
  83. // helper to get key status, default to enabled when missing
  84. getStatus := func(idx int) int {
  85. if statusList == nil {
  86. return common.ChannelStatusEnabled
  87. }
  88. if status, ok := statusList[idx]; ok {
  89. return status
  90. }
  91. return common.ChannelStatusEnabled
  92. }
  93. // Collect indexes of enabled keys
  94. enabledIdx := make([]int, 0, len(keys))
  95. for i := range keys {
  96. if getStatus(i) == common.ChannelStatusEnabled {
  97. enabledIdx = append(enabledIdx, i)
  98. }
  99. }
  100. // If no specific status list or none enabled, fall back to first key
  101. if len(enabledIdx) == 0 {
  102. return keys[0], nil
  103. }
  104. switch channel.ChannelInfo.MultiKeyMode {
  105. case constant.MultiKeyModeRandom:
  106. // Randomly pick one enabled key
  107. return keys[enabledIdx[rand.Intn(len(enabledIdx))]], nil
  108. case constant.MultiKeyModePolling:
  109. // Use channel-specific lock to ensure thread-safe polling
  110. lock := getChannelPollingLock(channel.Id)
  111. lock.Lock()
  112. defer lock.Unlock()
  113. channelInfo, err := CacheGetChannelInfo(channel.Id)
  114. if err != nil {
  115. return "", types.NewError(err, types.ErrorCodeGetChannelFailed)
  116. }
  117. //println("before polling index:", channel.ChannelInfo.MultiKeyPollingIndex)
  118. defer func() {
  119. if common.DebugEnabled {
  120. println(fmt.Sprintf("channel %d polling index: %d", channel.Id, channel.ChannelInfo.MultiKeyPollingIndex))
  121. }
  122. if !common.MemoryCacheEnabled {
  123. _ = channel.SaveChannelInfo()
  124. } else {
  125. // CacheUpdateChannel(channel)
  126. }
  127. }()
  128. // Start from the saved polling index and look for the next enabled key
  129. start := channelInfo.MultiKeyPollingIndex
  130. if start < 0 || start >= len(keys) {
  131. start = 0
  132. }
  133. for i := 0; i < len(keys); i++ {
  134. idx := (start + i) % len(keys)
  135. if getStatus(idx) == common.ChannelStatusEnabled {
  136. // update polling index for next call (point to the next position)
  137. channel.ChannelInfo.MultiKeyPollingIndex = (idx + 1) % len(keys)
  138. return keys[idx], nil
  139. }
  140. }
  141. // Fallback – should not happen, but return first enabled key
  142. return keys[enabledIdx[0]], nil
  143. default:
  144. // Unknown mode, default to first enabled key (or original key string)
  145. return keys[enabledIdx[0]], nil
  146. }
  147. }
  148. func (channel *Channel) SaveChannelInfo() error {
  149. return DB.Model(channel).Update("channel_info", channel.ChannelInfo).Error
  150. }
  151. func (channel *Channel) GetModels() []string {
  152. if channel.Models == "" {
  153. return []string{}
  154. }
  155. return strings.Split(strings.Trim(channel.Models, ","), ",")
  156. }
  157. func (channel *Channel) GetGroups() []string {
  158. if channel.Group == "" {
  159. return []string{}
  160. }
  161. groups := strings.Split(strings.Trim(channel.Group, ","), ",")
  162. for i, group := range groups {
  163. groups[i] = strings.TrimSpace(group)
  164. }
  165. return groups
  166. }
  167. func (channel *Channel) GetOtherInfo() map[string]interface{} {
  168. otherInfo := make(map[string]interface{})
  169. if channel.OtherInfo != "" {
  170. err := json.Unmarshal([]byte(channel.OtherInfo), &otherInfo)
  171. if err != nil {
  172. common.SysError("failed to unmarshal other info: " + err.Error())
  173. }
  174. }
  175. return otherInfo
  176. }
  177. func (channel *Channel) SetOtherInfo(otherInfo map[string]interface{}) {
  178. otherInfoBytes, err := json.Marshal(otherInfo)
  179. if err != nil {
  180. common.SysError("failed to marshal other info: " + err.Error())
  181. return
  182. }
  183. channel.OtherInfo = string(otherInfoBytes)
  184. }
  185. func (channel *Channel) GetTag() string {
  186. if channel.Tag == nil {
  187. return ""
  188. }
  189. return *channel.Tag
  190. }
  191. func (channel *Channel) SetTag(tag string) {
  192. channel.Tag = &tag
  193. }
  194. func (channel *Channel) GetAutoBan() bool {
  195. if channel.AutoBan == nil {
  196. return false
  197. }
  198. return *channel.AutoBan == 1
  199. }
  200. func (channel *Channel) Save() error {
  201. return DB.Save(channel).Error
  202. }
  203. func GetAllChannels(startIdx int, num int, selectAll bool, idSort bool) ([]*Channel, error) {
  204. var channels []*Channel
  205. var err error
  206. order := "priority desc"
  207. if idSort {
  208. order = "id desc"
  209. }
  210. if selectAll {
  211. err = DB.Order(order).Find(&channels).Error
  212. } else {
  213. err = DB.Order(order).Limit(num).Offset(startIdx).Omit("key").Find(&channels).Error
  214. }
  215. return channels, err
  216. }
  217. func GetChannelsByTag(tag string, idSort bool) ([]*Channel, error) {
  218. var channels []*Channel
  219. order := "priority desc"
  220. if idSort {
  221. order = "id desc"
  222. }
  223. err := DB.Where("tag = ?", tag).Order(order).Find(&channels).Error
  224. return channels, err
  225. }
  226. func SearchChannels(keyword string, group string, model string, idSort bool) ([]*Channel, error) {
  227. var channels []*Channel
  228. modelsCol := "`models`"
  229. // 如果是 PostgreSQL,使用双引号
  230. if common.UsingPostgreSQL {
  231. modelsCol = `"models"`
  232. }
  233. baseURLCol := "`base_url`"
  234. // 如果是 PostgreSQL,使用双引号
  235. if common.UsingPostgreSQL {
  236. baseURLCol = `"base_url"`
  237. }
  238. order := "priority desc"
  239. if idSort {
  240. order = "id desc"
  241. }
  242. // 构造基础查询
  243. baseQuery := DB.Model(&Channel{}).Omit("key")
  244. // 构造WHERE子句
  245. var whereClause string
  246. var args []interface{}
  247. if group != "" && group != "null" {
  248. var groupCondition string
  249. if common.UsingMySQL {
  250. groupCondition = `CONCAT(',', ` + commonGroupCol + `, ',') LIKE ?`
  251. } else {
  252. // sqlite, PostgreSQL
  253. groupCondition = `(',' || ` + commonGroupCol + ` || ',') LIKE ?`
  254. }
  255. whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + ` LIKE ? AND ` + groupCondition
  256. args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%", "%,"+group+",%")
  257. } else {
  258. whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + " LIKE ?"
  259. args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%")
  260. }
  261. // 执行查询
  262. err := baseQuery.Where(whereClause, args...).Order(order).Find(&channels).Error
  263. if err != nil {
  264. return nil, err
  265. }
  266. return channels, nil
  267. }
  268. func GetChannelById(id int, selectAll bool) (*Channel, error) {
  269. channel := &Channel{Id: id}
  270. var err error = nil
  271. if selectAll {
  272. err = DB.First(channel, "id = ?", id).Error
  273. } else {
  274. err = DB.Omit("key").First(channel, "id = ?", id).Error
  275. }
  276. if err != nil {
  277. return nil, err
  278. }
  279. if channel == nil {
  280. return nil, errors.New("channel not found")
  281. }
  282. return channel, nil
  283. }
  284. func BatchInsertChannels(channels []Channel) error {
  285. var err error
  286. err = DB.Create(&channels).Error
  287. if err != nil {
  288. return err
  289. }
  290. for _, channel_ := range channels {
  291. err = channel_.AddAbilities()
  292. if err != nil {
  293. return err
  294. }
  295. }
  296. return nil
  297. }
  298. func BatchDeleteChannels(ids []int) error {
  299. //使用事务 删除channel表和channel_ability表
  300. tx := DB.Begin()
  301. err := tx.Where("id in (?)", ids).Delete(&Channel{}).Error
  302. if err != nil {
  303. // 回滚事务
  304. tx.Rollback()
  305. return err
  306. }
  307. err = tx.Where("channel_id in (?)", ids).Delete(&Ability{}).Error
  308. if err != nil {
  309. // 回滚事务
  310. tx.Rollback()
  311. return err
  312. }
  313. // 提交事务
  314. tx.Commit()
  315. return err
  316. }
  317. func (channel *Channel) GetPriority() int64 {
  318. if channel.Priority == nil {
  319. return 0
  320. }
  321. return *channel.Priority
  322. }
  323. func (channel *Channel) GetWeight() int {
  324. if channel.Weight == nil {
  325. return 0
  326. }
  327. return int(*channel.Weight)
  328. }
  329. func (channel *Channel) GetBaseURL() string {
  330. if channel.BaseURL == nil {
  331. return ""
  332. }
  333. return *channel.BaseURL
  334. }
  335. func (channel *Channel) GetModelMapping() string {
  336. if channel.ModelMapping == nil {
  337. return ""
  338. }
  339. return *channel.ModelMapping
  340. }
  341. func (channel *Channel) GetStatusCodeMapping() string {
  342. if channel.StatusCodeMapping == nil {
  343. return ""
  344. }
  345. return *channel.StatusCodeMapping
  346. }
  347. func (channel *Channel) Insert() error {
  348. var err error
  349. err = DB.Create(channel).Error
  350. if err != nil {
  351. return err
  352. }
  353. err = channel.AddAbilities()
  354. return err
  355. }
  356. func (channel *Channel) Update() error {
  357. var err error
  358. err = DB.Model(channel).Updates(channel).Error
  359. if err != nil {
  360. return err
  361. }
  362. DB.Model(channel).First(channel, "id = ?", channel.Id)
  363. err = channel.UpdateAbilities(nil)
  364. return err
  365. }
  366. func (channel *Channel) UpdateResponseTime(responseTime int64) {
  367. err := DB.Model(channel).Select("response_time", "test_time").Updates(Channel{
  368. TestTime: common.GetTimestamp(),
  369. ResponseTime: int(responseTime),
  370. }).Error
  371. if err != nil {
  372. common.SysError("failed to update response time: " + err.Error())
  373. }
  374. }
  375. func (channel *Channel) UpdateBalance(balance float64) {
  376. err := DB.Model(channel).Select("balance_updated_time", "balance").Updates(Channel{
  377. BalanceUpdatedTime: common.GetTimestamp(),
  378. Balance: balance,
  379. }).Error
  380. if err != nil {
  381. common.SysError("failed to update balance: " + err.Error())
  382. }
  383. }
  384. func (channel *Channel) Delete() error {
  385. var err error
  386. err = DB.Delete(channel).Error
  387. if err != nil {
  388. return err
  389. }
  390. err = channel.DeleteAbilities()
  391. return err
  392. }
  393. var channelStatusLock sync.Mutex
  394. // channelPollingLocks stores locks for each channel.id to ensure thread-safe polling
  395. var channelPollingLocks sync.Map
  396. // getChannelPollingLock returns or creates a mutex for the given channel ID
  397. func getChannelPollingLock(channelId int) *sync.Mutex {
  398. if lock, exists := channelPollingLocks.Load(channelId); exists {
  399. return lock.(*sync.Mutex)
  400. }
  401. // Create new lock for this channel
  402. newLock := &sync.Mutex{}
  403. actual, _ := channelPollingLocks.LoadOrStore(channelId, newLock)
  404. return actual.(*sync.Mutex)
  405. }
  406. // CleanupChannelPollingLocks removes locks for channels that no longer exist
  407. // This is optional and can be called periodically to prevent memory leaks
  408. func CleanupChannelPollingLocks() {
  409. var activeChannelIds []int
  410. DB.Model(&Channel{}).Pluck("id", &activeChannelIds)
  411. activeChannelSet := make(map[int]bool)
  412. for _, id := range activeChannelIds {
  413. activeChannelSet[id] = true
  414. }
  415. channelPollingLocks.Range(func(key, value interface{}) bool {
  416. channelId := key.(int)
  417. if !activeChannelSet[channelId] {
  418. channelPollingLocks.Delete(channelId)
  419. }
  420. return true
  421. })
  422. }
  423. func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int) {
  424. keys := channel.getKeys()
  425. if len(keys) == 0 {
  426. channel.Status = status
  427. } else {
  428. var keyIndex int
  429. for i, key := range keys {
  430. if key == usingKey {
  431. keyIndex = i
  432. break
  433. }
  434. }
  435. if channel.ChannelInfo.MultiKeyStatusList == nil {
  436. channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
  437. }
  438. if status == common.ChannelStatusEnabled {
  439. delete(channel.ChannelInfo.MultiKeyStatusList, keyIndex)
  440. } else {
  441. channel.ChannelInfo.MultiKeyStatusList[keyIndex] = status
  442. }
  443. if len(channel.ChannelInfo.MultiKeyStatusList) >= channel.ChannelInfo.MultiKeySize {
  444. channel.Status = common.ChannelStatusAutoDisabled
  445. info := channel.GetOtherInfo()
  446. info["status_reason"] = "All keys are disabled"
  447. info["status_time"] = common.GetTimestamp()
  448. channel.SetOtherInfo(info)
  449. }
  450. }
  451. }
  452. func UpdateChannelStatus(channelId int, usingKey string, status int, reason string) bool {
  453. if common.MemoryCacheEnabled {
  454. channelStatusLock.Lock()
  455. defer channelStatusLock.Unlock()
  456. channelCache, _ := CacheGetChannel(channelId)
  457. if channelCache == nil {
  458. return false
  459. }
  460. if channelCache.ChannelInfo.IsMultiKey {
  461. // 如果是多Key模式,更新缓存中的状态
  462. handlerMultiKeyUpdate(channelCache, usingKey, status)
  463. //CacheUpdateChannel(channelCache)
  464. //return true
  465. } else {
  466. // 如果缓存渠道存在,且状态已是目标状态,直接返回
  467. if channelCache.Status == status {
  468. return false
  469. }
  470. // 如果缓存渠道不存在(说明已经被禁用),且要设置的状态不为启用,直接返回
  471. if status != common.ChannelStatusEnabled {
  472. return false
  473. }
  474. CacheUpdateChannelStatus(channelId, status)
  475. }
  476. }
  477. shouldUpdateAbilities := false
  478. defer func() {
  479. if shouldUpdateAbilities {
  480. err := UpdateAbilityStatus(channelId, status == common.ChannelStatusEnabled)
  481. if err != nil {
  482. common.SysError("failed to update ability status: " + err.Error())
  483. }
  484. }
  485. }()
  486. channel, err := GetChannelById(channelId, true)
  487. if err != nil {
  488. return false
  489. } else {
  490. if channel.Status == status {
  491. return false
  492. }
  493. if channel.ChannelInfo.IsMultiKey {
  494. beforeStatus := channel.Status
  495. handlerMultiKeyUpdate(channel, usingKey, status)
  496. if beforeStatus != channel.Status {
  497. shouldUpdateAbilities = true
  498. }
  499. } else {
  500. info := channel.GetOtherInfo()
  501. info["status_reason"] = reason
  502. info["status_time"] = common.GetTimestamp()
  503. channel.SetOtherInfo(info)
  504. channel.Status = status
  505. shouldUpdateAbilities = true
  506. }
  507. err = channel.Save()
  508. if err != nil {
  509. common.SysError("failed to update channel status: " + err.Error())
  510. return false
  511. }
  512. }
  513. return true
  514. }
  515. func EnableChannelByTag(tag string) error {
  516. err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusEnabled).Error
  517. if err != nil {
  518. return err
  519. }
  520. err = UpdateAbilityStatusByTag(tag, true)
  521. return err
  522. }
  523. func DisableChannelByTag(tag string) error {
  524. err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusManuallyDisabled).Error
  525. if err != nil {
  526. return err
  527. }
  528. err = UpdateAbilityStatusByTag(tag, false)
  529. return err
  530. }
  531. func EditChannelByTag(tag string, newTag *string, modelMapping *string, models *string, group *string, priority *int64, weight *uint) error {
  532. updateData := Channel{}
  533. shouldReCreateAbilities := false
  534. updatedTag := tag
  535. // 如果 newTag 不为空且不等于 tag,则更新 tag
  536. if newTag != nil && *newTag != tag {
  537. updateData.Tag = newTag
  538. updatedTag = *newTag
  539. }
  540. if modelMapping != nil && *modelMapping != "" {
  541. updateData.ModelMapping = modelMapping
  542. }
  543. if models != nil && *models != "" {
  544. shouldReCreateAbilities = true
  545. updateData.Models = *models
  546. }
  547. if group != nil && *group != "" {
  548. shouldReCreateAbilities = true
  549. updateData.Group = *group
  550. }
  551. if priority != nil {
  552. updateData.Priority = priority
  553. }
  554. if weight != nil {
  555. updateData.Weight = weight
  556. }
  557. err := DB.Model(&Channel{}).Where("tag = ?", tag).Updates(updateData).Error
  558. if err != nil {
  559. return err
  560. }
  561. if shouldReCreateAbilities {
  562. channels, err := GetChannelsByTag(updatedTag, false)
  563. if err == nil {
  564. for _, channel := range channels {
  565. err = channel.UpdateAbilities(nil)
  566. if err != nil {
  567. common.SysError("failed to update abilities: " + err.Error())
  568. }
  569. }
  570. }
  571. } else {
  572. err := UpdateAbilityByTag(tag, newTag, priority, weight)
  573. if err != nil {
  574. return err
  575. }
  576. }
  577. return nil
  578. }
  579. func UpdateChannelUsedQuota(id int, quota int) {
  580. if common.BatchUpdateEnabled {
  581. addNewRecord(BatchUpdateTypeChannelUsedQuota, id, quota)
  582. return
  583. }
  584. updateChannelUsedQuota(id, quota)
  585. }
  586. func updateChannelUsedQuota(id int, quota int) {
  587. err := DB.Model(&Channel{}).Where("id = ?", id).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error
  588. if err != nil {
  589. common.SysError("failed to update channel used quota: " + err.Error())
  590. }
  591. }
  592. func DeleteChannelByStatus(status int64) (int64, error) {
  593. result := DB.Where("status = ?", status).Delete(&Channel{})
  594. return result.RowsAffected, result.Error
  595. }
  596. func DeleteDisabledChannel() (int64, error) {
  597. result := DB.Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).Delete(&Channel{})
  598. return result.RowsAffected, result.Error
  599. }
  600. func GetPaginatedTags(offset int, limit int) ([]*string, error) {
  601. var tags []*string
  602. err := DB.Model(&Channel{}).Select("DISTINCT tag").Where("tag != ''").Offset(offset).Limit(limit).Find(&tags).Error
  603. return tags, err
  604. }
  605. func SearchTags(keyword string, group string, model string, idSort bool) ([]*string, error) {
  606. var tags []*string
  607. modelsCol := "`models`"
  608. // 如果是 PostgreSQL,使用双引号
  609. if common.UsingPostgreSQL {
  610. modelsCol = `"models"`
  611. }
  612. baseURLCol := "`base_url`"
  613. // 如果是 PostgreSQL,使用双引号
  614. if common.UsingPostgreSQL {
  615. baseURLCol = `"base_url"`
  616. }
  617. order := "priority desc"
  618. if idSort {
  619. order = "id desc"
  620. }
  621. // 构造基础查询
  622. baseQuery := DB.Model(&Channel{}).Omit("key")
  623. // 构造WHERE子句
  624. var whereClause string
  625. var args []interface{}
  626. if group != "" && group != "null" {
  627. var groupCondition string
  628. if common.UsingMySQL {
  629. groupCondition = `CONCAT(',', ` + commonGroupCol + `, ',') LIKE ?`
  630. } else {
  631. // sqlite, PostgreSQL
  632. groupCondition = `(',' || ` + commonGroupCol + ` || ',') LIKE ?`
  633. }
  634. whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + ` LIKE ? AND ` + groupCondition
  635. args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%", "%,"+group+",%")
  636. } else {
  637. whereClause = "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + " LIKE ?"
  638. args = append(args, common.String2Int(keyword), "%"+keyword+"%", keyword, "%"+keyword+"%", "%"+model+"%")
  639. }
  640. subQuery := baseQuery.Where(whereClause, args...).
  641. Select("tag").
  642. Where("tag != ''").
  643. Order(order)
  644. err := DB.Table("(?) as sub", subQuery).
  645. Select("DISTINCT tag").
  646. Find(&tags).Error
  647. if err != nil {
  648. return nil, err
  649. }
  650. return tags, nil
  651. }
  652. func (channel *Channel) ValidateSettings() error {
  653. channelParams := &dto.ChannelSettings{}
  654. if channel.Setting != nil && *channel.Setting != "" {
  655. err := json.Unmarshal([]byte(*channel.Setting), channelParams)
  656. if err != nil {
  657. return err
  658. }
  659. }
  660. return nil
  661. }
  662. func (channel *Channel) GetSetting() dto.ChannelSettings {
  663. setting := dto.ChannelSettings{}
  664. if channel.Setting != nil && *channel.Setting != "" {
  665. err := json.Unmarshal([]byte(*channel.Setting), &setting)
  666. if err != nil {
  667. common.SysError("failed to unmarshal setting: " + err.Error())
  668. channel.Setting = nil // 清空设置以避免后续错误
  669. _ = channel.Save() // 保存修改
  670. }
  671. }
  672. return setting
  673. }
  674. func (channel *Channel) SetSetting(setting dto.ChannelSettings) {
  675. settingBytes, err := json.Marshal(setting)
  676. if err != nil {
  677. common.SysError("failed to marshal setting: " + err.Error())
  678. return
  679. }
  680. channel.Setting = common.GetPointer[string](string(settingBytes))
  681. }
  682. func (channel *Channel) GetParamOverride() map[string]interface{} {
  683. paramOverride := make(map[string]interface{})
  684. if channel.ParamOverride != nil && *channel.ParamOverride != "" {
  685. err := json.Unmarshal([]byte(*channel.ParamOverride), &paramOverride)
  686. if err != nil {
  687. common.SysError("failed to unmarshal param override: " + err.Error())
  688. }
  689. }
  690. return paramOverride
  691. }
  692. func GetChannelsByIds(ids []int) ([]*Channel, error) {
  693. var channels []*Channel
  694. err := DB.Where("id in (?)", ids).Find(&channels).Error
  695. return channels, err
  696. }
  697. func BatchSetChannelTag(ids []int, tag *string) error {
  698. // 开启事务
  699. tx := DB.Begin()
  700. if tx.Error != nil {
  701. return tx.Error
  702. }
  703. // 更新标签
  704. err := tx.Model(&Channel{}).Where("id in (?)", ids).Update("tag", tag).Error
  705. if err != nil {
  706. tx.Rollback()
  707. return err
  708. }
  709. // update ability status
  710. channels, err := GetChannelsByIds(ids)
  711. if err != nil {
  712. tx.Rollback()
  713. return err
  714. }
  715. for _, channel := range channels {
  716. err = channel.UpdateAbilities(tx)
  717. if err != nil {
  718. tx.Rollback()
  719. return err
  720. }
  721. }
  722. // 提交事务
  723. return tx.Commit().Error
  724. }
  725. // CountAllChannels returns total channels in DB
  726. func CountAllChannels() (int64, error) {
  727. var total int64
  728. err := DB.Model(&Channel{}).Count(&total).Error
  729. return total, err
  730. }
  731. // CountAllTags returns number of non-empty distinct tags
  732. func CountAllTags() (int64, error) {
  733. var total int64
  734. err := DB.Model(&Channel{}).Where("tag is not null AND tag != ''").Distinct("tag").Count(&total).Error
  735. return total, err
  736. }
  737. // Get channels of specified type with pagination
  738. func GetChannelsByType(startIdx int, num int, idSort bool, channelType int) ([]*Channel, error) {
  739. var channels []*Channel
  740. order := "priority desc"
  741. if idSort {
  742. order = "id desc"
  743. }
  744. err := DB.Where("type = ?", channelType).Order(order).Limit(num).Offset(startIdx).Omit("key").Find(&channels).Error
  745. return channels, err
  746. }
  747. // Count channels of specific type
  748. func CountChannelsByType(channelType int) (int64, error) {
  749. var count int64
  750. err := DB.Model(&Channel{}).Where("type = ?", channelType).Count(&count).Error
  751. return count, err
  752. }
  753. // Return map[type]count for all channels
  754. func CountChannelsGroupByType() (map[int64]int64, error) {
  755. type result struct {
  756. Type int64 `gorm:"column:type"`
  757. Count int64 `gorm:"column:count"`
  758. }
  759. var results []result
  760. err := DB.Model(&Channel{}).Select("type, count(*) as count").Group("type").Find(&results).Error
  761. if err != nil {
  762. return nil, err
  763. }
  764. counts := make(map[int64]int64)
  765. for _, r := range results {
  766. counts[r.Type] = r.Count
  767. }
  768. return counts, nil
  769. }