group_repo.go 38 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200
  1. package repository
  2. import (
  3. "context"
  4. "database/sql"
  5. "encoding/json"
  6. "errors"
  7. "fmt"
  8. "sort"
  9. "strings"
  10. dbent "github.com/Wei-Shaw/sub2api/ent"
  11. "github.com/Wei-Shaw/sub2api/ent/group"
  12. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  13. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  14. "github.com/Wei-Shaw/sub2api/internal/service"
  15. "github.com/lib/pq"
  16. entsql "entgo.io/ent/dialect/sql"
  17. )
  18. type sqlExecutor interface {
  19. ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
  20. QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
  21. }
  22. type groupRepository struct {
  23. client *dbent.Client
  24. sql sqlExecutor
  25. }
  26. // lockLiveGroups makes account-group inserts participate in the same row-lock
  27. // protocol as guarded group deletion. FOR SHARE conflicts with the deleter's
  28. // FOR UPDATE lock, and READ COMMITTED rechecks deleted_at after any wait.
  29. func lockLiveGroups(ctx context.Context, exec sqlExecutor, groupIDs []int64) error {
  30. if len(groupIDs) == 0 {
  31. return nil
  32. }
  33. unique := make(map[int64]struct{}, len(groupIDs))
  34. for _, id := range groupIDs {
  35. unique[id] = struct{}{}
  36. }
  37. rows, err := exec.QueryContext(ctx, `/* account_group_live_group_lock */
  38. SELECT id FROM groups
  39. WHERE id = ANY($1) AND deleted_at IS NULL
  40. ORDER BY id
  41. FOR SHARE`, pq.Array(groupIDs))
  42. if err != nil {
  43. return err
  44. }
  45. defer func() { _ = rows.Close() }()
  46. locked := 0
  47. for rows.Next() {
  48. locked++
  49. }
  50. if err := rows.Err(); err != nil {
  51. return err
  52. }
  53. if locked != len(unique) {
  54. return service.ErrGroupNotFound
  55. }
  56. return nil
  57. }
  58. func NewGroupRepository(client *dbent.Client, sqlDB *sql.DB) service.GroupRepository {
  59. return newGroupRepositoryWithSQL(client, sqlDB)
  60. }
  61. // NewAdminGroupRepository exposes the atomic group-duplication capability as
  62. // an explicit dependency of the admin service.
  63. func NewAdminGroupRepository(client *dbent.Client, sqlDB *sql.DB) service.AdminGroupRepository {
  64. return newGroupRepositoryWithSQL(client, sqlDB)
  65. }
  66. func newGroupRepositoryWithSQL(client *dbent.Client, sqlq sqlExecutor) *groupRepository {
  67. return &groupRepository{client: client, sql: sqlq}
  68. }
  69. func (r *groupRepository) Create(ctx context.Context, groupIn *service.Group) error {
  70. if err := createGroupRecord(ctx, r.client, groupIn); err != nil {
  71. return err
  72. }
  73. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupIn.ID, nil); err != nil {
  74. logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue group create failed: group=%d err=%v", groupIn.ID, err)
  75. }
  76. return nil
  77. }
  78. func createGroupRecord(ctx context.Context, client *dbent.Client, groupIn *service.Group) error {
  79. if groupIn == nil {
  80. return errors.New("group is nil")
  81. }
  82. modelPricing, err := json.Marshal(groupIn.ModelPricing)
  83. if err != nil {
  84. return fmt.Errorf("marshal group model pricing: %w", err)
  85. }
  86. builder := client.Group.Create().
  87. SetName(groupIn.Name).
  88. SetDescription(groupIn.Description).
  89. SetPlatform(groupIn.Platform).
  90. SetRateMultiplier(groupIn.RateMultiplier).
  91. SetSortOrder(groupIn.SortOrder).
  92. SetIsExclusive(groupIn.IsExclusive).
  93. SetSecurityPolicyEnabled(groupIn.SecurityPolicyEnabled).
  94. SetSecurityPolicyMode(service.NormalizeSecurityPolicyMode(groupIn.SecurityPolicyMode)).
  95. SetSecurityPolicyEmailEnabled(groupIn.SecurityPolicyEmailEnabled).
  96. SetStatus(groupIn.Status).
  97. SetSubscriptionType(groupIn.SubscriptionType).
  98. SetNillableDailyLimitUsd(groupIn.DailyLimitUSD).
  99. SetNillableWeeklyLimitUsd(groupIn.WeeklyLimitUSD).
  100. SetNillableMonthlyLimitUsd(groupIn.MonthlyLimitUSD).
  101. SetAllowImageGeneration(groupIn.AllowImageGeneration).
  102. SetAllowBatchImageGeneration(groupIn.AllowBatchImageGeneration).
  103. SetImageRateIndependent(groupIn.ImageRateIndependent).
  104. SetImageRateMultiplier(groupIn.ImageRateMultiplier).
  105. SetNillableImagePrice1k(groupIn.ImagePrice1K).
  106. SetNillableImagePrice2k(groupIn.ImagePrice2K).
  107. SetNillableImagePrice4k(groupIn.ImagePrice4K).
  108. SetBatchImageDiscountMultiplier(groupIn.BatchImageDiscountMultiplier).
  109. SetBatchImageHoldMultiplier(groupIn.BatchImageHoldMultiplier).
  110. SetVideoRateIndependent(groupIn.VideoRateIndependent).
  111. SetVideoRateMultiplier(groupIn.VideoRateMultiplier).
  112. SetNillableVideoPrice480p(groupIn.VideoPrice480P).
  113. SetNillableVideoPrice720p(groupIn.VideoPrice720P).
  114. SetNillableVideoPrice1080p(groupIn.VideoPrice1080P).
  115. SetVideoModelPrices(service.NormalizeVideoModelPrices(groupIn.VideoModelPrices)).
  116. SetNillableWebSearchPricePerCall(groupIn.WebSearchPricePerCall).
  117. SetNillableSearchPricePer1k(groupIn.SearchPricePer1k).
  118. SetNillableAudioRealtimePricePerMin(groupIn.AudioRealtimePricePerMin).
  119. SetNillableAudioTtsPricePerMillionChars(groupIn.AudioTTSPricePerMillionChars).
  120. SetNillableAudioSttPricePerHour(groupIn.AudioSTTPricePerHour).
  121. SetLongContextPricingEnabled(groupIn.LongContextPricingEnabled).
  122. SetModelPricing(modelPricing).
  123. SetDefaultValidityDays(groupIn.DefaultValidityDays).
  124. SetClaudeCodeOnly(groupIn.ClaudeCodeOnly).
  125. SetNillableFallbackGroupID(groupIn.FallbackGroupID).
  126. SetNillableFallbackGroupIDOnInvalidRequest(groupIn.FallbackGroupIDOnInvalidRequest).
  127. SetModelRoutingEnabled(groupIn.ModelRoutingEnabled).
  128. SetMcpXMLInject(groupIn.MCPXMLInject).
  129. SetAllowMessagesDispatch(groupIn.AllowMessagesDispatch).
  130. SetAllowLive(groupIn.AllowLive).
  131. SetForceOpenaiFast(groupIn.ForceOpenAIFast).
  132. SetFreeOpenaiFast(groupIn.FreeOpenAIFast).
  133. SetRequireOauthOnly(groupIn.RequireOAuthOnly).
  134. SetRequirePrivacySet(groupIn.RequirePrivacySet).
  135. SetDefaultMappedModel(groupIn.DefaultMappedModel).
  136. SetMessagesDispatchModelConfig(groupIn.MessagesDispatchModelConfig).
  137. SetModelAllowlist(service.DomainGroupModelAllowlist(groupIn.ModelAllowlist)).
  138. SetCodexModelsManifestConfig(groupIn.CodexModelsManifestConfig).
  139. SetRpmLimit(groupIn.RPMLimit).
  140. SetMaxReasoningEffort(groupIn.MaxReasoningEffort).
  141. SetMaxReasoningEffortOverLimit(groupIn.MaxReasoningEffortOverLimit).
  142. SetReasoningEffortMappings(groupIn.ReasoningEffortMappings).
  143. SetPeakRateEnabled(groupIn.PeakRateEnabled).
  144. SetPeakStart(groupIn.PeakStart).
  145. SetPeakEnd(groupIn.PeakEnd).
  146. SetPeakRateMultiplier(groupIn.PeakRateMultiplier).
  147. SetProfitControlEnabled(groupIn.ProfitControlEnabled).
  148. SetProfitMinMargin(groupIn.ProfitMinMargin).
  149. SetProfitSafetyBuffer(groupIn.ProfitSafetyBuffer)
  150. if groupIn.DuplicateOperationID != "" {
  151. builder = builder.SetDuplicateOperationID(groupIn.DuplicateOperationID)
  152. }
  153. // 设置模型路由配置
  154. if groupIn.ModelRouting != nil {
  155. builder = builder.SetModelRouting(groupIn.ModelRouting)
  156. }
  157. // 设置支持的模型系列(始终设置,空数组表示不限制)
  158. builder = builder.SetSupportedModelScopes(groupIn.SupportedModelScopes)
  159. created, err := builder.Save(ctx)
  160. if err != nil {
  161. return translatePersistenceError(err, nil, service.ErrGroupExists)
  162. }
  163. groupIn.ID = created.ID
  164. groupIn.CreatedAt = created.CreatedAt
  165. groupIn.UpdatedAt = created.UpdatedAt
  166. return nil
  167. }
  168. func (r *groupRepository) FindByDuplicateOperationID(ctx context.Context, operationID string) (*service.Group, error) {
  169. operationID = strings.TrimSpace(operationID)
  170. if operationID == "" {
  171. return nil, nil
  172. }
  173. row, err := r.client.Group.Query().
  174. Where(group.DuplicateOperationIDEQ(operationID)).
  175. Order(dbent.Asc(group.FieldID)).
  176. First(ctx)
  177. if dbent.IsNotFound(err) {
  178. return nil, nil
  179. }
  180. if err != nil {
  181. return nil, fmt.Errorf("find group duplicate operation: %w", err)
  182. }
  183. return groupEntityToService(row), nil
  184. }
  185. // CreateFromSource atomically persists a copied group, clones the source
  186. // account bindings with their exact priorities, and writes its scheduler event.
  187. func (r *groupRepository) CreateFromSource(ctx context.Context, groupIn *service.Group, sourceGroupID int64) error {
  188. if groupIn == nil {
  189. return errors.New("group is nil")
  190. }
  191. tx, err := r.client.Tx(ctx)
  192. if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
  193. return err
  194. }
  195. var txClient *dbent.Client
  196. if err == nil {
  197. defer func() { _ = tx.Rollback() }()
  198. txClient = tx.Client()
  199. } else {
  200. // Reuse a caller-owned transaction when this repository is already transactional.
  201. txClient = r.client
  202. }
  203. if err := createGroupRecord(ctx, txClient, groupIn); err != nil {
  204. return err
  205. }
  206. result, err := txClient.ExecContext(
  207. ctx,
  208. `INSERT INTO account_groups (account_id, group_id, priority, created_at)
  209. SELECT ag.account_id, $2, ag.priority, NOW()
  210. FROM account_groups ag
  211. JOIN accounts a ON a.id = ag.account_id
  212. WHERE ag.group_id = $1
  213. AND a.deleted_at IS NULL
  214. AND (NOT $3 OR a.type <> $4)
  215. ON CONFLICT (account_id, group_id) DO NOTHING`,
  216. sourceGroupID,
  217. groupIn.ID,
  218. groupIn.RequireOAuthOnly,
  219. service.AccountTypeAPIKey,
  220. )
  221. if err != nil {
  222. return err
  223. }
  224. if count, countErr := result.RowsAffected(); countErr == nil {
  225. groupIn.AccountCount = count
  226. }
  227. if err := enqueueSchedulerOutbox(ctx, txClient, service.SchedulerOutboxEventGroupChanged, nil, &groupIn.ID, nil); err != nil {
  228. return err
  229. }
  230. if tx != nil {
  231. if err := tx.Commit(); err != nil {
  232. return err
  233. }
  234. }
  235. return nil
  236. }
  237. func (r *groupRepository) GetByID(ctx context.Context, id int64) (*service.Group, error) {
  238. out, err := r.GetByIDLite(ctx, id)
  239. if err != nil {
  240. return nil, err
  241. }
  242. counts, err := r.loadAccountCounts(ctx, []int64{out.ID})
  243. if err == nil {
  244. c := counts[out.ID]
  245. out.AccountCount = c.Total
  246. out.ActiveAccountCount = c.Active
  247. out.RateLimitedAccountCount = c.RateLimited
  248. }
  249. return out, nil
  250. }
  251. func (r *groupRepository) GetByIDLite(ctx context.Context, id int64) (*service.Group, error) {
  252. // AccountCount is intentionally not loaded here; use GetByID when needed.
  253. m, err := r.client.Group.Query().
  254. Where(group.IDEQ(id)).
  255. Only(ctx)
  256. if err != nil {
  257. return nil, translatePersistenceError(err, service.ErrGroupNotFound, nil)
  258. }
  259. return groupEntityToService(m), nil
  260. }
  261. func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) error {
  262. modelPricing, err := json.Marshal(groupIn.ModelPricing)
  263. if err != nil {
  264. return fmt.Errorf("marshal group model pricing: %w", err)
  265. }
  266. builder := r.client.Group.UpdateOneID(groupIn.ID).
  267. SetName(groupIn.Name).
  268. SetDescription(groupIn.Description).
  269. SetPlatform(groupIn.Platform).
  270. SetRateMultiplier(groupIn.RateMultiplier).
  271. SetIsExclusive(groupIn.IsExclusive).
  272. SetSecurityPolicyEnabled(groupIn.SecurityPolicyEnabled).
  273. SetSecurityPolicyMode(service.NormalizeSecurityPolicyMode(groupIn.SecurityPolicyMode)).
  274. SetSecurityPolicyEmailEnabled(groupIn.SecurityPolicyEmailEnabled).
  275. SetStatus(groupIn.Status).
  276. SetSubscriptionType(groupIn.SubscriptionType).
  277. SetNillableDailyLimitUsd(groupIn.DailyLimitUSD).
  278. SetNillableWeeklyLimitUsd(groupIn.WeeklyLimitUSD).
  279. SetNillableMonthlyLimitUsd(groupIn.MonthlyLimitUSD).
  280. SetAllowImageGeneration(groupIn.AllowImageGeneration).
  281. SetAllowBatchImageGeneration(groupIn.AllowBatchImageGeneration).
  282. SetImageRateIndependent(groupIn.ImageRateIndependent).
  283. SetImageRateMultiplier(groupIn.ImageRateMultiplier).
  284. SetNillableImagePrice1k(groupIn.ImagePrice1K).
  285. SetNillableImagePrice2k(groupIn.ImagePrice2K).
  286. SetNillableImagePrice4k(groupIn.ImagePrice4K).
  287. SetBatchImageDiscountMultiplier(groupIn.BatchImageDiscountMultiplier).
  288. SetBatchImageHoldMultiplier(groupIn.BatchImageHoldMultiplier).
  289. SetVideoRateIndependent(groupIn.VideoRateIndependent).
  290. SetVideoRateMultiplier(groupIn.VideoRateMultiplier).
  291. SetNillableVideoPrice480p(groupIn.VideoPrice480P).
  292. SetNillableVideoPrice720p(groupIn.VideoPrice720P).
  293. SetNillableVideoPrice1080p(groupIn.VideoPrice1080P).
  294. SetVideoModelPrices(service.NormalizeVideoModelPrices(groupIn.VideoModelPrices)).
  295. SetLongContextPricingEnabled(groupIn.LongContextPricingEnabled).
  296. SetModelPricing(modelPricing).
  297. SetDefaultValidityDays(groupIn.DefaultValidityDays).
  298. SetClaudeCodeOnly(groupIn.ClaudeCodeOnly).
  299. SetModelRoutingEnabled(groupIn.ModelRoutingEnabled).
  300. SetMcpXMLInject(groupIn.MCPXMLInject).
  301. SetAllowMessagesDispatch(groupIn.AllowMessagesDispatch).
  302. SetAllowLive(groupIn.AllowLive).
  303. SetForceOpenaiFast(groupIn.ForceOpenAIFast).
  304. SetFreeOpenaiFast(groupIn.FreeOpenAIFast).
  305. SetRequireOauthOnly(groupIn.RequireOAuthOnly).
  306. SetRequirePrivacySet(groupIn.RequirePrivacySet).
  307. SetDefaultMappedModel(groupIn.DefaultMappedModel).
  308. SetMessagesDispatchModelConfig(groupIn.MessagesDispatchModelConfig).
  309. SetModelAllowlist(service.DomainGroupModelAllowlist(groupIn.ModelAllowlist)).
  310. SetCodexModelsManifestConfig(groupIn.CodexModelsManifestConfig).
  311. SetRpmLimit(groupIn.RPMLimit).
  312. SetMaxReasoningEffort(groupIn.MaxReasoningEffort).
  313. SetMaxReasoningEffortOverLimit(groupIn.MaxReasoningEffortOverLimit).
  314. SetReasoningEffortMappings(groupIn.ReasoningEffortMappings).
  315. SetPeakRateEnabled(groupIn.PeakRateEnabled).
  316. SetPeakStart(groupIn.PeakStart).
  317. SetPeakEnd(groupIn.PeakEnd).
  318. SetPeakRateMultiplier(groupIn.PeakRateMultiplier).
  319. SetProfitControlEnabled(groupIn.ProfitControlEnabled).
  320. SetProfitMinMargin(groupIn.ProfitMinMargin).
  321. SetProfitSafetyBuffer(groupIn.ProfitSafetyBuffer)
  322. // 显式处理可空字段:nil 需要 clear,非 nil 需要 set。
  323. if groupIn.DailyLimitUSD != nil {
  324. builder = builder.SetDailyLimitUsd(*groupIn.DailyLimitUSD)
  325. } else {
  326. builder = builder.ClearDailyLimitUsd()
  327. }
  328. if groupIn.WeeklyLimitUSD != nil {
  329. builder = builder.SetWeeklyLimitUsd(*groupIn.WeeklyLimitUSD)
  330. } else {
  331. builder = builder.ClearWeeklyLimitUsd()
  332. }
  333. if groupIn.MonthlyLimitUSD != nil {
  334. builder = builder.SetMonthlyLimitUsd(*groupIn.MonthlyLimitUSD)
  335. } else {
  336. builder = builder.ClearMonthlyLimitUsd()
  337. }
  338. if groupIn.ImagePrice1K != nil {
  339. builder = builder.SetImagePrice1k(*groupIn.ImagePrice1K)
  340. } else {
  341. builder = builder.ClearImagePrice1k()
  342. }
  343. if groupIn.ImagePrice2K != nil {
  344. builder = builder.SetImagePrice2k(*groupIn.ImagePrice2K)
  345. } else {
  346. builder = builder.ClearImagePrice2k()
  347. }
  348. if groupIn.ImagePrice4K != nil {
  349. builder = builder.SetImagePrice4k(*groupIn.ImagePrice4K)
  350. } else {
  351. builder = builder.ClearImagePrice4k()
  352. }
  353. if groupIn.VideoPrice480P != nil {
  354. builder = builder.SetVideoPrice480p(*groupIn.VideoPrice480P)
  355. } else {
  356. builder = builder.ClearVideoPrice480p()
  357. }
  358. if groupIn.VideoPrice720P != nil {
  359. builder = builder.SetVideoPrice720p(*groupIn.VideoPrice720P)
  360. } else {
  361. builder = builder.ClearVideoPrice720p()
  362. }
  363. if groupIn.VideoPrice1080P != nil {
  364. builder = builder.SetVideoPrice1080p(*groupIn.VideoPrice1080P)
  365. } else {
  366. builder = builder.ClearVideoPrice1080p()
  367. }
  368. if groupIn.WebSearchPricePerCall != nil {
  369. builder = builder.SetWebSearchPricePerCall(*groupIn.WebSearchPricePerCall)
  370. } else {
  371. builder = builder.ClearWebSearchPricePerCall()
  372. }
  373. if groupIn.SearchPricePer1k != nil {
  374. builder = builder.SetSearchPricePer1k(*groupIn.SearchPricePer1k)
  375. } else {
  376. builder = builder.ClearSearchPricePer1k()
  377. }
  378. if groupIn.AudioRealtimePricePerMin != nil {
  379. builder = builder.SetAudioRealtimePricePerMin(*groupIn.AudioRealtimePricePerMin)
  380. } else {
  381. builder = builder.ClearAudioRealtimePricePerMin()
  382. }
  383. if groupIn.AudioTTSPricePerMillionChars != nil {
  384. builder = builder.SetAudioTtsPricePerMillionChars(*groupIn.AudioTTSPricePerMillionChars)
  385. } else {
  386. builder = builder.ClearAudioTtsPricePerMillionChars()
  387. }
  388. if groupIn.AudioSTTPricePerHour != nil {
  389. builder = builder.SetAudioSttPricePerHour(*groupIn.AudioSTTPricePerHour)
  390. } else {
  391. builder = builder.ClearAudioSttPricePerHour()
  392. }
  393. // 处理 FallbackGroupID:nil 时清除,否则设置
  394. if groupIn.FallbackGroupID != nil {
  395. builder = builder.SetFallbackGroupID(*groupIn.FallbackGroupID)
  396. } else {
  397. builder = builder.ClearFallbackGroupID()
  398. }
  399. // 处理 FallbackGroupIDOnInvalidRequest:nil 时清除,否则设置
  400. if groupIn.FallbackGroupIDOnInvalidRequest != nil {
  401. builder = builder.SetFallbackGroupIDOnInvalidRequest(*groupIn.FallbackGroupIDOnInvalidRequest)
  402. } else {
  403. builder = builder.ClearFallbackGroupIDOnInvalidRequest()
  404. }
  405. // 处理 ModelRouting:nil 时清除,否则设置
  406. if groupIn.ModelRouting != nil {
  407. builder = builder.SetModelRouting(groupIn.ModelRouting)
  408. } else {
  409. builder = builder.ClearModelRouting()
  410. }
  411. // 处理 SupportedModelScopes(始终设置,空数组表示不限制)
  412. builder = builder.SetSupportedModelScopes(groupIn.SupportedModelScopes)
  413. updated, err := builder.Save(ctx)
  414. if err != nil {
  415. return translatePersistenceError(err, service.ErrGroupNotFound, service.ErrGroupExists)
  416. }
  417. groupIn.UpdatedAt = updated.UpdatedAt
  418. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupIn.ID, nil); err != nil {
  419. logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue group update failed: group=%d err=%v", groupIn.ID, err)
  420. }
  421. return nil
  422. }
  423. func (r *groupRepository) Delete(ctx context.Context, id int64) error {
  424. _, err := r.client.Group.Delete().Where(group.IDEQ(id)).Exec(ctx)
  425. if err != nil {
  426. return translatePersistenceError(err, service.ErrGroupNotFound, nil)
  427. }
  428. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &id, nil); err != nil {
  429. logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue group delete failed: group=%d err=%v", id, err)
  430. }
  431. return nil
  432. }
  433. func (r *groupRepository) List(ctx context.Context, params pagination.PaginationParams) ([]service.Group, *pagination.PaginationResult, error) {
  434. return r.ListWithFilters(ctx, params, "", "", "", nil)
  435. }
  436. func (r *groupRepository) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]service.Group, *pagination.PaginationResult, error) {
  437. q := r.client.Group.Query()
  438. return r.listWithFiltersQuery(ctx, q, params, platform, status, search, isExclusive)
  439. }
  440. func (r *groupRepository) ListBindableWithFilters(ctx context.Context, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]service.Group, *pagination.PaginationResult, error) {
  441. q := r.client.Group.Query().Where(group.PlatformNEQ(service.PlatformComposite))
  442. return r.listWithFiltersQuery(ctx, q, params, platform, status, search, isExclusive)
  443. }
  444. func (r *groupRepository) listWithFiltersQuery(ctx context.Context, q *dbent.GroupQuery, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]service.Group, *pagination.PaginationResult, error) {
  445. if platform != "" {
  446. q = q.Where(group.PlatformEQ(platform))
  447. }
  448. if status != "" {
  449. q = q.Where(group.StatusEQ(status))
  450. }
  451. if search != "" {
  452. q = q.Where(group.Or(
  453. group.NameContainsFold(search),
  454. group.DescriptionContainsFold(search),
  455. ))
  456. }
  457. if isExclusive != nil {
  458. q = q.Where(group.IsExclusiveEQ(*isExclusive))
  459. }
  460. total, err := q.Clone().Count(ctx)
  461. if err != nil {
  462. return nil, nil, err
  463. }
  464. if strings.EqualFold(strings.TrimSpace(params.SortBy), "account_count") {
  465. return r.listWithAccountCountSort(ctx, q, params, total)
  466. }
  467. groupsQuery := q.
  468. Offset(params.Offset()).
  469. Limit(params.Limit())
  470. for _, order := range groupListOrder(params) {
  471. groupsQuery = groupsQuery.Order(order)
  472. }
  473. groups, err := groupsQuery.All(ctx)
  474. if err != nil {
  475. return nil, nil, err
  476. }
  477. groupIDs := make([]int64, 0, len(groups))
  478. outGroups := make([]service.Group, 0, len(groups))
  479. for i := range groups {
  480. g := groupEntityToService(groups[i])
  481. outGroups = append(outGroups, *g)
  482. groupIDs = append(groupIDs, g.ID)
  483. }
  484. counts, err := r.loadAccountCounts(ctx, groupIDs)
  485. if err == nil {
  486. for i := range outGroups {
  487. c := counts[outGroups[i].ID]
  488. outGroups[i].AccountCount = c.Total
  489. outGroups[i].ActiveAccountCount = c.Active
  490. outGroups[i].RateLimitedAccountCount = c.RateLimited
  491. }
  492. }
  493. return outGroups, paginationResultFromTotal(int64(total), params), nil
  494. }
  495. func (r *groupRepository) listWithAccountCountSort(ctx context.Context, q *dbent.GroupQuery, params pagination.PaginationParams, total int) ([]service.Group, *pagination.PaginationResult, error) {
  496. // 第一步:只查 ID + sort_order(轻量,不做分页 — 需要全量排序 account_count)。
  497. rows, err := q.Clone().
  498. Select(group.FieldID, group.FieldSortOrder).
  499. Order(dbent.Asc(group.FieldSortOrder), dbent.Asc(group.FieldID)).
  500. All(ctx)
  501. if err != nil {
  502. return nil, nil, err
  503. }
  504. type sortEntry struct {
  505. id int64
  506. sortOrder int
  507. accountCount int64
  508. }
  509. entries := make([]sortEntry, 0, len(rows))
  510. groupIDs := make([]int64, len(rows))
  511. for i, r := range rows {
  512. groupIDs[i] = r.ID
  513. entries = append(entries, sortEntry{id: r.ID, sortOrder: r.SortOrder})
  514. }
  515. // 第二步:批量加载 account counts(一次 SQL)。
  516. counts, err := r.loadAccountCounts(ctx, groupIDs)
  517. if err != nil {
  518. return nil, nil, err
  519. }
  520. for i := range entries {
  521. c := counts[entries[i].id]
  522. if c.Total > 0 {
  523. entries[i].accountCount = c.Total
  524. }
  525. }
  526. // 第三步:Go 侧排序(数据量 = Group 总数,通常 < 200,安全)。
  527. sortOrder := params.NormalizedSortOrder(pagination.SortOrderDesc)
  528. tieCmp := func(a, b sortEntry) bool {
  529. if a.sortOrder == b.sortOrder {
  530. return a.id < b.id
  531. }
  532. return a.sortOrder < b.sortOrder
  533. }
  534. sort.SliceStable(entries, func(i, j int) bool {
  535. if entries[i].accountCount == entries[j].accountCount {
  536. return tieCmp(entries[i], entries[j])
  537. }
  538. if sortOrder == pagination.SortOrderAsc {
  539. return entries[i].accountCount < entries[j].accountCount
  540. }
  541. return entries[i].accountCount > entries[j].accountCount
  542. })
  543. // 第四步:分页,只加载当前页需要的完整 Group。
  544. page := paginateSlice(entries, params)
  545. if len(page) == 0 {
  546. return nil, paginationResultFromTotal(int64(total), params), nil
  547. }
  548. pageIDs := make([]int64, len(page))
  549. pageIdx := make(map[int64]int, len(page))
  550. for i, e := range page {
  551. pageIDs[i] = e.id
  552. pageIdx[e.id] = i
  553. }
  554. groups, err := r.client.Group.Query().
  555. Where(group.IDIn(pageIDs...)).
  556. All(ctx)
  557. if err != nil {
  558. return nil, nil, err
  559. }
  560. outGroups := make([]service.Group, len(page))
  561. for i := range groups {
  562. g := groupEntityToService(groups[i])
  563. c := counts[g.ID]
  564. g.AccountCount = c.Total
  565. g.ActiveAccountCount = c.Active
  566. g.RateLimitedAccountCount = c.RateLimited
  567. if idx, ok := pageIdx[g.ID]; ok {
  568. outGroups[idx] = *g
  569. }
  570. }
  571. return outGroups, paginationResultFromTotal(int64(total), params), nil
  572. }
  573. func groupListOrder(params pagination.PaginationParams) []func(*entsql.Selector) {
  574. sortBy := strings.ToLower(strings.TrimSpace(params.SortBy))
  575. sortOrder := params.NormalizedSortOrder(pagination.SortOrderAsc)
  576. var field string
  577. tieField := group.FieldID
  578. defaultOrder := true
  579. switch sortBy {
  580. case "", "sort_order":
  581. field = group.FieldSortOrder
  582. case "name":
  583. field = group.FieldName
  584. defaultOrder = false
  585. case "platform":
  586. field = group.FieldPlatform
  587. defaultOrder = false
  588. case "billing_type", "subscription_type":
  589. field = group.FieldSubscriptionType
  590. defaultOrder = false
  591. case "rate_multiplier":
  592. field = group.FieldRateMultiplier
  593. defaultOrder = false
  594. case "is_exclusive":
  595. field = group.FieldIsExclusive
  596. defaultOrder = false
  597. case "status":
  598. field = group.FieldStatus
  599. defaultOrder = false
  600. case "created_at":
  601. field = group.FieldCreatedAt
  602. defaultOrder = false
  603. case "id":
  604. field = group.FieldID
  605. defaultOrder = false
  606. tieField = ""
  607. default:
  608. field = group.FieldSortOrder
  609. }
  610. if sortOrder == pagination.SortOrderDesc && sortBy != "" {
  611. if tieField == "" {
  612. return []func(*entsql.Selector){dbent.Desc(field)}
  613. }
  614. return []func(*entsql.Selector){dbent.Desc(field), dbent.Desc(tieField)}
  615. }
  616. if defaultOrder {
  617. return []func(*entsql.Selector){dbent.Asc(group.FieldSortOrder), dbent.Asc(group.FieldID)}
  618. }
  619. if tieField == "" {
  620. return []func(*entsql.Selector){dbent.Asc(field)}
  621. }
  622. return []func(*entsql.Selector){dbent.Asc(field), dbent.Asc(tieField)}
  623. }
  624. func (r *groupRepository) ListActive(ctx context.Context) ([]service.Group, error) {
  625. groups, err := r.client.Group.Query().
  626. Where(group.StatusEQ(service.StatusActive)).
  627. Order(dbent.Asc(group.FieldSortOrder), dbent.Asc(group.FieldID)).
  628. All(ctx)
  629. if err != nil {
  630. return nil, err
  631. }
  632. groupIDs := make([]int64, 0, len(groups))
  633. outGroups := make([]service.Group, 0, len(groups))
  634. for i := range groups {
  635. g := groupEntityToService(groups[i])
  636. outGroups = append(outGroups, *g)
  637. groupIDs = append(groupIDs, g.ID)
  638. }
  639. counts, err := r.loadAccountCounts(ctx, groupIDs)
  640. if err == nil {
  641. for i := range outGroups {
  642. c := counts[outGroups[i].ID]
  643. outGroups[i].AccountCount = c.Total
  644. outGroups[i].ActiveAccountCount = c.Active
  645. outGroups[i].RateLimitedAccountCount = c.RateLimited
  646. }
  647. }
  648. return outGroups, nil
  649. }
  650. func (r *groupRepository) ListActiveIDs(ctx context.Context) ([]int64, error) {
  651. if r.sql != nil {
  652. rows, err := r.sql.QueryContext(ctx, `
  653. SELECT id
  654. FROM groups
  655. WHERE status = $1
  656. AND deleted_at IS NULL
  657. ORDER BY sort_order ASC, id ASC
  658. `, service.StatusActive)
  659. if err != nil {
  660. return nil, err
  661. }
  662. defer func() { _ = rows.Close() }()
  663. ids := make([]int64, 0)
  664. for rows.Next() {
  665. var id int64
  666. if err := rows.Scan(&id); err != nil {
  667. return nil, err
  668. }
  669. ids = append(ids, id)
  670. }
  671. if err := rows.Err(); err != nil {
  672. return nil, err
  673. }
  674. return ids, nil
  675. }
  676. groups, err := r.client.Group.Query().
  677. Where(group.StatusEQ(service.StatusActive)).
  678. Select(group.FieldID).
  679. Order(dbent.Asc(group.FieldSortOrder), dbent.Asc(group.FieldID)).
  680. All(ctx)
  681. if err != nil {
  682. return nil, err
  683. }
  684. ids := make([]int64, 0, len(groups))
  685. for i := range groups {
  686. ids = append(ids, groups[i].ID)
  687. }
  688. return ids, nil
  689. }
  690. func (r *groupRepository) ListActiveByPlatform(ctx context.Context, platform string) ([]service.Group, error) {
  691. groups, err := r.client.Group.Query().
  692. Where(group.StatusEQ(service.StatusActive), group.PlatformEQ(platform)).
  693. Order(dbent.Asc(group.FieldSortOrder), dbent.Asc(group.FieldID)).
  694. All(ctx)
  695. if err != nil {
  696. return nil, err
  697. }
  698. groupIDs := make([]int64, 0, len(groups))
  699. outGroups := make([]service.Group, 0, len(groups))
  700. for i := range groups {
  701. g := groupEntityToService(groups[i])
  702. outGroups = append(outGroups, *g)
  703. groupIDs = append(groupIDs, g.ID)
  704. }
  705. counts, err := r.loadAccountCounts(ctx, groupIDs)
  706. if err == nil {
  707. for i := range outGroups {
  708. c := counts[outGroups[i].ID]
  709. outGroups[i].AccountCount = c.Total
  710. outGroups[i].ActiveAccountCount = c.Active
  711. outGroups[i].RateLimitedAccountCount = c.RateLimited
  712. }
  713. }
  714. return outGroups, nil
  715. }
  716. func (r *groupRepository) ExistsByName(ctx context.Context, name string) (bool, error) {
  717. return r.client.Group.Query().Where(group.NameEQ(name)).Exist(ctx)
  718. }
  719. // ExistsByIDs 批量检查分组是否存在(仅检查未软删除记录)。
  720. // 返回结构:map[groupID]exists。
  721. func (r *groupRepository) ExistsByIDs(ctx context.Context, ids []int64) (map[int64]bool, error) {
  722. result := make(map[int64]bool, len(ids))
  723. if len(ids) == 0 {
  724. return result, nil
  725. }
  726. uniqueIDs := make([]int64, 0, len(ids))
  727. seen := make(map[int64]struct{}, len(ids))
  728. for _, id := range ids {
  729. if id <= 0 {
  730. continue
  731. }
  732. if _, ok := seen[id]; ok {
  733. continue
  734. }
  735. seen[id] = struct{}{}
  736. uniqueIDs = append(uniqueIDs, id)
  737. result[id] = false
  738. }
  739. if len(uniqueIDs) == 0 {
  740. return result, nil
  741. }
  742. rows, err := r.sql.QueryContext(ctx, `
  743. SELECT id
  744. FROM groups
  745. WHERE id = ANY($1) AND deleted_at IS NULL
  746. `, pq.Array(uniqueIDs))
  747. if err != nil {
  748. return nil, err
  749. }
  750. defer func() { _ = rows.Close() }()
  751. for rows.Next() {
  752. var id int64
  753. if err := rows.Scan(&id); err != nil {
  754. return nil, err
  755. }
  756. result[id] = true
  757. }
  758. if err := rows.Err(); err != nil {
  759. return nil, err
  760. }
  761. return result, nil
  762. }
  763. func (r *groupRepository) GetAccountCount(ctx context.Context, groupID int64) (total int64, active int64, err error) {
  764. var rateLimited int64
  765. err = scanSingleRow(ctx, r.sql,
  766. fmt.Sprintf(`SELECT
  767. COUNT(*) FILTER (WHERE a.deleted_at IS NULL),
  768. COUNT(*) FILTER (WHERE %s),
  769. COUNT(*) FILTER (WHERE %s)
  770. FROM account_groups ag JOIN accounts a ON a.id = ag.account_id
  771. WHERE ag.group_id = $1`, groupAccountAvailableSQL, groupAccountTemporarilyLimitedSQL),
  772. []any{groupID}, &total, &active, &rateLimited)
  773. return
  774. }
  775. func (r *groupRepository) DeleteAccountGroupsByGroupID(ctx context.Context, groupID int64) (int64, error) {
  776. res, err := r.sql.ExecContext(ctx, "DELETE FROM account_groups WHERE group_id = $1", groupID)
  777. if err != nil {
  778. return 0, err
  779. }
  780. affected, _ := res.RowsAffected()
  781. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupID, nil); err != nil {
  782. logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue group account clear failed: group=%d err=%v", groupID, err)
  783. }
  784. return affected, nil
  785. }
  786. func (r *groupRepository) DeleteCascade(ctx context.Context, id int64) ([]int64, error) {
  787. return r.deleteCascade(ctx, id, false)
  788. }
  789. func (r *groupRepository) DeleteCascadeIfEmpty(ctx context.Context, id int64) ([]int64, error) {
  790. return r.deleteCascade(ctx, id, true)
  791. }
  792. func (r *groupRepository) deleteCascade(ctx context.Context, id int64, requireEmpty bool) ([]int64, error) {
  793. // 使用 ent 事务统一包裹:避免手工基于 *sql.Tx 构造 ent client 带来的驱动断言问题,
  794. // 同时保证级联删除的原子性。
  795. tx, err := r.client.Tx(ctx)
  796. if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
  797. return nil, err
  798. }
  799. exec := r.client
  800. txClient := r.client
  801. if err == nil {
  802. defer func() { _ = tx.Rollback() }()
  803. exec = tx.Client()
  804. txClient = exec
  805. }
  806. // err 为 dbent.ErrTxStarted 时,复用当前 client 参与同一事务。
  807. // Lock the group row to avoid concurrent writes while we cascade.
  808. // 这里使用 exec.QueryContext 手动扫描,确保同一事务内加锁并能区分"未找到"与其他错误。
  809. rows, err := exec.QueryContext(ctx, "SELECT id, subscription_type FROM groups WHERE id = $1 AND deleted_at IS NULL FOR UPDATE", id)
  810. if err != nil {
  811. return nil, err
  812. }
  813. var lockedID int64
  814. var subscriptionType string
  815. if rows.Next() {
  816. if err := rows.Scan(&lockedID, &subscriptionType); err != nil {
  817. _ = rows.Close()
  818. return nil, err
  819. }
  820. }
  821. if err := rows.Close(); err != nil {
  822. return nil, err
  823. }
  824. if err := rows.Err(); err != nil {
  825. return nil, err
  826. }
  827. if lockedID == 0 {
  828. return nil, service.ErrGroupNotFound
  829. }
  830. if requireEmpty {
  831. var hasAccount bool
  832. if err := scanSingleRow(ctx, exec, `SELECT EXISTS (
  833. SELECT 1 FROM account_groups ag
  834. JOIN accounts a ON a.id = ag.account_id
  835. WHERE ag.group_id = $1 AND a.deleted_at IS NULL
  836. )`, []any{id}, &hasAccount); err != nil {
  837. return nil, err
  838. }
  839. if hasAccount {
  840. return nil, service.ErrGroupNotEmpty
  841. }
  842. }
  843. var affectedUserIDs []int64
  844. if subscriptionType == service.SubscriptionTypeSubscription {
  845. // 只查询未软删除的订阅,避免通知已取消订阅的用户
  846. rows, err := exec.QueryContext(ctx, "SELECT user_id FROM user_subscriptions WHERE group_id = $1 AND deleted_at IS NULL", id)
  847. if err != nil {
  848. return nil, err
  849. }
  850. for rows.Next() {
  851. var userID int64
  852. if scanErr := rows.Scan(&userID); scanErr != nil {
  853. _ = rows.Close()
  854. return nil, scanErr
  855. }
  856. affectedUserIDs = append(affectedUserIDs, userID)
  857. }
  858. if err := rows.Close(); err != nil {
  859. return nil, err
  860. }
  861. if err := rows.Err(); err != nil {
  862. return nil, err
  863. }
  864. // 软删除订阅:设置 deleted_at 而非硬删除
  865. if _, err := exec.ExecContext(ctx, "UPDATE user_subscriptions SET deleted_at = NOW() WHERE group_id = $1 AND deleted_at IS NULL", id); err != nil {
  866. return nil, err
  867. }
  868. }
  869. // 2. Remove the group id from user_allowed_groups join table.
  870. // Legacy users.allowed_groups 列已弃用,不再同步。
  871. if _, err := exec.ExecContext(ctx, "DELETE FROM user_allowed_groups WHERE group_id = $1", id); err != nil {
  872. return nil, err
  873. }
  874. // 3. Delete account_groups join rows.
  875. if _, err := exec.ExecContext(ctx, "DELETE FROM account_groups WHERE group_id = $1", id); err != nil {
  876. return nil, err
  877. }
  878. // 4. Soft-delete composite model routes owned by this group.
  879. if _, err := exec.ExecContext(ctx, "UPDATE composite_model_routes SET deleted_at = NOW() WHERE group_id = $1 AND deleted_at IS NULL", id); err != nil {
  880. return nil, err
  881. }
  882. // 5. Soft-delete group itself.
  883. if _, err := txClient.Group.Delete().Where(group.IDEQ(id)).Exec(ctx); err != nil {
  884. return nil, err
  885. }
  886. if tx != nil {
  887. if err := tx.Commit(); err != nil {
  888. return nil, err
  889. }
  890. }
  891. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &id, nil); err != nil {
  892. logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue group cascade delete failed: group=%d err=%v", id, err)
  893. }
  894. return affectedUserIDs, nil
  895. }
  896. type groupAccountCounts struct {
  897. Total int64
  898. Active int64
  899. RateLimited int64
  900. }
  901. const (
  902. // 分组页的"可用"账号数必须与账号仓储的 ListSchedulableByGroupID 过滤口径一致。
  903. groupAccountAvailableSQL = `a.deleted_at IS NULL
  904. AND a.status = 'active'
  905. AND a.schedulable = true
  906. AND (a.expires_at IS NULL OR a.expires_at > NOW() OR a.auto_pause_on_expired = FALSE)
  907. AND (a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW())
  908. AND (a.overload_until IS NULL OR a.overload_until <= NOW())
  909. AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW())`
  910. // 这里沿用历史字段名 RateLimitedAccountCount,但统计的是会让账号暂时退出调度的时间窗口。
  911. groupAccountTemporarilyLimitedSQL = `a.deleted_at IS NULL
  912. AND a.status = 'active'
  913. AND a.schedulable = true
  914. AND (a.expires_at IS NULL OR a.expires_at > NOW() OR a.auto_pause_on_expired = FALSE)
  915. AND (
  916. a.rate_limit_reset_at > NOW() OR
  917. a.overload_until > NOW() OR
  918. a.temp_unschedulable_until > NOW()
  919. )`
  920. )
  921. func (r *groupRepository) loadAccountCounts(ctx context.Context, groupIDs []int64) (counts map[int64]groupAccountCounts, err error) {
  922. counts = make(map[int64]groupAccountCounts, len(groupIDs))
  923. if len(groupIDs) == 0 {
  924. return counts, nil
  925. }
  926. rows, err := r.sql.QueryContext(
  927. ctx,
  928. fmt.Sprintf(`SELECT ag.group_id,
  929. COUNT(*) FILTER (WHERE a.deleted_at IS NULL) AS total,
  930. COUNT(*) FILTER (WHERE %s) AS active,
  931. COUNT(*) FILTER (WHERE %s) AS rate_limited
  932. FROM account_groups ag
  933. JOIN accounts a ON a.id = ag.account_id
  934. WHERE ag.group_id = ANY($1)
  935. GROUP BY ag.group_id`, groupAccountAvailableSQL, groupAccountTemporarilyLimitedSQL),
  936. pq.Array(groupIDs),
  937. )
  938. if err != nil {
  939. return nil, err
  940. }
  941. defer func() {
  942. if closeErr := rows.Close(); closeErr != nil && err == nil {
  943. err = closeErr
  944. counts = nil
  945. }
  946. }()
  947. for rows.Next() {
  948. var groupID int64
  949. var c groupAccountCounts
  950. if err = rows.Scan(&groupID, &c.Total, &c.Active, &c.RateLimited); err != nil {
  951. return nil, err
  952. }
  953. counts[groupID] = c
  954. }
  955. if err = rows.Err(); err != nil {
  956. return nil, err
  957. }
  958. return counts, nil
  959. }
  960. // GetAccountIDsByGroupIDs 获取多个分组的所有账号 ID(去重)
  961. func (r *groupRepository) GetAccountIDsByGroupIDs(ctx context.Context, groupIDs []int64) ([]int64, error) {
  962. if len(groupIDs) == 0 {
  963. return nil, nil
  964. }
  965. rows, err := r.sql.QueryContext(
  966. ctx,
  967. "SELECT DISTINCT account_id FROM account_groups WHERE group_id = ANY($1) ORDER BY account_id",
  968. pq.Array(groupIDs),
  969. )
  970. if err != nil {
  971. return nil, err
  972. }
  973. defer func() { _ = rows.Close() }()
  974. var accountIDs []int64
  975. for rows.Next() {
  976. var accountID int64
  977. if err := rows.Scan(&accountID); err != nil {
  978. return nil, err
  979. }
  980. accountIDs = append(accountIDs, accountID)
  981. }
  982. if err := rows.Err(); err != nil {
  983. return nil, err
  984. }
  985. return accountIDs, nil
  986. }
  987. // BindAccountsToGroup 将多个账号绑定到指定分组(批量插入,忽略已存在的绑定)
  988. func (r *groupRepository) BindAccountsToGroup(ctx context.Context, groupID int64, accountIDs []int64) error {
  989. if len(accountIDs) == 0 {
  990. return nil
  991. }
  992. tx, err := r.client.Tx(ctx)
  993. if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
  994. return err
  995. }
  996. exec := sqlExecutor(r.client)
  997. if tx != nil {
  998. defer func() { _ = tx.Rollback() }()
  999. exec = tx.Client()
  1000. }
  1001. if err := lockLiveGroups(ctx, exec, []int64{groupID}); err != nil {
  1002. return err
  1003. }
  1004. // 使用 INSERT ... ON CONFLICT DO NOTHING 忽略已存在的绑定
  1005. _, err = exec.ExecContext(
  1006. ctx,
  1007. `INSERT INTO account_groups (account_id, group_id, priority, created_at)
  1008. SELECT unnest($1::bigint[]), $2, 50, NOW()
  1009. ON CONFLICT (account_id, group_id) DO NOTHING`,
  1010. pq.Array(accountIDs),
  1011. groupID,
  1012. )
  1013. if err != nil {
  1014. return err
  1015. }
  1016. if tx != nil {
  1017. if err := tx.Commit(); err != nil {
  1018. return err
  1019. }
  1020. }
  1021. // 发送调度器事件
  1022. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupID, nil); err != nil {
  1023. logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue bind accounts to group failed: group=%d err=%v", groupID, err)
  1024. }
  1025. return nil
  1026. }
  1027. // UpdateSortOrders 批量更新分组排序
  1028. func (r *groupRepository) UpdateSortOrders(ctx context.Context, updates []service.GroupSortOrderUpdate) error {
  1029. if len(updates) == 0 {
  1030. return nil
  1031. }
  1032. // 去重后保留最后一次排序值,避免重复 ID 造成 CASE 分支冲突。
  1033. sortOrderByID := make(map[int64]int, len(updates))
  1034. groupIDs := make([]int64, 0, len(updates))
  1035. for _, u := range updates {
  1036. if u.ID <= 0 {
  1037. continue
  1038. }
  1039. if _, exists := sortOrderByID[u.ID]; !exists {
  1040. groupIDs = append(groupIDs, u.ID)
  1041. }
  1042. sortOrderByID[u.ID] = u.SortOrder
  1043. }
  1044. if len(groupIDs) == 0 {
  1045. return nil
  1046. }
  1047. // 与旧实现保持一致:任何不存在/已删除的分组都返回 not found,且不执行更新。
  1048. var existingCount int
  1049. if err := scanSingleRow(
  1050. ctx,
  1051. r.sql,
  1052. `SELECT COUNT(*) FROM groups WHERE deleted_at IS NULL AND id = ANY($1)`,
  1053. []any{pq.Array(groupIDs)},
  1054. &existingCount,
  1055. ); err != nil {
  1056. return err
  1057. }
  1058. if existingCount != len(groupIDs) {
  1059. return service.ErrGroupNotFound
  1060. }
  1061. args := make([]any, 0, len(groupIDs)*2+1)
  1062. caseClauses := make([]string, 0, len(groupIDs))
  1063. placeholder := 1
  1064. for _, id := range groupIDs {
  1065. caseClauses = append(caseClauses, fmt.Sprintf("WHEN $%d THEN $%d", placeholder, placeholder+1))
  1066. args = append(args, id, sortOrderByID[id])
  1067. placeholder += 2
  1068. }
  1069. args = append(args, pq.Array(groupIDs))
  1070. query := fmt.Sprintf(`
  1071. UPDATE groups
  1072. SET sort_order = CASE id
  1073. %s
  1074. ELSE sort_order
  1075. END
  1076. WHERE deleted_at IS NULL AND id = ANY($%d)
  1077. `, strings.Join(caseClauses, "\n\t\t\t"), placeholder)
  1078. result, err := r.sql.ExecContext(ctx, query, args...)
  1079. if err != nil {
  1080. return err
  1081. }
  1082. affected, err := result.RowsAffected()
  1083. if err != nil {
  1084. return err
  1085. }
  1086. if affected != int64(len(groupIDs)) {
  1087. return service.ErrGroupNotFound
  1088. }
  1089. for _, id := range groupIDs {
  1090. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &id, nil); err != nil {
  1091. logger.LegacyPrintf("repository.group", "[SchedulerOutbox] enqueue group sort update failed: group=%d err=%v", id, err)
  1092. }
  1093. }
  1094. return nil
  1095. }