channel.go 23 KB

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