account_repo.go 125 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021302230233024302530263027302830293030303130323033303430353036303730383039304030413042304330443045304630473048304930503051305230533054305530563057305830593060306130623063306430653066306730683069307030713072307330743075307630773078307930803081308230833084308530863087308830893090309130923093309430953096309730983099310031013102310331043105310631073108310931103111311231133114311531163117311831193120312131223123312431253126312731283129313031313132313331343135313631373138313931403141314231433144314531463147314831493150315131523153315431553156315731583159316031613162316331643165316631673168316931703171317231733174317531763177317831793180318131823183318431853186318731883189319031913192319331943195319631973198319932003201320232033204320532063207320832093210321132123213321432153216321732183219322032213222322332243225322632273228322932303231323232333234323532363237323832393240324132423243324432453246324732483249325032513252325332543255325632573258325932603261326232633264326532663267326832693270327132723273327432753276327732783279328032813282328332843285328632873288328932903291329232933294329532963297329832993300330133023303330433053306330733083309331033113312331333143315331633173318331933203321332233233324332533263327332833293330333133323333333433353336333733383339334033413342334333443345334633473348334933503351335233533354335533563357335833593360336133623363336433653366336733683369337033713372337333743375337633773378337933803381338233833384338533863387338833893390339133923393339433953396339733983399340034013402340334043405340634073408340934103411341234133414341534163417341834193420342134223423342434253426342734283429343034313432343334343435343634373438343934403441344234433444344534463447344834493450345134523453345434553456345734583459346034613462346334643465346634673468346934703471347234733474347534763477347834793480348134823483348434853486348734883489349034913492349334943495349634973498349935003501350235033504350535063507350835093510351135123513351435153516351735183519352035213522352335243525352635273528352935303531353235333534353535363537353835393540354135423543354435453546354735483549355035513552355335543555355635573558355935603561356235633564356535663567356835693570357135723573357435753576357735783579358035813582358335843585358635873588358935903591359235933594359535963597359835993600360136023603360436053606360736083609361036113612361336143615361636173618361936203621362236233624362536263627362836293630363136323633363436353636363736383639364036413642364336443645364636473648364936503651365236533654365536563657365836593660366136623663366436653666366736683669367036713672367336743675367636773678367936803681368236833684368536863687368836893690369136923693369436953696369736983699370037013702370337043705370637073708370937103711371237133714371537163717371837193720372137223723372437253726372737283729373037313732373337343735373637373738373937403741374237433744374537463747374837493750375137523753375437553756375737583759376037613762376337643765376637673768376937703771377237733774377537763777377837793780378137823783378437853786378737883789379037913792379337943795379637973798379938003801380238033804380538063807380838093810381138123813381438153816381738183819382038213822382338243825382638273828382938303831383238333834383538363837383838393840384138423843384438453846384738483849385038513852385338543855385638573858385938603861386238633864386538663867386838693870387138723873387438753876387738783879388038813882388338843885388638873888388938903891389238933894389538963897389838993900390139023903390439053906390739083909391039113912391339143915
  1. // Package repository 实现数据访问层(Repository Pattern)。
  2. //
  3. // 该包提供了与数据库交互的所有操作,包括 CRUD、复杂查询和批量操作。
  4. // 采用 Repository 模式将数据访问逻辑与业务逻辑分离,便于测试和维护。
  5. //
  6. // 主要特性:
  7. // - 使用 Ent ORM 进行类型安全的数据库操作
  8. // - 对于复杂查询(如批量更新、聚合统计)使用原生 SQL
  9. // - 提供统一的错误翻译机制,将数据库错误转换为业务错误
  10. // - 支持软删除,所有查询自动过滤已删除记录
  11. package repository
  12. import (
  13. "context"
  14. "database/sql"
  15. "encoding/json"
  16. "errors"
  17. "math/rand/v2"
  18. "strconv"
  19. "strings"
  20. "time"
  21. dbent "github.com/Wei-Shaw/sub2api/ent"
  22. dbaccount "github.com/Wei-Shaw/sub2api/ent/account"
  23. dbaccountgroup "github.com/Wei-Shaw/sub2api/ent/accountgroup"
  24. dbgroup "github.com/Wei-Shaw/sub2api/ent/group"
  25. dbpredicate "github.com/Wei-Shaw/sub2api/ent/predicate"
  26. dbproxy "github.com/Wei-Shaw/sub2api/ent/proxy"
  27. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  28. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  29. "github.com/Wei-Shaw/sub2api/internal/service"
  30. "github.com/lib/pq"
  31. entsql "entgo.io/ent/dialect/sql"
  32. "entgo.io/ent/dialect/sql/sqljson"
  33. )
  34. // accountRepository 实现 service.AccountRepository 接口。
  35. // 提供 AI API 账户的完整数据访问功能。
  36. //
  37. // 设计说明:
  38. // - client: Ent 客户端,用于类型安全的 ORM 操作
  39. // - sql: 原生 SQL 执行器,用于复杂查询和批量操作
  40. // - schedulerCache: 调度器缓存,用于在账号状态变更时同步快照
  41. type accountRepository struct {
  42. client *dbent.Client // Ent ORM 客户端
  43. sql sqlExecutor // 原生 SQL 执行接口
  44. // schedulerCache 用于在账号状态变更时主动同步快照到缓存,
  45. // 确保粘性会话能及时感知账号不可用状态。
  46. // Used to proactively sync account snapshot to cache when status changes,
  47. // ensuring sticky sessions can promptly detect unavailable accounts.
  48. schedulerCache service.SchedulerCache
  49. }
  50. var schedulerNeutralExtraKeyPrefixes = []string{
  51. "codex_primary_",
  52. "codex_secondary_",
  53. "codex_5h_",
  54. "codex_7d_",
  55. "codex_reset_credit_",
  56. "passive_usage_",
  57. "upstream_billing_probe",
  58. "upstream_billing_rate_sync",
  59. "ollama_cloud_usage",
  60. }
  61. var schedulerNeutralExtraKeys = map[string]struct{}{
  62. "codex_usage_updated_at": {},
  63. "grok_billing_snapshot": {},
  64. "session_window_utilization": {},
  65. }
  66. const postgresParameterBatchSize = 50000
  67. const codexFingerprintSeedCanonicalPattern = "^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
  68. const codexFingerprintNilSeed = "00000000-0000-0000-0000-000000000000"
  69. func codexFingerprintSeedValidSQL(extraExpr string) string {
  70. value := "(" + extraExpr + " ->> 'codex_fingerprint_seed')"
  71. return "(" + value + " ~ '" + codexFingerprintSeedCanonicalPattern + "' AND " + value + " <> '" + codexFingerprintNilSeed + "')"
  72. }
  73. func ensureCodexFingerprintSeedSQL(extraExpr string) string {
  74. return "CASE WHEN platform = 'openai' AND type = 'oauth' THEN " +
  75. "jsonb_set(" + extraExpr + ", '{codex_fingerprint_seed}', " +
  76. "CASE WHEN " + codexFingerprintSeedValidSQL("extra") +
  77. " THEN to_jsonb(extra ->> 'codex_fingerprint_seed') ELSE to_jsonb(gen_random_uuid()::text) END, true) " +
  78. "ELSE " + extraExpr + " END"
  79. }
  80. func stripCodexFingerprintSeedFromExtraUpdate(extra map[string]any) map[string]any {
  81. if extra == nil {
  82. return nil
  83. }
  84. if _, exists := extra["codex_fingerprint_seed"]; !exists {
  85. return extra
  86. }
  87. stripped := make(map[string]any, len(extra)-1)
  88. for key, value := range extra {
  89. if key == "codex_fingerprint_seed" {
  90. continue
  91. }
  92. stripped[key] = value
  93. }
  94. return stripped
  95. }
  96. // NewAccountRepository 创建账户仓储实例。
  97. // 这是对外暴露的构造函数,返回接口类型以便于依赖注入。
  98. func NewAccountRepository(client *dbent.Client, sqlDB *sql.DB, schedulerCache service.SchedulerCache) service.AccountRepository {
  99. return newAccountRepositoryWithSQL(client, sqlDB, schedulerCache)
  100. }
  101. // NewAdminAccountRepository exposes the account repository's atomic duplication capability
  102. // as an explicit dependency of the admin service.
  103. func NewAdminAccountRepository(client *dbent.Client, sqlDB *sql.DB, schedulerCache service.SchedulerCache) service.AdminAccountRepository {
  104. return newAccountRepositoryWithSQL(client, sqlDB, schedulerCache)
  105. }
  106. // newAccountRepositoryWithSQL 是内部构造函数,支持依赖注入 SQL 执行器。
  107. // 这种设计便于单元测试时注入 mock 对象。
  108. func newAccountRepositoryWithSQL(client *dbent.Client, sqlq sqlExecutor, schedulerCache service.SchedulerCache) *accountRepository {
  109. return &accountRepository{client: client, sql: sqlq, schedulerCache: schedulerCache}
  110. }
  111. func (r *accountRepository) Create(ctx context.Context, account *service.Account) error {
  112. if err := createAccountRecord(ctx, r.client, account); err != nil {
  113. return err
  114. }
  115. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &account.ID, nil, buildSchedulerGroupPayload(account.GroupIDs)); err != nil {
  116. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue account create failed: account=%d err=%v", account.ID, err)
  117. }
  118. return nil
  119. }
  120. func createAccountRecord(ctx context.Context, client *dbent.Client, account *service.Account) error {
  121. if account == nil {
  122. return service.ErrAccountNilInput
  123. }
  124. service.PrepareNewAccountProtection(account)
  125. builder := client.Account.Create().
  126. SetName(account.Name).
  127. SetNillableNotes(account.Notes).
  128. SetPlatform(account.Platform).
  129. SetType(account.Type).
  130. SetCredentials(normalizeJSONMap(account.Credentials)).
  131. SetExtra(normalizeJSONMap(account.Extra)).
  132. SetConcurrency(account.Concurrency).
  133. SetPriority(account.Priority).
  134. SetStatus(account.Status).
  135. SetErrorMessage(account.ErrorMessage).
  136. SetSchedulable(account.Schedulable).
  137. SetAutoPauseOnExpired(account.AutoPauseOnExpired)
  138. if account.RateMultiplier != nil {
  139. builder.SetRateMultiplier(*account.RateMultiplier)
  140. }
  141. if account.LoadFactor != nil {
  142. builder.SetLoadFactor(*account.LoadFactor)
  143. }
  144. if account.ProxyID != nil {
  145. builder.SetProxyID(*account.ProxyID)
  146. }
  147. if account.LastUsedAt != nil {
  148. builder.SetLastUsedAt(*account.LastUsedAt)
  149. }
  150. if account.ExpiresAt != nil {
  151. builder.SetExpiresAt(*account.ExpiresAt)
  152. }
  153. if account.RateLimitedAt != nil {
  154. builder.SetRateLimitedAt(*account.RateLimitedAt)
  155. }
  156. if account.RateLimitResetAt != nil {
  157. builder.SetRateLimitResetAt(*account.RateLimitResetAt)
  158. }
  159. if account.OverloadUntil != nil {
  160. builder.SetOverloadUntil(*account.OverloadUntil)
  161. }
  162. if account.SessionWindowStart != nil {
  163. builder.SetSessionWindowStart(*account.SessionWindowStart)
  164. }
  165. if account.SessionWindowEnd != nil {
  166. builder.SetSessionWindowEnd(*account.SessionWindowEnd)
  167. }
  168. if account.SessionWindowStatus != "" {
  169. builder.SetSessionWindowStatus(account.SessionWindowStatus)
  170. }
  171. builder.SetQuotaDimension(dbaccount.QuotaDimension(account.QuotaDimensionOrDefault()))
  172. if account.ParentAccountID != nil {
  173. builder.SetParentAccountID(*account.ParentAccountID)
  174. }
  175. created, err := builder.Save(ctx)
  176. if err != nil {
  177. return translatePersistenceError(err, service.ErrAccountNotFound, nil)
  178. }
  179. account.ID = created.ID
  180. account.CreatedAt = created.CreatedAt
  181. account.UpdatedAt = created.UpdatedAt
  182. return nil
  183. }
  184. // CreateWithAccountGroups atomically persists an account, its exact per-group priorities,
  185. // and the scheduler outbox event used to publish the new routing snapshot.
  186. func (r *accountRepository) CreateWithAccountGroups(ctx context.Context, account *service.Account, groups []service.AccountGroup) error {
  187. if account == nil {
  188. return service.ErrAccountNilInput
  189. }
  190. tx, err := r.client.Tx(ctx)
  191. if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
  192. return err
  193. }
  194. var txClient *dbent.Client
  195. if err == nil {
  196. defer func() { _ = tx.Rollback() }()
  197. txClient = tx.Client()
  198. } else {
  199. // Reuse a caller-owned transaction when this repository is already transactional.
  200. txClient = r.client
  201. }
  202. groupIDs := make([]int64, 0, len(groups))
  203. for i := range groups {
  204. groupIDs = append(groupIDs, groups[i].GroupID)
  205. }
  206. if err := lockLiveGroups(ctx, txClient, groupIDs); err != nil {
  207. return err
  208. }
  209. if err := createAccountRecord(ctx, txClient, account); err != nil {
  210. return err
  211. }
  212. if len(groups) > 0 {
  213. builders := make([]*dbent.AccountGroupCreate, 0, len(groups))
  214. for i := range groups {
  215. groups[i].AccountID = account.ID
  216. builders = append(builders, txClient.AccountGroup.Create().
  217. SetAccountID(account.ID).
  218. SetGroupID(groups[i].GroupID).
  219. SetPriority(groups[i].Priority),
  220. )
  221. }
  222. if _, err := txClient.AccountGroup.CreateBulk(builders...).Save(ctx); err != nil {
  223. return err
  224. }
  225. }
  226. account.GroupIDs = groupIDs
  227. account.AccountGroups = append([]service.AccountGroup(nil), groups...)
  228. if err := enqueueSchedulerOutbox(ctx, txClient, service.SchedulerOutboxEventAccountChanged, &account.ID, nil, buildSchedulerGroupPayload(groupIDs)); err != nil {
  229. return err
  230. }
  231. if tx != nil {
  232. if err := tx.Commit(); err != nil {
  233. return err
  234. }
  235. }
  236. return nil
  237. }
  238. func (r *accountRepository) GetByID(ctx context.Context, id int64) (*service.Account, error) {
  239. m, err := r.client.Account.Query().Where(dbaccount.IDEQ(id)).Only(ctx)
  240. if err != nil {
  241. return nil, translatePersistenceError(err, service.ErrAccountNotFound, nil)
  242. }
  243. accounts, err := r.accountsToService(ctx, []*dbent.Account{m})
  244. if err != nil {
  245. return nil, err
  246. }
  247. if len(accounts) == 0 {
  248. return nil, service.ErrAccountNotFound
  249. }
  250. return &accounts[0], nil
  251. }
  252. // SelectRandomActiveProxy chooses one active, non-expired proxy for an
  253. // account configured with proxy_mode=random. The query is intentionally
  254. // performed at request hydration time so the account can rotate across the
  255. // current pool without persisting a random choice.
  256. //
  257. // Sampling uses COUNT + OFFSET instead of ORDER BY RANDOM(): the latter
  258. // sorts the whole proxy table on every request and degrades as the pool
  259. // grows, while the count/offset pair is two indexed queries.
  260. func (r *accountRepository) SelectRandomActiveProxy(ctx context.Context) (*service.Proxy, error) {
  261. if r == nil || r.client == nil {
  262. return nil, nil
  263. }
  264. eligible := func() *dbent.ProxyQuery {
  265. return r.client.Proxy.Query().
  266. Where(
  267. dbproxy.StatusEQ(service.StatusActive),
  268. dbproxy.DeletedAtIsNil(),
  269. dbproxy.Or(
  270. dbproxy.ExpiresAtIsNil(),
  271. dbproxy.ExpiresAtGT(time.Now()),
  272. ),
  273. )
  274. }
  275. total, err := eligible().Count(ctx)
  276. if err != nil {
  277. return nil, err
  278. }
  279. if total == 0 {
  280. return nil, nil
  281. }
  282. offset := 0
  283. if total > 1 {
  284. offset = rand.IntN(total)
  285. }
  286. // Deterministic ID order keeps OFFSET sampling well-defined. A row that
  287. // disappears between COUNT and SELECT surfaces as NotFound and is
  288. // treated like an empty pool so callers fail closed.
  289. item, err := eligible().Order(dbproxy.ByID()).Offset(offset).First(ctx)
  290. if err != nil {
  291. if dbent.IsNotFound(err) {
  292. return nil, nil
  293. }
  294. return nil, err
  295. }
  296. return proxyEntityToService(item), nil
  297. }
  298. func (r *accountRepository) GetByIDs(ctx context.Context, ids []int64) ([]*service.Account, error) {
  299. if len(ids) == 0 {
  300. return []*service.Account{}, nil
  301. }
  302. // De-duplicate while preserving order of first occurrence.
  303. uniqueIDs := make([]int64, 0, len(ids))
  304. seen := make(map[int64]struct{}, len(ids))
  305. for _, id := range ids {
  306. if id <= 0 {
  307. continue
  308. }
  309. if _, ok := seen[id]; ok {
  310. continue
  311. }
  312. seen[id] = struct{}{}
  313. uniqueIDs = append(uniqueIDs, id)
  314. }
  315. if len(uniqueIDs) == 0 {
  316. return []*service.Account{}, nil
  317. }
  318. entAccounts, err := r.client.Account.
  319. Query().
  320. Where(dbaccount.IDIn(uniqueIDs...)).
  321. WithProxy().
  322. All(ctx)
  323. if err != nil {
  324. return nil, err
  325. }
  326. if len(entAccounts) == 0 {
  327. return []*service.Account{}, nil
  328. }
  329. accountIDs := make([]int64, 0, len(entAccounts))
  330. entByID := make(map[int64]*dbent.Account, len(entAccounts))
  331. for _, acc := range entAccounts {
  332. entByID[acc.ID] = acc
  333. accountIDs = append(accountIDs, acc.ID)
  334. }
  335. groupsByAccount, groupIDsByAccount, accountGroupsByAccount, err := r.loadAccountGroups(ctx, accountIDs)
  336. if err != nil {
  337. return nil, err
  338. }
  339. outByID := make(map[int64]*service.Account, len(entAccounts))
  340. for _, entAcc := range entAccounts {
  341. out := accountEntityToService(entAcc)
  342. if out == nil {
  343. continue
  344. }
  345. // Prefer the preloaded proxy edge when available.
  346. if entAcc.Edges.Proxy != nil {
  347. out.Proxy = proxyEntityToService(entAcc.Edges.Proxy)
  348. }
  349. if groups, ok := groupsByAccount[entAcc.ID]; ok {
  350. out.Groups = groups
  351. }
  352. if groupIDs, ok := groupIDsByAccount[entAcc.ID]; ok {
  353. out.GroupIDs = groupIDs
  354. }
  355. if ags, ok := accountGroupsByAccount[entAcc.ID]; ok {
  356. out.AccountGroups = ags
  357. }
  358. outByID[entAcc.ID] = out
  359. }
  360. // Preserve input order (first occurrence), and ignore missing IDs.
  361. out := make([]*service.Account, 0, len(uniqueIDs))
  362. for _, id := range uniqueIDs {
  363. if _, ok := entByID[id]; !ok {
  364. continue
  365. }
  366. if acc, ok := outByID[id]; ok && acc != nil {
  367. out = append(out, acc)
  368. }
  369. }
  370. return out, nil
  371. }
  372. // ExistsByID 检查指定 ID 的账号是否存在。
  373. // 相比 GetByID,此方法性能更优,因为:
  374. // - 使用 Exist() 方法生成 SELECT EXISTS 查询,只返回布尔值
  375. // - 不加载完整的账号实体及其关联数据(Groups、Proxy 等)
  376. // - 适用于删除前的存在性检查等只需判断有无的场景
  377. func (r *accountRepository) ExistsByID(ctx context.Context, id int64) (bool, error) {
  378. exists, err := r.client.Account.Query().Where(dbaccount.IDEQ(id)).Exist(ctx)
  379. if err != nil {
  380. return false, err
  381. }
  382. return exists, nil
  383. }
  384. func (r *accountRepository) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*service.Account, error) {
  385. if crsAccountID == "" {
  386. return nil, nil
  387. }
  388. // 使用 sqljson.ValueEQ 生成 JSON 路径过滤,避免手写 SQL 片段导致语法兼容问题。
  389. // 排除 spark 影子账号(parent_account_id 非空):影子不持凭据,绝不能被 CRS 当作普通账号
  390. // 更新而覆盖 type/credentials/proxy。即便影子 Extra 被误写入 crs_account_id 也不会命中
  391. // (外审第7轮 P1)。
  392. m, err := r.client.Account.Query().
  393. Where(dbaccount.ParentAccountIDIsNil()).
  394. Where(func(s *entsql.Selector) {
  395. s.Where(sqljson.ValueEQ(dbaccount.FieldExtra, crsAccountID, sqljson.Path("crs_account_id")))
  396. }).
  397. Only(ctx)
  398. if err != nil {
  399. if dbent.IsNotFound(err) {
  400. return nil, nil
  401. }
  402. return nil, err
  403. }
  404. accounts, err := r.accountsToService(ctx, []*dbent.Account{m})
  405. if err != nil {
  406. return nil, err
  407. }
  408. if len(accounts) == 0 {
  409. return nil, nil
  410. }
  411. return &accounts[0], nil
  412. }
  413. func (r *accountRepository) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
  414. // parent_account_id IS NULL 排除 spark 影子账号:影子不是 CRS 账号,绝不能进 CRS 同步映射
  415. // (否则会被当普通账号更新而覆盖 type/credentials/proxy)(外审第7轮 P1)。
  416. rows, err := r.sql.QueryContext(ctx, `
  417. SELECT id, extra->>'crs_account_id'
  418. FROM accounts
  419. WHERE deleted_at IS NULL
  420. AND parent_account_id IS NULL
  421. AND extra->>'crs_account_id' IS NOT NULL
  422. AND extra->>'crs_account_id' != ''
  423. `)
  424. if err != nil {
  425. return nil, err
  426. }
  427. defer func() { _ = rows.Close() }()
  428. result := make(map[string]int64)
  429. for rows.Next() {
  430. var id int64
  431. var crsID string
  432. if err := rows.Scan(&id, &crsID); err != nil {
  433. return nil, err
  434. }
  435. result[crsID] = id
  436. }
  437. if err := rows.Err(); err != nil {
  438. return nil, err
  439. }
  440. return result, nil
  441. }
  442. func (r *accountRepository) Update(ctx context.Context, account *service.Account) error {
  443. return r.updateAccount(ctx, account, nil, nil, account.RateMultiplier)
  444. }
  445. // UpdateWithAccountBillingSettings applies an admin account edit while
  446. // preserving a concurrently probe-synchronized rate unless the request
  447. // explicitly includes a manual rate.
  448. func (r *accountRepository) UpdateWithAccountBillingSettings(
  449. ctx context.Context,
  450. account *service.Account,
  451. probeEnabled *bool,
  452. rateSyncEnabled *bool,
  453. rateMultiplier *float64,
  454. ) error {
  455. return r.updateAccount(ctx, account, probeEnabled, rateSyncEnabled, rateMultiplier)
  456. }
  457. func (r *accountRepository) updateAccount(
  458. ctx context.Context,
  459. account *service.Account,
  460. explicitProbeEnabled *bool,
  461. explicitRateSyncEnabled *bool,
  462. explicitRateMultiplier *float64,
  463. ) error {
  464. if account == nil {
  465. return nil
  466. }
  467. baseCtx := ctx
  468. contextTx := dbent.TxFromContext(ctx)
  469. client := r.client
  470. var tx *dbent.Tx
  471. if contextTx != nil {
  472. client = contextTx.Client()
  473. } else {
  474. var err error
  475. tx, err = r.client.Tx(ctx)
  476. if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
  477. return err
  478. }
  479. if tx != nil {
  480. defer func() { _ = tx.Rollback() }()
  481. ctx = dbent.NewTxContext(ctx, tx)
  482. client = tx.Client()
  483. }
  484. }
  485. updated, err := r.updateLockedAccount(
  486. ctx,
  487. client,
  488. account,
  489. explicitProbeEnabled,
  490. explicitRateSyncEnabled,
  491. explicitRateMultiplier,
  492. )
  493. if err != nil {
  494. return translatePersistenceError(err, service.ErrAccountNotFound, nil)
  495. }
  496. if err := enqueueSchedulerOutbox(ctx, client, service.SchedulerOutboxEventAccountChanged, &account.ID, nil, buildSchedulerGroupPayload(account.GroupIDs)); err != nil {
  497. return err
  498. }
  499. if tx != nil {
  500. if err := tx.Commit(); err != nil {
  501. return err
  502. }
  503. }
  504. account.UpdatedAt = updated.UpdatedAt
  505. // 普通账号编辑(如 model_mapping / credentials)也需要立即刷新单账号快照,
  506. // 否则网关在 outbox worker 延迟或异常时仍可能读到旧配置。
  507. if contextTx == nil {
  508. r.syncSchedulerAccountSnapshot(baseCtx, account.ID)
  509. }
  510. return nil
  511. }
  512. func (r *accountRepository) updateLockedAccount(
  513. ctx context.Context,
  514. client *dbent.Client,
  515. account *service.Account,
  516. explicitProbeEnabled *bool,
  517. explicitRateSyncEnabled *bool,
  518. explicitRateMultiplier *float64,
  519. ) (*dbent.Account, error) {
  520. extra, err := lockAndMergeAccountProbeExtra(ctx, client, account, explicitProbeEnabled, explicitRateSyncEnabled)
  521. if err != nil {
  522. return nil, err
  523. }
  524. account.Extra = extra
  525. if err := preserveLockedAccountProtection(ctx, client, account); err != nil {
  526. return nil, err
  527. }
  528. extra = normalizeJSONMap(account.Extra)
  529. schedulable := account.Schedulable
  530. if account.Status == service.StatusError {
  531. schedulable = false
  532. }
  533. builder := client.Account.UpdateOneID(account.ID).
  534. SetName(account.Name).
  535. SetNillableNotes(account.Notes).
  536. SetPlatform(account.Platform).
  537. SetType(account.Type).
  538. SetCredentials(normalizeJSONMap(account.Credentials)).
  539. SetExtra(extra).
  540. SetConcurrency(account.Concurrency).
  541. SetPriority(account.Priority).
  542. SetStatus(account.Status).
  543. SetErrorMessage(account.ErrorMessage).
  544. SetSchedulable(schedulable).
  545. SetAutoPauseOnExpired(account.AutoPauseOnExpired)
  546. if explicitRateMultiplier != nil {
  547. builder.SetRateMultiplier(*explicitRateMultiplier)
  548. }
  549. if account.LoadFactor != nil {
  550. builder.SetLoadFactor(*account.LoadFactor)
  551. } else {
  552. builder.ClearLoadFactor()
  553. }
  554. if account.IsRandomProxy() {
  555. builder.ClearProxyID()
  556. } else if account.ProxyID != nil {
  557. builder.SetProxyID(*account.ProxyID)
  558. } else {
  559. builder.ClearProxyID()
  560. }
  561. if account.LastUsedAt != nil {
  562. builder.SetLastUsedAt(*account.LastUsedAt)
  563. } else {
  564. builder.ClearLastUsedAt()
  565. }
  566. if account.ExpiresAt != nil {
  567. builder.SetExpiresAt(*account.ExpiresAt)
  568. } else {
  569. builder.ClearExpiresAt()
  570. }
  571. if account.RateLimitedAt != nil {
  572. builder.SetRateLimitedAt(*account.RateLimitedAt)
  573. } else {
  574. builder.ClearRateLimitedAt()
  575. }
  576. if account.RateLimitResetAt != nil {
  577. builder.SetRateLimitResetAt(*account.RateLimitResetAt)
  578. } else {
  579. builder.ClearRateLimitResetAt()
  580. }
  581. if account.OverloadUntil != nil {
  582. builder.SetOverloadUntil(*account.OverloadUntil)
  583. } else {
  584. builder.ClearOverloadUntil()
  585. }
  586. if account.SessionWindowStart != nil {
  587. builder.SetSessionWindowStart(*account.SessionWindowStart)
  588. } else {
  589. builder.ClearSessionWindowStart()
  590. }
  591. if account.SessionWindowEnd != nil {
  592. builder.SetSessionWindowEnd(*account.SessionWindowEnd)
  593. } else {
  594. builder.ClearSessionWindowEnd()
  595. }
  596. if account.SessionWindowStatus != "" {
  597. builder.SetSessionWindowStatus(account.SessionWindowStatus)
  598. } else {
  599. builder.ClearSessionWindowStatus()
  600. }
  601. if account.Notes == nil {
  602. builder.ClearNotes()
  603. }
  604. builder.SetQuotaDimension(dbaccount.QuotaDimension(account.QuotaDimensionOrDefault()))
  605. builder.SetNillableParentAccountID(account.ParentAccountID)
  606. return builder.Save(ctx)
  607. }
  608. func lockAndMergeAccountProbeExtra(
  609. ctx context.Context,
  610. client *dbent.Client,
  611. account *service.Account,
  612. explicitProbeEnabled *bool,
  613. explicitRateSyncEnabled *bool,
  614. ) (map[string]any, error) {
  615. credentials, err := json.Marshal(normalizeJSONMap(account.Credentials))
  616. if err != nil {
  617. return nil, err
  618. }
  619. var proxyID any
  620. if account.ProxyID != nil {
  621. proxyID = *account.ProxyID
  622. }
  623. rows, err := client.QueryContext(ctx, `
  624. SELECT
  625. platform = $2
  626. AND type = $3
  627. AND credentials = $4::jsonb
  628. AND proxy_id IS NOT DISTINCT FROM $5,
  629. COALESCE(
  630. platform IN (`+ollamaCloudUsagePlatformsSQL+`)
  631. AND $2 IN (`+ollamaCloudUsagePlatformsSQL+`)
  632. AND type = 'apikey'
  633. AND $3 = 'apikey'
  634. AND credentials -> 'api_key' IS NOT DISTINCT FROM $4::jsonb -> 'api_key'
  635. AND `+ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'")+`
  636. AND `+ollamaCloudBaseURLMatchesSQL("$4::jsonb ->> 'base_url'")+`,
  637. false
  638. ),
  639. proxy_id IS NOT DISTINCT FROM $5,
  640. extra -> 'upstream_billing_probe_enabled',
  641. extra -> 'upstream_billing_rate_sync_enabled',
  642. extra -> 'upstream_billing_probe',
  643. extra -> 'ollama_cloud_usage_session',
  644. extra -> 'ollama_cloud_usage_auto_refresh',
  645. extra -> 'ollama_cloud_usage_snapshot'
  646. FROM accounts
  647. WHERE id = $1 AND deleted_at IS NULL
  648. FOR NO KEY UPDATE
  649. `, account.ID, account.Platform, account.Type, string(credentials), proxyID)
  650. if err != nil {
  651. return nil, err
  652. }
  653. defer func() { _ = rows.Close() }()
  654. if !rows.Next() {
  655. if err := rows.Err(); err != nil {
  656. return nil, err
  657. }
  658. return nil, service.ErrAccountNotFound
  659. }
  660. var (
  661. identityUnchanged bool
  662. ollamaGroupIdentityUnchanged bool
  663. ollamaProxyIdentityUnchanged bool
  664. currentEnabled []byte
  665. currentRateSyncEnabled []byte
  666. currentSnapshot []byte
  667. currentOllamaSession []byte
  668. currentOllamaAutoRefresh []byte
  669. currentOllamaSnapshot []byte
  670. )
  671. if err := rows.Scan(
  672. &identityUnchanged,
  673. &ollamaGroupIdentityUnchanged,
  674. &ollamaProxyIdentityUnchanged,
  675. &currentEnabled,
  676. &currentRateSyncEnabled,
  677. &currentSnapshot,
  678. &currentOllamaSession,
  679. &currentOllamaAutoRefresh,
  680. &currentOllamaSnapshot,
  681. ); err != nil {
  682. return nil, err
  683. }
  684. if err := rows.Err(); err != nil {
  685. return nil, err
  686. }
  687. extra := copyJSONMap(normalizeJSONMap(account.Extra))
  688. for _, key := range []string{
  689. service.UpstreamBillingProbeEnabledExtraKey,
  690. service.UpstreamBillingRateSyncEnabledExtraKey,
  691. service.UpstreamBillingProbeExtraKey,
  692. service.OllamaCloudUsageSessionExtraKey,
  693. service.OllamaCloudUsageAutoRefreshExtraKey,
  694. service.OllamaCloudUsageSnapshotExtraKey,
  695. } {
  696. delete(extra, key)
  697. }
  698. probeAccount := service.IsUpstreamBillingProbeIdentity(account.Platform, account.Type)
  699. probeEnabled := false
  700. probeEnabledPresent := false
  701. if probeAccount {
  702. if enabled, ok, err := decodeAccountExtraJSON(currentEnabled); err != nil {
  703. return nil, err
  704. } else if value, isBool := enabled.(bool); ok && isBool {
  705. probeEnabled = value
  706. probeEnabledPresent = true
  707. }
  708. if explicitProbeEnabled != nil {
  709. probeEnabled = *explicitProbeEnabled
  710. probeEnabledPresent = true
  711. }
  712. }
  713. rateSyncEnabled := false
  714. rateSyncEnabledPresent := false
  715. if probeAccount {
  716. if enabled, ok, err := decodeAccountExtraJSON(currentRateSyncEnabled); err != nil {
  717. return nil, err
  718. } else if value, isBool := enabled.(bool); ok && isBool {
  719. rateSyncEnabled = value
  720. rateSyncEnabledPresent = true
  721. }
  722. if explicitRateSyncEnabled != nil {
  723. rateSyncEnabled = *explicitRateSyncEnabled
  724. rateSyncEnabledPresent = true
  725. }
  726. if explicitProbeEnabled != nil && !*explicitProbeEnabled {
  727. rateSyncEnabled = false
  728. rateSyncEnabledPresent = true
  729. }
  730. // 同步依赖探测,方向是单向的:探测关闭(或探测键缺失)一律把同步归零。
  731. // 不做反向推导——由 rate_sync=true 推出 probe=true 会让一条"同步开、探测键
  732. // 缺失"的僵尸记录在任意一次无关编辑时静默打开周期性外呼。需要同时打开两个
  733. // 开关的调用方(管理端编辑)自己显式传 explicitProbeEnabled=true。
  734. if !probeEnabled {
  735. rateSyncEnabled = false
  736. }
  737. if probeEnabledPresent {
  738. extra[service.UpstreamBillingProbeEnabledExtraKey] = probeEnabled
  739. }
  740. if rateSyncEnabledPresent {
  741. extra[service.UpstreamBillingRateSyncEnabledExtraKey] = rateSyncEnabled
  742. }
  743. }
  744. probeExplicitlyDisabled := probeEnabledPresent && !probeEnabled
  745. if identityUnchanged && !probeExplicitlyDisabled {
  746. if snapshot, ok, err := decodeAccountExtraJSON(currentSnapshot); err != nil {
  747. return nil, err
  748. } else if ok {
  749. extra[service.UpstreamBillingProbeExtraKey] = snapshot
  750. }
  751. }
  752. if service.IsOllamaCloudUsageAccount(account) && ollamaGroupIdentityUnchanged {
  753. for key, raw := range map[string][]byte{
  754. service.OllamaCloudUsageSessionExtraKey: currentOllamaSession,
  755. service.OllamaCloudUsageAutoRefreshExtraKey: currentOllamaAutoRefresh,
  756. } {
  757. if value, ok, err := decodeAccountExtraJSON(raw); err != nil {
  758. return nil, err
  759. } else if ok {
  760. extra[key] = value
  761. }
  762. }
  763. if ollamaProxyIdentityUnchanged {
  764. if snapshot, ok, err := decodeAccountExtraJSON(currentOllamaSnapshot); err != nil {
  765. return nil, err
  766. } else if ok {
  767. extra[service.OllamaCloudUsageSnapshotExtraKey] = snapshot
  768. }
  769. }
  770. }
  771. return extra, nil
  772. }
  773. func decodeAccountExtraJSON(raw []byte) (any, bool, error) {
  774. if len(raw) == 0 || string(raw) == "null" {
  775. return nil, false, nil
  776. }
  777. var value any
  778. if err := json.Unmarshal(raw, &value); err != nil {
  779. return nil, false, err
  780. }
  781. return value, true, nil
  782. }
  783. func (r *accountRepository) UpdateCredentials(ctx context.Context, id int64, credentials map[string]any) error {
  784. payload, err := json.Marshal(normalizeJSONMap(credentials))
  785. if err != nil {
  786. return err
  787. }
  788. baseCtx := ctx
  789. contextTx := dbent.TxFromContext(ctx)
  790. client := r.client
  791. var tx *dbent.Tx
  792. if contextTx != nil {
  793. client = contextTx.Client()
  794. } else if r.client != nil {
  795. var txErr error
  796. tx, txErr = r.client.Tx(ctx)
  797. if txErr != nil && !errors.Is(txErr, dbent.ErrTxStarted) {
  798. return txErr
  799. }
  800. if tx != nil {
  801. defer func() { _ = tx.Rollback() }()
  802. ctx = dbent.NewTxContext(ctx, tx)
  803. client = tx.Client()
  804. }
  805. }
  806. result, err := client.ExecContext(ctx, `
  807. UPDATE accounts
  808. SET
  809. credentials = $1::jsonb,
  810. extra = CASE
  811. -- 凭证整体未变化 ⇒ Ollama 组身份必然未变化;顶层 DISTINCT 守卫防止
  812. -- 非 Ollama 账号的无变化持久化误清探测快照或重写 NULL extra。
  813. WHEN platform IN (`+ollamaCloudUsagePlatformsSQL+`)
  814. AND type = 'apikey'
  815. AND credentials IS DISTINCT FROM $1::jsonb
  816. AND (
  817. credentials -> 'api_key' IS DISTINCT FROM $1::jsonb -> 'api_key'
  818. OR NOT (
  819. `+ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'")+`
  820. AND `+ollamaCloudBaseURLMatchesSQL("$1::jsonb ->> 'base_url'")+`
  821. )
  822. )
  823. THEN COALESCE(extra, '{}'::jsonb)
  824. - 'upstream_billing_probe'
  825. - 'ollama_cloud_usage_session'
  826. - 'ollama_cloud_usage_auto_refresh'
  827. - 'ollama_cloud_usage_snapshot'
  828. -- 上游倍率探测已放宽到全部 API-key 平台:凭证变化即视为探测
  829. -- 身份变化,丢弃 stale 快照。
  830. WHEN type = 'apikey'
  831. AND credentials IS DISTINCT FROM $1::jsonb
  832. THEN COALESCE(extra, '{}'::jsonb) - 'upstream_billing_probe'
  833. ELSE extra
  834. END,
  835. updated_at = NOW()
  836. WHERE id = $2 AND deleted_at IS NULL
  837. `, string(payload), id)
  838. if err != nil {
  839. return err
  840. }
  841. affected, err := result.RowsAffected()
  842. if err != nil {
  843. return err
  844. }
  845. if affected == 0 {
  846. return service.ErrAccountNotFound
  847. }
  848. if err := enqueueSchedulerOutbox(ctx, client, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  849. return err
  850. }
  851. if tx != nil {
  852. if err := tx.Commit(); err != nil {
  853. return err
  854. }
  855. }
  856. if contextTx == nil {
  857. r.syncSchedulerAccountSnapshot(baseCtx, id)
  858. }
  859. return nil
  860. }
  861. func (r *accountRepository) Delete(ctx context.Context, id int64) error {
  862. groupIDs, err := r.loadAccountGroupIDs(ctx, id)
  863. if err != nil {
  864. return err
  865. }
  866. // 使用事务保证账号与关联分组的删除原子性
  867. tx, err := r.client.Tx(ctx)
  868. if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
  869. return err
  870. }
  871. var txClient *dbent.Client
  872. if err == nil {
  873. defer func() { _ = tx.Rollback() }()
  874. txClient = tx.Client()
  875. } else {
  876. // 已处于外部事务中(ErrTxStarted),复用当前 client
  877. txClient = r.client
  878. }
  879. if _, err := txClient.AccountGroup.Delete().Where(dbaccountgroup.AccountIDEQ(id)).Exec(ctx); err != nil {
  880. return err
  881. }
  882. if _, err := txClient.ExecContext(ctx, "DELETE FROM scheduled_test_plans WHERE account_id = $1", id); err != nil {
  883. return err
  884. }
  885. if _, err := txClient.Account.Delete().Where(dbaccount.IDEQ(id)).Exec(ctx); err != nil {
  886. return err
  887. }
  888. if tx != nil {
  889. if err := tx.Commit(); err != nil {
  890. return err
  891. }
  892. }
  893. r.deleteSchedulerAccountSnapshot(ctx, id)
  894. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, buildSchedulerGroupPayload(groupIDs)); err != nil {
  895. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue account delete failed: account=%d err=%v", id, err)
  896. }
  897. return nil
  898. }
  899. func (r *accountRepository) List(ctx context.Context, params pagination.PaginationParams) ([]service.Account, *pagination.PaginationResult, error) {
  900. return r.ListWithFilters(ctx, params, "", "", "", "", 0, "")
  901. }
  902. func (r *accountRepository) accountListFilteredQuery(platform, accountType, status, search string, groupID int64, privacyMode string) *dbent.AccountQuery {
  903. q := r.client.Account.Query()
  904. if platform != "" {
  905. q = q.Where(dbaccount.PlatformEQ(platform))
  906. }
  907. if accountType != "" {
  908. q = q.Where(dbaccount.TypeEQ(accountType))
  909. }
  910. if status != "" {
  911. switch status {
  912. case service.StatusActive:
  913. q = q.Where(
  914. dbaccount.StatusEQ(status),
  915. dbaccount.SchedulableEQ(true),
  916. dbaccount.Or(
  917. dbaccount.RateLimitResetAtIsNil(),
  918. dbaccount.RateLimitResetAtLTE(time.Now()),
  919. ),
  920. dbpredicate.Account(func(s *entsql.Selector) {
  921. col := s.C("temp_unschedulable_until")
  922. s.Where(entsql.Or(
  923. entsql.IsNull(col),
  924. entsql.LTE(col, entsql.Expr("NOW()")),
  925. ))
  926. }),
  927. )
  928. case "rate_limited":
  929. q = q.Where(
  930. dbaccount.StatusEQ(service.StatusActive),
  931. dbaccount.RateLimitResetAtGT(time.Now()),
  932. dbpredicate.Account(func(s *entsql.Selector) {
  933. col := s.C("temp_unschedulable_until")
  934. s.Where(entsql.Or(
  935. entsql.IsNull(col),
  936. entsql.LTE(col, entsql.Expr("NOW()")),
  937. ))
  938. }),
  939. )
  940. case "temp_unschedulable":
  941. q = q.Where(
  942. dbaccount.StatusEQ(service.StatusActive),
  943. dbpredicate.Account(func(s *entsql.Selector) {
  944. col := s.C("temp_unschedulable_until")
  945. s.Where(entsql.And(
  946. entsql.Not(entsql.IsNull(col)),
  947. entsql.GT(col, entsql.Expr("NOW()")),
  948. ))
  949. }),
  950. )
  951. case "unschedulable":
  952. q = q.Where(
  953. dbaccount.StatusEQ(service.StatusActive),
  954. dbaccount.SchedulableEQ(false),
  955. dbaccount.Or(
  956. dbaccount.RateLimitResetAtIsNil(),
  957. dbaccount.RateLimitResetAtLTE(time.Now()),
  958. ),
  959. dbpredicate.Account(func(s *entsql.Selector) {
  960. col := s.C("temp_unschedulable_until")
  961. s.Where(entsql.Or(
  962. entsql.IsNull(col),
  963. entsql.LTE(col, entsql.Expr("NOW()")),
  964. ))
  965. }),
  966. )
  967. default:
  968. q = q.Where(dbaccount.StatusEQ(status))
  969. }
  970. }
  971. if search != "" {
  972. q = q.Where(dbaccount.NameContainsFold(search))
  973. }
  974. if groupID == service.AccountListGroupUngrouped {
  975. q = q.Where(dbaccount.Not(dbaccount.HasAccountGroups()))
  976. } else if groupID > 0 {
  977. q = q.Where(dbaccount.HasAccountGroupsWith(dbaccountgroup.GroupIDEQ(groupID)))
  978. }
  979. if privacyMode != "" {
  980. q = q.Where(dbpredicate.Account(func(s *entsql.Selector) {
  981. path := sqljson.Path("privacy_mode")
  982. switch privacyMode {
  983. case service.AccountPrivacyModeUnsetFilter:
  984. s.Where(entsql.Or(
  985. entsql.Not(sqljson.HasKey(dbaccount.FieldExtra, path)),
  986. sqljson.ValueEQ(dbaccount.FieldExtra, "", path),
  987. ))
  988. default:
  989. s.Where(sqljson.ValueEQ(dbaccount.FieldExtra, privacyMode, path))
  990. }
  991. }))
  992. }
  993. return q
  994. }
  995. func (r *accountRepository) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]service.Account, *pagination.PaginationResult, error) {
  996. q := r.accountListFilteredQuery(platform, accountType, status, search, groupID, privacyMode)
  997. // Clone before Count so interceptor-appended predicates (SoftDeleteMixin's
  998. // deleted_at IS NULL) don't accumulate on the shared builder and pollute the
  999. // subsequent list query. Same pattern used in group_repo/promo_code_repo/user_repo
  1000. // (P1-03 audit fix, commit 2588fa6a).
  1001. total, err := q.Clone().Count(ctx)
  1002. if err != nil {
  1003. return nil, nil, err
  1004. }
  1005. accountsQuery := q.
  1006. Offset(params.Offset()).
  1007. Limit(params.Limit())
  1008. for _, order := range accountListOrder(params) {
  1009. accountsQuery = accountsQuery.Order(order)
  1010. }
  1011. accounts, err := accountsQuery.All(ctx)
  1012. if err != nil {
  1013. return nil, nil, err
  1014. }
  1015. outAccounts, err := r.accountsToService(ctx, accounts)
  1016. if err != nil {
  1017. return nil, nil, err
  1018. }
  1019. return outAccounts, paginationResultFromTotal(int64(total), params), nil
  1020. }
  1021. func (r *accountRepository) ListAllWithFilters(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]service.Account, error) {
  1022. accounts, err := r.accountListFilteredQuery(platform, accountType, status, search, groupID, privacyMode).All(ctx)
  1023. if err != nil {
  1024. return nil, err
  1025. }
  1026. return r.accountsToService(ctx, accounts)
  1027. }
  1028. func (r *accountRepository) ListOpsAccountsForStats(ctx context.Context, platformFilter string, groupIDFilter *int64) ([]service.Account, error) {
  1029. if r == nil || r.client == nil {
  1030. return []service.Account{}, nil
  1031. }
  1032. q := r.client.Account.Query()
  1033. if platformFilter = strings.TrimSpace(platformFilter); platformFilter != "" {
  1034. q = q.Where(dbaccount.PlatformEQ(platformFilter))
  1035. }
  1036. if groupIDFilter != nil && *groupIDFilter > 0 {
  1037. q = q.Where(dbaccount.HasAccountGroupsWith(dbaccountgroup.GroupIDEQ(*groupIDFilter)))
  1038. }
  1039. accounts, err := q.
  1040. Select(
  1041. dbaccount.FieldID,
  1042. dbaccount.FieldName,
  1043. dbaccount.FieldPlatform,
  1044. dbaccount.FieldConcurrency,
  1045. dbaccount.FieldLoadFactor,
  1046. dbaccount.FieldStatus,
  1047. dbaccount.FieldErrorMessage,
  1048. dbaccount.FieldSchedulable,
  1049. dbaccount.FieldRateLimitResetAt,
  1050. dbaccount.FieldOverloadUntil,
  1051. dbaccount.FieldTempUnschedulableUntil,
  1052. ).
  1053. Order(dbent.Asc(dbaccount.FieldID)).
  1054. All(ctx)
  1055. if err != nil {
  1056. return nil, err
  1057. }
  1058. return r.accountsToService(ctx, accounts)
  1059. }
  1060. func accountListOrder(params pagination.PaginationParams) []func(*entsql.Selector) {
  1061. sortBy := strings.ToLower(strings.TrimSpace(params.SortBy))
  1062. sortOrder := params.NormalizedSortOrder(pagination.SortOrderAsc)
  1063. if sortBy == "upstream_billing_rate" {
  1064. direction := "ASC"
  1065. tieOrder := entsql.Asc
  1066. if sortOrder == pagination.SortOrderDesc {
  1067. direction = "DESC"
  1068. tieOrder = entsql.Desc
  1069. }
  1070. return []func(*entsql.Selector){func(s *entsql.Selector) {
  1071. extra := s.C(dbaccount.FieldExtra)
  1072. expression := upstreamBillingRateSortExpression(extra)
  1073. s.OrderExpr(entsql.Expr(expression + " " + direction + " NULLS LAST"))
  1074. s.OrderBy(tieOrder(s.C(dbaccount.FieldID)))
  1075. }}
  1076. }
  1077. field := dbaccount.FieldName
  1078. defaultOrder := true
  1079. switch sortBy {
  1080. case "", "name":
  1081. field = dbaccount.FieldName
  1082. case "id":
  1083. field = dbaccount.FieldID
  1084. defaultOrder = false
  1085. case "status":
  1086. field = dbaccount.FieldStatus
  1087. defaultOrder = false
  1088. case "schedulable":
  1089. field = dbaccount.FieldSchedulable
  1090. defaultOrder = false
  1091. case "priority":
  1092. field = dbaccount.FieldPriority
  1093. defaultOrder = false
  1094. case "rate_multiplier":
  1095. field = dbaccount.FieldRateMultiplier
  1096. defaultOrder = false
  1097. case "last_used_at":
  1098. field = dbaccount.FieldLastUsedAt
  1099. defaultOrder = false
  1100. case "expires_at":
  1101. field = dbaccount.FieldExpiresAt
  1102. defaultOrder = false
  1103. case "created_at":
  1104. field = dbaccount.FieldCreatedAt
  1105. defaultOrder = false
  1106. }
  1107. if sortOrder == pagination.SortOrderDesc {
  1108. return []func(*entsql.Selector){dbent.Desc(field), dbent.Desc(dbaccount.FieldID)}
  1109. }
  1110. if defaultOrder {
  1111. return []func(*entsql.Selector){dbent.Asc(dbaccount.FieldName), dbent.Asc(dbaccount.FieldID)}
  1112. }
  1113. return []func(*entsql.Selector){dbent.Asc(field), dbent.Asc(dbaccount.FieldID)}
  1114. }
  1115. func upstreamBillingRateSortExpression(extra string) string {
  1116. status := extra + " #>> '{upstream_billing_probe,status}'"
  1117. effectiveJSON := extra + " #> '{upstream_billing_probe,data,effective_rate_multiplier}'"
  1118. effective := extra + " #>> '{upstream_billing_probe,data,effective_rate_multiplier}'"
  1119. resolvedJSON := extra + " #> '{upstream_billing_probe,data,resolved_rate_multiplier}'"
  1120. resolved := extra + " #>> '{upstream_billing_probe,data,resolved_rate_multiplier}'"
  1121. peakEnabledJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_enabled}'"
  1122. peakEnabled := extra + " #>> '{upstream_billing_probe,data,peak_rate_enabled}'"
  1123. peakStart := extra + " #>> '{upstream_billing_probe,data,peak_start}'"
  1124. peakEnd := extra + " #>> '{upstream_billing_probe,data,peak_end}'"
  1125. peakMultiplierJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_multiplier}'"
  1126. peakMultiplier := extra + " #>> '{upstream_billing_probe,data,peak_rate_multiplier}'"
  1127. peakMultiplierValue := "(CASE WHEN jsonb_typeof(" + peakMultiplierJSON + ") = 'number' THEN (" + peakMultiplier + ")::numeric END)"
  1128. billingScope := extra + " #>> '{upstream_billing_probe,data,billing_scope}'"
  1129. timezone := extra + " #>> '{upstream_billing_probe,data,timezone}'"
  1130. validClock := "'^([01][0-9]|2[0-3]):[0-5][0-9]$'"
  1131. startMinute := "(CASE WHEN " + peakStart + " ~ " + validClock + " THEN split_part(" + peakStart + ", ':', 1)::numeric * 60 + split_part(" + peakStart + ", ':', 2)::numeric END)"
  1132. endMinute := "(CASE WHEN " + peakEnd + " ~ " + validClock + " THEN split_part(" + peakEnd + ", ':', 1)::numeric * 60 + split_part(" + peakEnd + ", ':', 2)::numeric END)"
  1133. localMinute := "(EXTRACT(HOUR FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))) * 60 + EXTRACT(MINUTE FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))))"
  1134. validPeakWindow := peakStart + " ~ " + validClock + " AND " +
  1135. peakEnd + " ~ " + validClock + " AND " +
  1136. startMinute + " < " + endMinute
  1137. validPeakConfig := validPeakWindow + " AND " + peakMultiplierValue + " >= 0 AND " +
  1138. "EXISTS (SELECT 1 FROM pg_timezone_names WHERE name = " + timezone + ")"
  1139. dynamicRate := "CASE WHEN " + peakEnabled + " = 'false' THEN (" + resolved + ")::numeric WHEN " + peakEnabled + " = 'true' AND " + validPeakConfig +
  1140. " THEN (" + resolved + ")::numeric * CASE WHEN " + localMinute + " >= " + startMinute + " AND " + localMinute + " < " + endMinute +
  1141. " THEN " + peakMultiplierValue + " ELSE 1 END ELSE NULL END"
  1142. legacySnapshot := "jsonb_typeof(" + resolvedJSON + ") IS NULL AND jsonb_typeof(" + peakEnabledJSON + ") IS NULL"
  1143. return "CASE WHEN " + status + " IN ('ok', 'failed') AND (jsonb_typeof(" + resolvedJSON + ") = 'number' OR jsonb_typeof(" + effectiveJSON + ") = 'number') THEN CASE WHEN jsonb_typeof(" +
  1144. resolvedJSON + ") = 'number' AND jsonb_typeof(" + peakEnabledJSON + ") = 'boolean' THEN CASE WHEN " + billingScope + " = 'token' THEN " + dynamicRate + " ELSE NULL END WHEN " + legacySnapshot +
  1145. " AND jsonb_typeof(" + effectiveJSON + ") = 'number' THEN (" + effective + ")::numeric END END"
  1146. }
  1147. func (r *accountRepository) ListByGroup(ctx context.Context, groupID int64) ([]service.Account, error) {
  1148. accounts, err := r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{
  1149. status: service.StatusActive,
  1150. })
  1151. if err != nil {
  1152. return nil, err
  1153. }
  1154. return accounts, nil
  1155. }
  1156. func (r *accountRepository) ListActive(ctx context.Context) ([]service.Account, error) {
  1157. accounts, err := r.client.Account.Query().
  1158. Where(dbaccount.StatusEQ(service.StatusActive)).
  1159. Order(dbent.Asc(dbaccount.FieldPriority)).
  1160. All(ctx)
  1161. if err != nil {
  1162. return nil, err
  1163. }
  1164. return r.accountsToService(ctx, accounts)
  1165. }
  1166. func (r *accountRepository) ListOAuthRefreshCandidatePage(ctx context.Context, options service.OAuthRefreshPageOptions) (*service.OAuthRefreshCandidatePage, error) {
  1167. if r.sql == nil {
  1168. return nil, errors.New("account repository SQL executor not configured")
  1169. }
  1170. if len(options.Platforms) == 0 {
  1171. return nil, errors.New("oauth refresh candidate platforms cannot be empty")
  1172. }
  1173. if options.Limit <= 0 || options.Limit > 1000 {
  1174. return nil, errors.New("oauth refresh candidate page limit must be between 1 and 1000")
  1175. }
  1176. // (cond) IS NOT TRUE 把 NULL 和 FALSE 都视为"可被刷新"。直接写
  1177. // NOT (a AND b) 在 PG 三值逻辑下会把 a 或 b 为 NULL 的行(即绝大多数
  1178. // 健康账号:temp_unschedulable_until=NULL)也排除,导致后台 token
  1179. // 刷新工作器漏掉所有正常账号 → access_token 到期后请求开始 401。
  1180. query := `
  1181. SELECT id
  1182. FROM accounts
  1183. WHERE deleted_at IS NULL
  1184. AND schedulable = TRUE
  1185. AND platform = ANY($1)
  1186. AND id > $2`
  1187. if options.ActiveOnly {
  1188. query += `
  1189. AND status = 'active'`
  1190. }
  1191. if options.IncludeSetupToken {
  1192. query += `
  1193. AND type IN ('oauth', 'setup-token')`
  1194. } else {
  1195. query += `
  1196. AND type = 'oauth'`
  1197. }
  1198. if options.RequireRefreshToken {
  1199. query += `
  1200. AND credentials ? 'refresh_token'
  1201. AND btrim(credentials->>'refresh_token') <> ''`
  1202. }
  1203. if options.ExcludeRetryCooldown {
  1204. query += `
  1205. AND (
  1206. temp_unschedulable_until > NOW()
  1207. AND temp_unschedulable_reason LIKE 'token refresh retry exhausted:%'
  1208. ) IS NOT TRUE`
  1209. }
  1210. query += `
  1211. ORDER BY id ASC
  1212. LIMIT $3`
  1213. rows, err := r.sql.QueryContext(ctx, query, pq.Array(options.Platforms), options.AfterID, options.Limit)
  1214. if err != nil {
  1215. return nil, err
  1216. }
  1217. defer func() { _ = rows.Close() }()
  1218. var ids []int64
  1219. for rows.Next() {
  1220. var id int64
  1221. if err := rows.Scan(&id); err != nil {
  1222. return nil, err
  1223. }
  1224. ids = append(ids, id)
  1225. }
  1226. if err := rows.Err(); err != nil {
  1227. return nil, err
  1228. }
  1229. if len(ids) == 0 {
  1230. return &service.OAuthRefreshCandidatePage{Accounts: []service.Account{}}, nil
  1231. }
  1232. accounts, err := r.GetByIDs(ctx, ids)
  1233. if err != nil {
  1234. return nil, err
  1235. }
  1236. accountsByID := make(map[int64]*service.Account, len(accounts))
  1237. for _, account := range accounts {
  1238. if account != nil {
  1239. accountsByID[account.ID] = account
  1240. }
  1241. }
  1242. out := make([]service.Account, 0, len(accounts))
  1243. for _, id := range ids {
  1244. if account := accountsByID[id]; account != nil {
  1245. out = append(out, *account)
  1246. }
  1247. }
  1248. page := &service.OAuthRefreshCandidatePage{
  1249. Accounts: out,
  1250. HasMore: len(ids) == options.Limit,
  1251. }
  1252. if len(ids) > 0 {
  1253. page.NextAfterID = ids[len(ids)-1]
  1254. }
  1255. return page, nil
  1256. }
  1257. func (r *accountRepository) ListByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
  1258. accounts, err := r.client.Account.Query().
  1259. Where(
  1260. dbaccount.PlatformEQ(platform),
  1261. dbaccount.StatusEQ(service.StatusActive),
  1262. ).
  1263. Order(dbent.Asc(dbaccount.FieldPriority)).
  1264. All(ctx)
  1265. if err != nil {
  1266. return nil, err
  1267. }
  1268. return r.accountsToService(ctx, accounts)
  1269. }
  1270. func (r *accountRepository) UpdateLastUsed(ctx context.Context, id int64) error {
  1271. now := time.Now()
  1272. _, err := r.client.Account.Update().
  1273. Where(dbaccount.IDEQ(id)).
  1274. SetLastUsedAt(now).
  1275. Save(ctx)
  1276. if err != nil {
  1277. return err
  1278. }
  1279. payload := map[string]any{
  1280. "last_used": map[string]int64{
  1281. strconv.FormatInt(id, 10): now.Unix(),
  1282. },
  1283. }
  1284. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountLastUsed, &id, nil, payload); err != nil {
  1285. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue last used failed: account=%d err=%v", id, err)
  1286. }
  1287. return nil
  1288. }
  1289. func (r *accountRepository) BatchUpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
  1290. if len(updates) == 0 {
  1291. return nil
  1292. }
  1293. ids := make([]int64, 0, len(updates))
  1294. args := make([]any, 0, len(updates)*2+1)
  1295. caseSQL := "UPDATE accounts SET last_used_at = CASE id"
  1296. idx := 1
  1297. for id, ts := range updates {
  1298. caseSQL += " WHEN $" + itoa(idx) + " THEN $" + itoa(idx+1) + "::timestamptz"
  1299. args = append(args, id, ts)
  1300. ids = append(ids, id)
  1301. idx += 2
  1302. }
  1303. caseSQL += " END, updated_at = NOW() WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL"
  1304. args = append(args, pq.Array(ids))
  1305. _, err := r.sql.ExecContext(ctx, caseSQL, args...)
  1306. if err != nil {
  1307. return err
  1308. }
  1309. lastUsedPayload := make(map[string]int64, len(updates))
  1310. for id, ts := range updates {
  1311. lastUsedPayload[strconv.FormatInt(id, 10)] = ts.Unix()
  1312. }
  1313. payload := map[string]any{"last_used": lastUsedPayload}
  1314. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountLastUsed, nil, nil, payload); err != nil {
  1315. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue batch last used failed: err=%v", err)
  1316. }
  1317. return nil
  1318. }
  1319. func (r *accountRepository) SetError(ctx context.Context, id int64, errorMsg string) error {
  1320. _, err := r.client.Account.Update().
  1321. Where(dbaccount.IDEQ(id)).
  1322. SetStatus(service.StatusError).
  1323. SetErrorMessage(errorMsg).
  1324. SetSchedulable(false).
  1325. Save(ctx)
  1326. if err != nil {
  1327. return err
  1328. }
  1329. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  1330. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue set error failed: account=%d err=%v", id, err)
  1331. }
  1332. r.syncSchedulerAccountSnapshot(ctx, id)
  1333. return nil
  1334. }
  1335. func (r *accountRepository) SetGrokCredentialErrorIfMatch(
  1336. ctx context.Context,
  1337. id int64,
  1338. snapshot service.GrokCredentialMutationSnapshot,
  1339. errorMsg string,
  1340. ) (bool, error) {
  1341. result, err := r.sql.ExecContext(ctx, `
  1342. WITH updated AS (
  1343. UPDATE accounts AS a
  1344. SET status = $1,
  1345. error_message = $2,
  1346. schedulable = false,
  1347. updated_at = NOW()
  1348. WHERE a.id = $3
  1349. AND a.deleted_at IS NULL
  1350. AND a.status = $4
  1351. AND a.platform = $5
  1352. AND a.type = $6
  1353. AND a.schedulable IS TRUE
  1354. AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW())
  1355. AND (a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW())
  1356. AND (a.overload_until IS NULL OR a.overload_until <= NOW())
  1357. AND (a.auto_pause_on_expired IS NOT TRUE OR a.expires_at IS NULL OR a.expires_at > NOW())
  1358. AND a.credentials = $7::jsonb
  1359. AND a.proxy_id IS NOT DISTINCT FROM $8
  1360. AND ($2 <> $9 OR (
  1361. a.proxy_id IS NOT NULL AND NOT EXISTS (
  1362. SELECT 1 FROM proxies p WHERE p.id = a.proxy_id AND p.deleted_at IS NULL
  1363. )
  1364. ))
  1365. RETURNING a.id
  1366. )
  1367. INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
  1368. SELECT $10, updated.id, NULL, NULL FROM updated
  1369. `, service.StatusError, errorMsg, id, service.StatusActive, service.PlatformGrok, service.AccountTypeOAuth,
  1370. snapshot.CredentialsJSON, snapshot.ProxyID, string(service.GrokCredentialReasonProxyInvalid),
  1371. service.SchedulerOutboxEventAccountChanged)
  1372. if err != nil {
  1373. return false, err
  1374. }
  1375. affected, err := result.RowsAffected()
  1376. if err != nil || affected == 0 {
  1377. return false, err
  1378. }
  1379. r.syncSchedulerAccountSnapshotDetached(ctx, id)
  1380. return true, nil
  1381. }
  1382. // SetGrokOAuthErrorIfCredentialsUnchanged atomically quarantines a structurally
  1383. // invalid Grok OAuth account only if it is still active and its complete JSONB
  1384. // credential document matches the state observed by reconciliation. Exact
  1385. // JSONB equality includes _token_version when present and prevents a concurrent
  1386. // reauthorization from being overwritten by a stale check-then-mutate path.
  1387. func (r *accountRepository) SetGrokOAuthErrorIfCredentialsUnchanged(
  1388. ctx context.Context,
  1389. id int64,
  1390. expectedCredentials map[string]any,
  1391. errorMsg string,
  1392. ) (bool, error) {
  1393. if r == nil || r.sql == nil {
  1394. return false, errors.New("account repository SQL executor is not configured")
  1395. }
  1396. expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials))
  1397. if err != nil {
  1398. return false, err
  1399. }
  1400. result, err := r.sql.ExecContext(ctx, `
  1401. WITH updated AS (
  1402. UPDATE accounts AS a
  1403. SET status = $1,
  1404. error_message = $2,
  1405. schedulable = FALSE,
  1406. updated_at = NOW()
  1407. WHERE a.id = $3
  1408. AND a.deleted_at IS NULL
  1409. AND a.platform = $4
  1410. AND a.type = $5
  1411. AND a.status = $6
  1412. AND a.credentials = $7::jsonb
  1413. AND NULLIF(BTRIM(a.credentials->>'refresh_token'), '') IS NULL
  1414. RETURNING a.id
  1415. )
  1416. INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
  1417. SELECT $8, updated.id, NULL, NULL FROM updated
  1418. `,
  1419. service.StatusError,
  1420. errorMsg,
  1421. id,
  1422. service.PlatformGrok,
  1423. service.AccountTypeOAuth,
  1424. service.StatusActive,
  1425. string(expectedJSON),
  1426. service.SchedulerOutboxEventAccountChanged,
  1427. )
  1428. if err != nil {
  1429. return false, err
  1430. }
  1431. rowsAffected, err := result.RowsAffected()
  1432. if err != nil {
  1433. return false, err
  1434. }
  1435. if rowsAffected == 0 {
  1436. return false, nil
  1437. }
  1438. r.syncSchedulerAccountSnapshotDetached(ctx, id)
  1439. return true, nil
  1440. }
  1441. // UpdateGrokOAuthCredentialsIfUnchanged persists provider-issued replacement
  1442. // credentials only while the complete Grok OAuth credential document and
  1443. // proxy still match the fresh snapshot used by the upstream refresh call. The
  1444. // scheduler outbox insert is part of the same PostgreSQL statement, so a
  1445. // durable invalidation failure rolls the credential update back as well.
  1446. func (r *accountRepository) UpdateGrokOAuthCredentialsIfUnchanged(
  1447. ctx context.Context,
  1448. id int64,
  1449. expectedCredentials map[string]any,
  1450. expectedProxyID *int64,
  1451. credentials map[string]any,
  1452. ) (bool, error) {
  1453. if r == nil || r.sql == nil {
  1454. return false, errors.New("account repository SQL executor is not configured")
  1455. }
  1456. expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials))
  1457. if err != nil {
  1458. return false, err
  1459. }
  1460. credentialsJSON, err := json.Marshal(normalizeJSONMap(credentials))
  1461. if err != nil {
  1462. return false, err
  1463. }
  1464. result, err := r.sql.ExecContext(ctx, `
  1465. WITH updated AS (
  1466. UPDATE accounts AS a
  1467. SET credentials = $1::jsonb,
  1468. updated_at = NOW()
  1469. WHERE a.id = $2
  1470. AND a.deleted_at IS NULL
  1471. AND a.platform = $3
  1472. AND a.type = $4
  1473. AND a.credentials = $5::jsonb
  1474. AND a.proxy_id IS NOT DISTINCT FROM $6
  1475. RETURNING a.id
  1476. )
  1477. INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
  1478. SELECT $7, updated.id, NULL, NULL FROM updated
  1479. `,
  1480. string(credentialsJSON),
  1481. id,
  1482. service.PlatformGrok,
  1483. service.AccountTypeOAuth,
  1484. string(expectedJSON),
  1485. expectedProxyID,
  1486. service.SchedulerOutboxEventAccountChanged,
  1487. )
  1488. if err != nil {
  1489. return false, err
  1490. }
  1491. rowsAffected, err := result.RowsAffected()
  1492. if err != nil {
  1493. return false, err
  1494. }
  1495. if rowsAffected == 0 {
  1496. return false, nil
  1497. }
  1498. r.syncSchedulerAccountSnapshotDetached(ctx, id)
  1499. return true, nil
  1500. }
  1501. // SetGrokOAuthRefreshErrorIfCredentialsUnchanged is the background-refresh
  1502. // counterpart to reconciliation's stricter missing-refresh-token mutation. It
  1503. // matches the complete credential document used by the failed upstream attempt
  1504. // but deliberately does not require the refresh token to be absent.
  1505. func (r *accountRepository) SetGrokOAuthRefreshErrorIfCredentialsUnchanged(
  1506. ctx context.Context,
  1507. id int64,
  1508. expectedCredentials map[string]any,
  1509. expectedProxyID *int64,
  1510. errorMsg string,
  1511. ) (bool, error) {
  1512. if r == nil || r.sql == nil {
  1513. return false, errors.New("account repository SQL executor is not configured")
  1514. }
  1515. expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials))
  1516. if err != nil {
  1517. return false, err
  1518. }
  1519. result, err := r.sql.ExecContext(ctx, `
  1520. WITH updated AS (
  1521. UPDATE accounts AS a
  1522. SET status = $1,
  1523. error_message = $2,
  1524. schedulable = FALSE,
  1525. updated_at = NOW()
  1526. WHERE a.id = $3
  1527. AND a.deleted_at IS NULL
  1528. AND a.platform = $4
  1529. AND a.type = $5
  1530. AND a.status = $6
  1531. AND a.credentials = $7::jsonb
  1532. AND a.proxy_id IS NOT DISTINCT FROM $8
  1533. RETURNING a.id
  1534. )
  1535. INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
  1536. SELECT $9, updated.id, NULL, NULL FROM updated
  1537. `,
  1538. service.StatusError,
  1539. errorMsg,
  1540. id,
  1541. service.PlatformGrok,
  1542. service.AccountTypeOAuth,
  1543. service.StatusActive,
  1544. string(expectedJSON),
  1545. expectedProxyID,
  1546. service.SchedulerOutboxEventAccountChanged,
  1547. )
  1548. if err != nil {
  1549. return false, err
  1550. }
  1551. rowsAffected, err := result.RowsAffected()
  1552. if err != nil {
  1553. return false, err
  1554. }
  1555. if rowsAffected == 0 {
  1556. return false, nil
  1557. }
  1558. r.syncSchedulerAccountSnapshotDetached(ctx, id)
  1559. return true, nil
  1560. }
  1561. // SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged applies a bounded
  1562. // transient refresh quarantine only while the active Grok OAuth credential
  1563. // document still matches the exact upstream attempt.
  1564. func (r *accountRepository) SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged(
  1565. ctx context.Context,
  1566. id int64,
  1567. expectedCredentials map[string]any,
  1568. expectedProxyID *int64,
  1569. until time.Time,
  1570. reason string,
  1571. ) (bool, error) {
  1572. if r == nil || r.sql == nil {
  1573. return false, errors.New("account repository SQL executor is not configured")
  1574. }
  1575. expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials))
  1576. if err != nil {
  1577. return false, err
  1578. }
  1579. result, err := r.sql.ExecContext(ctx, `
  1580. WITH updated AS (
  1581. UPDATE accounts AS a
  1582. SET temp_unschedulable_until = $1,
  1583. temp_unschedulable_reason = $2,
  1584. updated_at = NOW()
  1585. WHERE a.id = $3
  1586. AND a.deleted_at IS NULL
  1587. AND a.platform = $4
  1588. AND a.type = $5
  1589. AND a.status = $6
  1590. AND a.credentials = $7::jsonb
  1591. AND a.proxy_id IS NOT DISTINCT FROM $8
  1592. AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until < $1)
  1593. RETURNING a.id
  1594. )
  1595. INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
  1596. SELECT $9, updated.id, NULL, NULL FROM updated
  1597. `,
  1598. until,
  1599. reason,
  1600. id,
  1601. service.PlatformGrok,
  1602. service.AccountTypeOAuth,
  1603. service.StatusActive,
  1604. string(expectedJSON),
  1605. expectedProxyID,
  1606. service.SchedulerOutboxEventAccountChanged,
  1607. )
  1608. if err != nil {
  1609. return false, err
  1610. }
  1611. rowsAffected, err := result.RowsAffected()
  1612. if err != nil {
  1613. return false, err
  1614. }
  1615. if rowsAffected == 0 {
  1616. return false, nil
  1617. }
  1618. r.syncSchedulerAccountSnapshotDetached(ctx, id)
  1619. return true, nil
  1620. }
  1621. // syncSchedulerAccountSnapshot 在账号状态变更时主动同步快照到调度器缓存。
  1622. // 当账号被设置为错误、禁用、不可调度或临时不可调度时调用,
  1623. // 确保调度器和粘性会话逻辑能及时感知账号的最新状态,避免继续使用不可用账号。
  1624. //
  1625. // syncSchedulerAccountSnapshot proactively syncs account snapshot to scheduler cache
  1626. // when account status changes. Called when account is set to error, disabled,
  1627. // unschedulable, or temporarily unschedulable, ensuring scheduler and sticky session
  1628. // logic can promptly detect the latest account state and avoid using unavailable accounts.
  1629. func (r *accountRepository) syncSchedulerAccountSnapshot(ctx context.Context, accountID int64) {
  1630. if r == nil || r.schedulerCache == nil || accountID <= 0 {
  1631. return
  1632. }
  1633. account, err := r.GetByID(ctx, accountID)
  1634. if err != nil {
  1635. logger.LegacyPrintf("repository.account", "[Scheduler] sync account snapshot read failed: id=%d err=%v", accountID, err)
  1636. return
  1637. }
  1638. if err := r.schedulerCache.SetAccount(ctx, account); err != nil {
  1639. logger.LegacyPrintf("repository.account", "[Scheduler] sync account snapshot write failed: id=%d err=%v", accountID, err)
  1640. }
  1641. }
  1642. func (r *accountRepository) syncSchedulerAccountSnapshotDetached(ctx context.Context, accountID int64) {
  1643. base := context.Background()
  1644. if ctx != nil {
  1645. base = context.WithoutCancel(ctx)
  1646. }
  1647. propagationCtx, cancel := context.WithTimeout(base, 2*time.Second)
  1648. defer cancel()
  1649. r.syncSchedulerAccountSnapshot(propagationCtx, accountID)
  1650. }
  1651. func (r *accountRepository) deleteSchedulerAccountSnapshot(ctx context.Context, accountID int64) {
  1652. if r == nil || r.schedulerCache == nil || accountID <= 0 {
  1653. return
  1654. }
  1655. if err := r.schedulerCache.DeleteAccount(ctx, accountID); err != nil {
  1656. logger.LegacyPrintf("repository.account", "[Scheduler] delete account snapshot failed: id=%d err=%v", accountID, err)
  1657. }
  1658. }
  1659. func (r *accountRepository) syncSchedulerAccountSnapshots(ctx context.Context, accountIDs []int64) {
  1660. if r == nil || r.schedulerCache == nil || len(accountIDs) == 0 {
  1661. return
  1662. }
  1663. uniqueIDs := make([]int64, 0, len(accountIDs))
  1664. seen := make(map[int64]struct{}, len(accountIDs))
  1665. for _, id := range accountIDs {
  1666. if id <= 0 {
  1667. continue
  1668. }
  1669. if _, exists := seen[id]; exists {
  1670. continue
  1671. }
  1672. seen[id] = struct{}{}
  1673. uniqueIDs = append(uniqueIDs, id)
  1674. }
  1675. if len(uniqueIDs) == 0 {
  1676. return
  1677. }
  1678. accounts, err := r.GetByIDs(ctx, uniqueIDs)
  1679. if err != nil {
  1680. logger.LegacyPrintf("repository.account", "[Scheduler] batch sync account snapshot read failed: count=%d err=%v", len(uniqueIDs), err)
  1681. return
  1682. }
  1683. for _, account := range accounts {
  1684. if account == nil {
  1685. continue
  1686. }
  1687. if err := r.schedulerCache.SetAccount(ctx, account); err != nil {
  1688. logger.LegacyPrintf("repository.account", "[Scheduler] batch sync account snapshot write failed: id=%d err=%v", account.ID, err)
  1689. }
  1690. }
  1691. }
  1692. func (r *accountRepository) ClearError(ctx context.Context, id int64) error {
  1693. _, err := r.client.Account.Update().
  1694. Where(dbaccount.IDEQ(id)).
  1695. SetStatus(service.StatusActive).
  1696. SetErrorMessage("").
  1697. Save(ctx)
  1698. if err != nil {
  1699. return err
  1700. }
  1701. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  1702. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue clear error failed: account=%d err=%v", id, err)
  1703. }
  1704. r.syncSchedulerAccountSnapshot(ctx, id)
  1705. return nil
  1706. }
  1707. func (r *accountRepository) AddToGroup(ctx context.Context, accountID, groupID int64, priority int) error {
  1708. tx, err := r.client.Tx(ctx)
  1709. if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
  1710. return err
  1711. }
  1712. client := r.client
  1713. if tx != nil {
  1714. defer func() { _ = tx.Rollback() }()
  1715. client = tx.Client()
  1716. }
  1717. if err := lockLiveGroups(ctx, client, []int64{groupID}); err != nil {
  1718. return err
  1719. }
  1720. _, err = client.AccountGroup.Create().
  1721. SetAccountID(accountID).
  1722. SetGroupID(groupID).
  1723. SetPriority(priority).
  1724. Save(ctx)
  1725. if err != nil {
  1726. return err
  1727. }
  1728. if tx != nil {
  1729. if err := tx.Commit(); err != nil {
  1730. return err
  1731. }
  1732. }
  1733. payload := buildSchedulerGroupPayload([]int64{groupID})
  1734. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountGroupsChanged, &accountID, nil, payload); err != nil {
  1735. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue add to group failed: account=%d group=%d err=%v", accountID, groupID, err)
  1736. }
  1737. return nil
  1738. }
  1739. func (r *accountRepository) RemoveFromGroup(ctx context.Context, accountID, groupID int64) error {
  1740. _, err := r.client.AccountGroup.Delete().
  1741. Where(
  1742. dbaccountgroup.AccountIDEQ(accountID),
  1743. dbaccountgroup.GroupIDEQ(groupID),
  1744. ).
  1745. Exec(ctx)
  1746. if err != nil {
  1747. return err
  1748. }
  1749. payload := buildSchedulerGroupPayload([]int64{groupID})
  1750. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountGroupsChanged, &accountID, nil, payload); err != nil {
  1751. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue remove from group failed: account=%d group=%d err=%v", accountID, groupID, err)
  1752. }
  1753. return nil
  1754. }
  1755. func (r *accountRepository) GetGroups(ctx context.Context, accountID int64) ([]service.Group, error) {
  1756. groups, err := r.client.Group.Query().
  1757. Where(
  1758. dbgroup.HasAccountsWith(dbaccount.IDEQ(accountID)),
  1759. ).
  1760. All(ctx)
  1761. if err != nil {
  1762. return nil, err
  1763. }
  1764. outGroups := make([]service.Group, 0, len(groups))
  1765. for i := range groups {
  1766. outGroups = append(outGroups, *groupEntityToService(groups[i]))
  1767. }
  1768. return outGroups, nil
  1769. }
  1770. func (r *accountRepository) BindGroups(ctx context.Context, accountID int64, groupIDs []int64) error {
  1771. existingGroupIDs, err := r.loadAccountGroupIDs(ctx, accountID)
  1772. if err != nil {
  1773. return err
  1774. }
  1775. // 使用事务保证删除旧绑定与创建新绑定的原子性
  1776. tx, err := r.client.Tx(ctx)
  1777. if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
  1778. return err
  1779. }
  1780. var txClient *dbent.Client
  1781. if err == nil {
  1782. defer func() { _ = tx.Rollback() }()
  1783. txClient = tx.Client()
  1784. } else {
  1785. // 已处于外部事务中(ErrTxStarted),复用当前 client
  1786. txClient = r.client
  1787. }
  1788. if err := lockLiveGroups(ctx, txClient, groupIDs); err != nil {
  1789. return err
  1790. }
  1791. if _, err := txClient.AccountGroup.Delete().Where(dbaccountgroup.AccountIDEQ(accountID)).Exec(ctx); err != nil {
  1792. return err
  1793. }
  1794. if len(groupIDs) == 0 {
  1795. if tx != nil {
  1796. return tx.Commit()
  1797. }
  1798. return nil
  1799. }
  1800. builders := make([]*dbent.AccountGroupCreate, 0, len(groupIDs))
  1801. for i, groupID := range groupIDs {
  1802. builders = append(builders, txClient.AccountGroup.Create().
  1803. SetAccountID(accountID).
  1804. SetGroupID(groupID).
  1805. SetPriority(i+1),
  1806. )
  1807. }
  1808. if _, err := txClient.AccountGroup.CreateBulk(builders...).Save(ctx); err != nil {
  1809. return err
  1810. }
  1811. if tx != nil {
  1812. if err := tx.Commit(); err != nil {
  1813. return err
  1814. }
  1815. }
  1816. payload := buildSchedulerGroupPayload(mergeGroupIDs(existingGroupIDs, groupIDs))
  1817. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountGroupsChanged, &accountID, nil, payload); err != nil {
  1818. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue bind groups failed: account=%d err=%v", accountID, err)
  1819. }
  1820. return nil
  1821. }
  1822. func (r *accountRepository) ListSchedulable(ctx context.Context) ([]service.Account, error) {
  1823. accounts, err := r.schedulableAccountsQuery(time.Now()).All(ctx)
  1824. if err != nil {
  1825. return nil, err
  1826. }
  1827. return r.accountsToService(ctx, accounts)
  1828. }
  1829. func (r *accountRepository) ListSchedulableAccountLoads(ctx context.Context) ([]service.AccountWithConcurrency, error) {
  1830. accounts, err := r.schedulableAccountsQuery(time.Now()).
  1831. Select(
  1832. dbaccount.FieldID,
  1833. dbaccount.FieldConcurrency,
  1834. dbaccount.FieldLoadFactor,
  1835. ).
  1836. All(ctx)
  1837. if err != nil {
  1838. return nil, err
  1839. }
  1840. loads := make([]service.AccountWithConcurrency, 0, len(accounts))
  1841. for _, account := range accounts {
  1842. projection := service.Account{
  1843. ID: account.ID,
  1844. Concurrency: account.Concurrency,
  1845. LoadFactor: account.LoadFactor,
  1846. }
  1847. loads = append(loads, service.AccountWithConcurrency{
  1848. ID: account.ID,
  1849. MaxConcurrency: projection.EffectiveLoadFactor(),
  1850. })
  1851. }
  1852. return loads, nil
  1853. }
  1854. func (r *accountRepository) schedulableAccountsQuery(now time.Time) *dbent.AccountQuery {
  1855. return r.client.Account.Query().
  1856. Where(
  1857. dbaccount.StatusEQ(service.StatusActive),
  1858. dbaccount.SchedulableEQ(true),
  1859. tempUnschedulablePredicate(),
  1860. notExpiredPredicate(now),
  1861. dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
  1862. dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
  1863. ).
  1864. Order(dbent.Asc(dbaccount.FieldPriority))
  1865. }
  1866. func (r *accountRepository) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]service.Account, error) {
  1867. return r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{
  1868. status: service.StatusActive,
  1869. schedulable: true,
  1870. })
  1871. }
  1872. func (r *accountRepository) ListSchedulableCapacityByGroupIDs(ctx context.Context, groupIDs []int64) ([]service.GroupAccountCapacityRow, error) {
  1873. groupIDs = uniquePositiveInt64s(groupIDs)
  1874. if len(groupIDs) == 0 {
  1875. return []service.GroupAccountCapacityRow{}, nil
  1876. }
  1877. if r.sql == nil {
  1878. rows := make([]service.GroupAccountCapacityRow, 0)
  1879. for _, groupID := range groupIDs {
  1880. accounts, err := r.ListSchedulableByGroupID(ctx, groupID)
  1881. if err != nil {
  1882. return nil, err
  1883. }
  1884. for i := range accounts {
  1885. acc := &accounts[i]
  1886. rows = append(rows, service.GroupAccountCapacityRow{
  1887. GroupID: groupID,
  1888. AccountID: acc.ID,
  1889. Concurrency: acc.Concurrency,
  1890. Extra: copyJSONMap(acc.Extra),
  1891. SessionWindowStart: acc.SessionWindowStart,
  1892. SessionWindowEnd: acc.SessionWindowEnd,
  1893. SessionWindowStatus: acc.SessionWindowStatus,
  1894. })
  1895. }
  1896. }
  1897. return rows, nil
  1898. }
  1899. rows, err := r.sql.QueryContext(ctx, `
  1900. SELECT
  1901. ag.group_id,
  1902. a.id AS account_id,
  1903. a.concurrency,
  1904. COALESCE(a.extra, '{}'::jsonb)::text AS extra,
  1905. a.session_window_start,
  1906. a.session_window_end,
  1907. COALESCE(a.session_window_status, '') AS session_window_status
  1908. FROM account_groups ag
  1909. JOIN accounts a ON a.id = ag.account_id
  1910. WHERE ag.group_id = ANY($1)
  1911. AND a.deleted_at IS NULL
  1912. AND a.status = $2
  1913. AND a.schedulable = TRUE
  1914. AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= $3)
  1915. AND (a.expires_at IS NULL OR a.expires_at > $3 OR a.auto_pause_on_expired = FALSE)
  1916. AND (a.overload_until IS NULL OR a.overload_until <= $3)
  1917. AND (a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= $3)
  1918. ORDER BY ag.group_id ASC, ag.priority ASC, a.priority ASC, a.id ASC
  1919. `, pq.Array(groupIDs), service.StatusActive, time.Now())
  1920. if err != nil {
  1921. return nil, err
  1922. }
  1923. defer func() { _ = rows.Close() }()
  1924. out := make([]service.GroupAccountCapacityRow, 0)
  1925. for rows.Next() {
  1926. var row service.GroupAccountCapacityRow
  1927. var extraRaw string
  1928. if err := rows.Scan(
  1929. &row.GroupID,
  1930. &row.AccountID,
  1931. &row.Concurrency,
  1932. &extraRaw,
  1933. &row.SessionWindowStart,
  1934. &row.SessionWindowEnd,
  1935. &row.SessionWindowStatus,
  1936. ); err != nil {
  1937. return nil, err
  1938. }
  1939. if extraRaw != "" && extraRaw != "null" {
  1940. var extra map[string]any
  1941. if err := json.Unmarshal([]byte(extraRaw), &extra); err != nil {
  1942. return nil, err
  1943. }
  1944. row.Extra = extra
  1945. }
  1946. out = append(out, row)
  1947. }
  1948. if err := rows.Err(); err != nil {
  1949. return nil, err
  1950. }
  1951. return out, nil
  1952. }
  1953. func (r *accountRepository) ListSchedulableByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
  1954. now := time.Now()
  1955. accounts, err := r.client.Account.Query().
  1956. Where(
  1957. dbaccount.PlatformEQ(platform),
  1958. dbaccount.StatusEQ(service.StatusActive),
  1959. dbaccount.SchedulableEQ(true),
  1960. tempUnschedulablePredicate(),
  1961. notExpiredPredicate(now),
  1962. dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
  1963. dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
  1964. ).
  1965. Order(dbent.Asc(dbaccount.FieldPriority)).
  1966. All(ctx)
  1967. if err != nil {
  1968. return nil, err
  1969. }
  1970. return r.accountsToService(ctx, accounts)
  1971. }
  1972. func (r *accountRepository) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]service.Account, error) {
  1973. // 单平台查询复用多平台逻辑,保持过滤条件与排序策略一致。
  1974. return r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{
  1975. status: service.StatusActive,
  1976. schedulable: true,
  1977. platforms: []string{platform},
  1978. })
  1979. }
  1980. func (r *accountRepository) ListSchedulableByPlatforms(ctx context.Context, platforms []string) ([]service.Account, error) {
  1981. if len(platforms) == 0 {
  1982. return nil, nil
  1983. }
  1984. // 仅返回可调度的活跃账号,并过滤处于过载/限流窗口的账号。
  1985. // 代理与分组信息统一在 accountsToService 中批量加载,避免 N+1 查询。
  1986. now := time.Now()
  1987. accounts, err := r.client.Account.Query().
  1988. Where(
  1989. dbaccount.PlatformIn(platforms...),
  1990. dbaccount.StatusEQ(service.StatusActive),
  1991. dbaccount.SchedulableEQ(true),
  1992. tempUnschedulablePredicate(),
  1993. notExpiredPredicate(now),
  1994. dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
  1995. dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
  1996. ).
  1997. Order(dbent.Asc(dbaccount.FieldPriority)).
  1998. All(ctx)
  1999. if err != nil {
  2000. return nil, err
  2001. }
  2002. return r.accountsToService(ctx, accounts)
  2003. }
  2004. func (r *accountRepository) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
  2005. now := time.Now()
  2006. accounts, err := r.client.Account.Query().
  2007. Where(
  2008. dbaccount.PlatformEQ(platform),
  2009. dbaccount.StatusEQ(service.StatusActive),
  2010. dbaccount.SchedulableEQ(true),
  2011. dbaccount.Not(dbaccount.HasAccountGroups()),
  2012. tempUnschedulablePredicate(),
  2013. notExpiredPredicate(now),
  2014. dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
  2015. dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
  2016. ).
  2017. Order(dbent.Asc(dbaccount.FieldPriority)).
  2018. All(ctx)
  2019. if err != nil {
  2020. return nil, err
  2021. }
  2022. return r.accountsToService(ctx, accounts)
  2023. }
  2024. func (r *accountRepository) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]service.Account, error) {
  2025. if len(platforms) == 0 {
  2026. return nil, nil
  2027. }
  2028. now := time.Now()
  2029. accounts, err := r.client.Account.Query().
  2030. Where(
  2031. dbaccount.PlatformIn(platforms...),
  2032. dbaccount.StatusEQ(service.StatusActive),
  2033. dbaccount.SchedulableEQ(true),
  2034. dbaccount.Not(dbaccount.HasAccountGroups()),
  2035. tempUnschedulablePredicate(),
  2036. notExpiredPredicate(now),
  2037. dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
  2038. dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
  2039. ).
  2040. Order(dbent.Asc(dbaccount.FieldPriority)).
  2041. All(ctx)
  2042. if err != nil {
  2043. return nil, err
  2044. }
  2045. return r.accountsToService(ctx, accounts)
  2046. }
  2047. func (r *accountRepository) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]service.Account, error) {
  2048. if len(platforms) == 0 {
  2049. return nil, nil
  2050. }
  2051. // 复用按分组查询逻辑,保证分组优先级 + 账号优先级的排序与筛选一致。
  2052. return r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{
  2053. status: service.StatusActive,
  2054. schedulable: true,
  2055. platforms: platforms,
  2056. })
  2057. }
  2058. // ListModelAvailabilityCandidates returns the persistently configured account
  2059. // pool used to decide whether a model is supported. Unlike scheduling queries,
  2060. // it intentionally ignores transient runtime state (rate limits, overload,
  2061. // temporary unschedulability, and expiry windows).
  2062. func (r *accountRepository) ListModelAvailabilityCandidates(
  2063. ctx context.Context,
  2064. groupID *int64,
  2065. platforms []string,
  2066. includeGrouped bool,
  2067. ) ([]service.Account, error) {
  2068. if len(platforms) == 0 {
  2069. return []service.Account{}, nil
  2070. }
  2071. if groupID != nil {
  2072. return r.queryAccountsByGroup(ctx, *groupID, accountGroupQueryOptions{
  2073. status: service.StatusActive,
  2074. schedulable: true,
  2075. ignoreTransientState: true,
  2076. platforms: platforms,
  2077. })
  2078. }
  2079. preds := []dbpredicate.Account{
  2080. dbaccount.StatusEQ(service.StatusActive),
  2081. dbaccount.SchedulableEQ(true),
  2082. dbaccount.PlatformIn(platforms...),
  2083. }
  2084. if !includeGrouped {
  2085. preds = append(preds, dbaccount.Not(dbaccount.HasAccountGroups()))
  2086. }
  2087. accounts, err := r.client.Account.Query().
  2088. Where(preds...).
  2089. Order(dbent.Asc(dbaccount.FieldPriority)).
  2090. All(ctx)
  2091. if err != nil {
  2092. return nil, err
  2093. }
  2094. return r.accountsToService(ctx, accounts)
  2095. }
  2096. func (r *accountRepository) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
  2097. // Every upstream observation extends a limit. Only an explicit recovery
  2098. // operation may clear it; late short responses must not undo a longer wait.
  2099. return r.SetRateLimitedIfLater(ctx, id, resetAt)
  2100. }
  2101. // SetRateLimitedIfLater atomically extends an account-level rate limit. Grok
  2102. // requests may finish concurrently, so an older response must not overwrite a
  2103. // later reset boundary observed by another request or instance.
  2104. func (r *accountRepository) SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error {
  2105. now := time.Now()
  2106. updated, err := r.client.Account.Update().
  2107. Where(
  2108. dbaccount.IDEQ(id),
  2109. dbaccount.Or(
  2110. dbaccount.RateLimitResetAtIsNil(),
  2111. dbaccount.RateLimitResetAtLT(resetAt),
  2112. ),
  2113. ).
  2114. SetRateLimitedAt(now).
  2115. SetRateLimitResetAt(resetAt).
  2116. Save(ctx)
  2117. if err != nil {
  2118. return err
  2119. }
  2120. if updated == 0 {
  2121. // This instance may not have observed the later value written elsewhere.
  2122. // Refresh its local scheduler snapshot even though no outbox event is needed.
  2123. r.syncSchedulerAccountSnapshot(ctx, id)
  2124. return nil
  2125. }
  2126. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2127. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue extended rate limit failed: account=%d err=%v", id, err)
  2128. }
  2129. r.syncSchedulerAccountSnapshot(ctx, id)
  2130. return nil
  2131. }
  2132. // ClearRateLimitIfObserved clears exactly the Grok rate-limit generation seen
  2133. // by a successful request. Matching both timestamps prevents a stale success
  2134. // from erasing a later clear/re-arm generation with an equal or shorter reset.
  2135. func (r *accountRepository) ClearRateLimitIfObserved(ctx context.Context, id int64, observedLimitedAt, observedResetAt time.Time) (bool, error) {
  2136. updated, err := r.client.Account.Update().
  2137. Where(
  2138. dbaccount.IDEQ(id),
  2139. dbaccount.PlatformEQ(service.PlatformGrok),
  2140. dbaccount.TypeEQ(service.AccountTypeOAuth),
  2141. dbaccount.RateLimitedAtEQ(observedLimitedAt),
  2142. dbaccount.RateLimitResetAtEQ(observedResetAt),
  2143. ).
  2144. ClearRateLimitedAt().
  2145. ClearRateLimitResetAt().
  2146. Save(ctx)
  2147. if err != nil {
  2148. return false, err
  2149. }
  2150. if updated == 0 {
  2151. r.syncSchedulerAccountSnapshot(ctx, id)
  2152. return false, nil
  2153. }
  2154. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2155. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue observed rate-limit clear failed: account=%d err=%v", id, err)
  2156. }
  2157. r.syncSchedulerAccountSnapshot(ctx, id)
  2158. return true, nil
  2159. }
  2160. func (r *accountRepository) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error {
  2161. if scope == "" {
  2162. return nil
  2163. }
  2164. now := time.Now().UTC()
  2165. payload := map[string]string{
  2166. "rate_limited_at": now.Format(time.RFC3339),
  2167. "rate_limit_reset_at": resetAt.UTC().Format(time.RFC3339),
  2168. }
  2169. if len(reason) > 0 {
  2170. if value := strings.TrimSpace(reason[0]); value != "" {
  2171. payload["reason"] = value
  2172. }
  2173. }
  2174. raw, err := json.Marshal(payload)
  2175. if err != nil {
  2176. return err
  2177. }
  2178. client := clientFromContext(ctx, r.client)
  2179. result, err := client.ExecContext(
  2180. ctx,
  2181. `UPDATE accounts SET
  2182. extra = jsonb_set(
  2183. jsonb_set(COALESCE(extra, '{}'::jsonb), '{model_rate_limits}'::text[], COALESCE(extra->'model_rate_limits', '{}'::jsonb), true),
  2184. ARRAY['model_rate_limits', $1]::text[],
  2185. $2::jsonb,
  2186. true
  2187. ),
  2188. updated_at = NOW()
  2189. WHERE id = $3 AND deleted_at IS NULL`,
  2190. scope,
  2191. raw,
  2192. id,
  2193. )
  2194. if err != nil {
  2195. return err
  2196. }
  2197. affected, err := result.RowsAffected()
  2198. if err != nil {
  2199. return err
  2200. }
  2201. if affected == 0 {
  2202. return service.ErrAccountNotFound
  2203. }
  2204. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2205. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue model rate limit failed: account=%d err=%v", id, err)
  2206. }
  2207. r.syncSchedulerAccountSnapshot(ctx, id)
  2208. return nil
  2209. }
  2210. func (r *accountRepository) SetOverloaded(ctx context.Context, id int64, until time.Time) error {
  2211. _, err := r.client.Account.Update().
  2212. Where(dbaccount.IDEQ(id)).
  2213. SetOverloadUntil(until).
  2214. Save(ctx)
  2215. if err != nil {
  2216. return err
  2217. }
  2218. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2219. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue overload failed: account=%d err=%v", id, err)
  2220. }
  2221. r.syncSchedulerAccountSnapshot(ctx, id)
  2222. return nil
  2223. }
  2224. func (r *accountRepository) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
  2225. result, err := r.sql.ExecContext(ctx, `
  2226. UPDATE accounts
  2227. SET temp_unschedulable_until = $1,
  2228. temp_unschedulable_reason = $2,
  2229. updated_at = NOW()
  2230. WHERE id = $3
  2231. AND deleted_at IS NULL
  2232. AND (temp_unschedulable_until IS NULL OR temp_unschedulable_until < $1)
  2233. `, until, reason, id)
  2234. if err != nil {
  2235. return err
  2236. }
  2237. affected, err := result.RowsAffected()
  2238. if err != nil {
  2239. return err
  2240. }
  2241. if affected <= 0 {
  2242. return nil
  2243. }
  2244. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2245. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue temp unschedulable failed: account=%d err=%v", id, err)
  2246. }
  2247. r.syncSchedulerAccountSnapshot(ctx, id)
  2248. return nil
  2249. }
  2250. func (r *accountRepository) SetGrokCredentialTempUnschedulableIfMatch(
  2251. ctx context.Context,
  2252. id int64,
  2253. snapshot service.GrokCredentialMutationSnapshot,
  2254. until time.Time,
  2255. reason string,
  2256. ) (bool, error) {
  2257. result, err := r.sql.ExecContext(ctx, `
  2258. WITH updated AS (
  2259. UPDATE accounts AS a
  2260. SET temp_unschedulable_until = CASE
  2261. WHEN a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until < $1 THEN $1
  2262. ELSE a.temp_unschedulable_until
  2263. END,
  2264. temp_unschedulable_reason = $2,
  2265. updated_at = NOW()
  2266. WHERE a.id = $3
  2267. AND a.deleted_at IS NULL
  2268. AND a.status = $4
  2269. AND a.platform = $5
  2270. AND a.type = $6
  2271. AND a.schedulable IS TRUE
  2272. AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW())
  2273. AND (a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW())
  2274. AND (a.overload_until IS NULL OR a.overload_until <= NOW())
  2275. AND (a.auto_pause_on_expired IS NOT TRUE OR a.expires_at IS NULL OR a.expires_at > NOW())
  2276. AND a.credentials = $7::jsonb
  2277. AND a.proxy_id IS NOT DISTINCT FROM $8
  2278. RETURNING a.id
  2279. )
  2280. INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
  2281. SELECT $9, updated.id, NULL, NULL FROM updated
  2282. `, until, reason, id, service.StatusActive, service.PlatformGrok, service.AccountTypeOAuth,
  2283. snapshot.CredentialsJSON, snapshot.ProxyID, service.SchedulerOutboxEventAccountChanged)
  2284. if err != nil {
  2285. return false, err
  2286. }
  2287. affected, err := result.RowsAffected()
  2288. if err != nil || affected == 0 {
  2289. return false, err
  2290. }
  2291. r.syncSchedulerAccountSnapshotDetached(ctx, id)
  2292. return true, nil
  2293. }
  2294. func (r *accountRepository) ClearTempUnschedulable(ctx context.Context, id int64) error {
  2295. _, err := r.sql.ExecContext(ctx, `
  2296. UPDATE accounts
  2297. SET temp_unschedulable_until = NULL,
  2298. temp_unschedulable_reason = NULL,
  2299. updated_at = NOW()
  2300. WHERE id = $1
  2301. AND deleted_at IS NULL
  2302. `, id)
  2303. if err != nil {
  2304. return err
  2305. }
  2306. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2307. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue clear temp unschedulable failed: account=%d err=%v", id, err)
  2308. }
  2309. r.syncSchedulerAccountSnapshot(ctx, id)
  2310. return nil
  2311. }
  2312. func (r *accountRepository) ClearRateLimit(ctx context.Context, id int64) error {
  2313. _, err := r.client.Account.Update().
  2314. Where(dbaccount.IDEQ(id)).
  2315. ClearRateLimitedAt().
  2316. ClearRateLimitResetAt().
  2317. ClearOverloadUntil().
  2318. Save(ctx)
  2319. if err != nil {
  2320. return err
  2321. }
  2322. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2323. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue clear rate limit failed: account=%d err=%v", id, err)
  2324. }
  2325. r.syncSchedulerAccountSnapshot(ctx, id)
  2326. return nil
  2327. }
  2328. func (r *accountRepository) ClearAntigravityQuotaScopes(ctx context.Context, id int64) error {
  2329. client := clientFromContext(ctx, r.client)
  2330. result, err := client.ExecContext(
  2331. ctx,
  2332. "UPDATE accounts SET extra = COALESCE(extra, '{}'::jsonb) - 'antigravity_quota_scopes', updated_at = NOW() WHERE id = $1 AND deleted_at IS NULL",
  2333. id,
  2334. )
  2335. if err != nil {
  2336. return err
  2337. }
  2338. affected, err := result.RowsAffected()
  2339. if err != nil {
  2340. return err
  2341. }
  2342. if affected == 0 {
  2343. return service.ErrAccountNotFound
  2344. }
  2345. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2346. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue clear quota scopes failed: account=%d err=%v", id, err)
  2347. }
  2348. return nil
  2349. }
  2350. func (r *accountRepository) ClearModelRateLimits(ctx context.Context, id int64) error {
  2351. client := clientFromContext(ctx, r.client)
  2352. result, err := client.ExecContext(
  2353. ctx,
  2354. "UPDATE accounts SET extra = COALESCE(extra, '{}'::jsonb) - 'model_rate_limits', updated_at = NOW() WHERE id = $1 AND deleted_at IS NULL",
  2355. id,
  2356. )
  2357. if err != nil {
  2358. return err
  2359. }
  2360. affected, err := result.RowsAffected()
  2361. if err != nil {
  2362. return err
  2363. }
  2364. if affected == 0 {
  2365. return service.ErrAccountNotFound
  2366. }
  2367. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2368. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue clear model rate limit failed: account=%d err=%v", id, err)
  2369. }
  2370. r.syncSchedulerAccountSnapshot(ctx, id)
  2371. return nil
  2372. }
  2373. func (r *accountRepository) UpdateSessionWindow(ctx context.Context, id int64, start, end *time.Time, status string) error {
  2374. builder := r.client.Account.Update().
  2375. Where(dbaccount.IDEQ(id)).
  2376. SetSessionWindowStatus(status)
  2377. if start != nil {
  2378. builder.SetSessionWindowStart(*start)
  2379. }
  2380. if end != nil {
  2381. builder.SetSessionWindowEnd(*end)
  2382. }
  2383. _, err := builder.Save(ctx)
  2384. if err != nil {
  2385. return err
  2386. }
  2387. // 触发调度器缓存更新(仅当窗口时间有变化时)
  2388. if start != nil || end != nil {
  2389. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2390. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue session window update failed: account=%d err=%v", id, err)
  2391. }
  2392. }
  2393. return nil
  2394. }
  2395. func (r *accountRepository) UpdateSessionWindowEnd(ctx context.Context, id int64, end time.Time) error {
  2396. _, err := r.client.Account.Update().
  2397. Where(dbaccount.IDEQ(id)).
  2398. SetSessionWindowEnd(end).
  2399. Save(ctx)
  2400. if err != nil {
  2401. return err
  2402. }
  2403. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2404. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue session window end update failed: account=%d err=%v", id, err)
  2405. }
  2406. return nil
  2407. }
  2408. func (r *accountRepository) SetSchedulable(ctx context.Context, id int64, schedulable bool) error {
  2409. _, err := r.client.Account.Update().
  2410. Where(dbaccount.IDEQ(id)).
  2411. SetSchedulable(schedulable).
  2412. Save(ctx)
  2413. if err != nil {
  2414. return err
  2415. }
  2416. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2417. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue schedulable change failed: account=%d err=%v", id, err)
  2418. }
  2419. if !schedulable {
  2420. r.syncSchedulerAccountSnapshot(ctx, id)
  2421. }
  2422. return nil
  2423. }
  2424. func (r *accountRepository) AutoPauseExpiredAccounts(ctx context.Context, now time.Time) (int64, error) {
  2425. rows, err := r.sql.QueryContext(ctx, `
  2426. UPDATE accounts
  2427. SET schedulable = FALSE,
  2428. updated_at = NOW()
  2429. WHERE deleted_at IS NULL
  2430. AND schedulable = TRUE
  2431. AND auto_pause_on_expired = TRUE
  2432. AND expires_at IS NOT NULL
  2433. AND expires_at <= $1
  2434. RETURNING id
  2435. `, now)
  2436. if err != nil {
  2437. return 0, err
  2438. }
  2439. defer func() {
  2440. _ = rows.Close()
  2441. }()
  2442. accountIDs := make([]int64, 0)
  2443. for rows.Next() {
  2444. var accountID int64
  2445. if err := rows.Scan(&accountID); err != nil {
  2446. return 0, err
  2447. }
  2448. accountIDs = append(accountIDs, accountID)
  2449. }
  2450. if err := rows.Err(); err != nil {
  2451. return 0, err
  2452. }
  2453. if len(accountIDs) > 0 {
  2454. // 只刷新本次暂停的账号及其所属分组,避免少量账号到期触发所有调度桶重建。
  2455. payload := map[string]any{"account_ids": accountIDs}
  2456. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil {
  2457. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue auto pause account changes failed: err=%v", err)
  2458. }
  2459. }
  2460. return int64(len(accountIDs)), nil
  2461. }
  2462. func (r *accountRepository) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
  2463. updates = stripCodexFingerprintSeedFromExtraUpdate(updates)
  2464. if len(updates) == 0 {
  2465. return nil
  2466. }
  2467. // 使用 JSONB 合并操作实现原子更新,避免读-改-写的并发丢失更新问题
  2468. payload, err := json.Marshal(updates)
  2469. if err != nil {
  2470. return err
  2471. }
  2472. clearProbeSnapshot := upstreamBillingProbeExplicitlyDisabled(updates) || upstreamBillingProbeSnapshotClearRequested(updates)
  2473. durableSchedulerChange := shouldEnqueueSchedulerOutboxForExtraUpdates(updates) || clearProbeSnapshot
  2474. baseCtx := ctx
  2475. contextTx := dbent.TxFromContext(ctx)
  2476. client := clientFromContext(ctx, r.client)
  2477. var tx *dbent.Tx
  2478. if durableSchedulerChange && contextTx == nil {
  2479. var txErr error
  2480. tx, txErr = r.client.Tx(ctx)
  2481. if txErr != nil && !errors.Is(txErr, dbent.ErrTxStarted) {
  2482. return txErr
  2483. }
  2484. if tx != nil {
  2485. defer func() { _ = tx.Rollback() }()
  2486. ctx = dbent.NewTxContext(ctx, tx)
  2487. client = tx.Client()
  2488. }
  2489. }
  2490. extraExpression := "COALESCE(extra, '{}'::jsonb) || $1::jsonb"
  2491. if clearProbeSnapshot {
  2492. extraExpression = "(" + extraExpression + ") - 'upstream_billing_probe'"
  2493. }
  2494. if service.ShouldEnsureCodexFingerprintSeedForExtraUpdates(updates) {
  2495. extraExpression = ensureCodexFingerprintSeedSQL(extraExpression)
  2496. }
  2497. extraExpression = preserveProtectionExtraSQL(ctx, extraExpression)
  2498. result, err := client.ExecContext(
  2499. ctx,
  2500. "UPDATE accounts SET extra = "+extraExpression+", updated_at = NOW() WHERE id = $2 AND deleted_at IS NULL",
  2501. string(payload), id,
  2502. )
  2503. if err != nil {
  2504. return err
  2505. }
  2506. affected, err := result.RowsAffected()
  2507. if err != nil {
  2508. return err
  2509. }
  2510. if affected == 0 {
  2511. return service.ErrAccountNotFound
  2512. }
  2513. if durableSchedulerChange {
  2514. if err := enqueueSchedulerOutbox(ctx, client, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  2515. return err
  2516. }
  2517. if tx != nil {
  2518. if err := tx.Commit(); err != nil {
  2519. return err
  2520. }
  2521. }
  2522. if contextTx == nil {
  2523. r.syncSchedulerAccountSnapshot(baseCtx, id)
  2524. }
  2525. } else {
  2526. // 观测型 extra 字段不需要触发 bucket 重建,但仍同步单账号快照,
  2527. // 让 sticky session / GetAccount 命中缓存时也能读到最新数据,
  2528. // 同时避免缓存局部 patch 覆盖掉并发写入的其它账号字段。
  2529. if dbent.TxFromContext(ctx) == nil {
  2530. r.syncSchedulerAccountSnapshot(ctx, id)
  2531. }
  2532. }
  2533. return nil
  2534. }
  2535. // UpdateUpstreamBillingProbeSnapshot stores a probe result only while the
  2536. // network identity used by that probe is still current.
  2537. func (r *accountRepository) UpdateUpstreamBillingProbeSnapshot(
  2538. ctx context.Context,
  2539. account *service.Account,
  2540. snapshot *service.UpstreamBillingProbeSnapshot,
  2541. rateMultiplier *float64,
  2542. ) error {
  2543. if account == nil || snapshot == nil {
  2544. return service.ErrAccountNilInput
  2545. }
  2546. if snapshot.Status != service.UpstreamBillingProbeStatusOK {
  2547. rateMultiplier = nil
  2548. }
  2549. if dbent.TxFromContext(ctx) == nil {
  2550. tx, err := r.client.Tx(ctx)
  2551. if errors.Is(err, dbent.ErrTxStarted) {
  2552. return r.updateUpstreamBillingProbeSnapshotInTx(ctx, account, snapshot, rateMultiplier)
  2553. }
  2554. if err != nil {
  2555. return err
  2556. }
  2557. defer func() { _ = tx.Rollback() }()
  2558. if err := r.updateUpstreamBillingProbeSnapshotInTx(dbent.NewTxContext(ctx, tx), account, snapshot, rateMultiplier); err != nil {
  2559. return err
  2560. }
  2561. if err := tx.Commit(); err != nil {
  2562. return err
  2563. }
  2564. // The durable outbox event is committed with the snapshot. This direct
  2565. // cache write only reduces visibility latency on the current instance.
  2566. r.syncSchedulerAccountSnapshot(ctx, account.ID)
  2567. return nil
  2568. }
  2569. return r.updateUpstreamBillingProbeSnapshotInTx(ctx, account, snapshot, rateMultiplier)
  2570. }
  2571. func (r *accountRepository) updateUpstreamBillingProbeSnapshotInTx(
  2572. ctx context.Context,
  2573. account *service.Account,
  2574. snapshot *service.UpstreamBillingProbeSnapshot,
  2575. rateMultiplier *float64,
  2576. ) error {
  2577. payload, err := json.Marshal(map[string]any{service.UpstreamBillingProbeExtraKey: snapshot})
  2578. if err != nil {
  2579. return err
  2580. }
  2581. credentials, err := json.Marshal(account.Credentials)
  2582. if err != nil {
  2583. return err
  2584. }
  2585. var expectedSnapshot any
  2586. if account.Extra != nil {
  2587. expectedSnapshot = account.Extra[service.UpstreamBillingProbeExtraKey]
  2588. }
  2589. expectedSnapshotJSON, err := json.Marshal(expectedSnapshot)
  2590. if err != nil {
  2591. return err
  2592. }
  2593. var expectedEnabled any
  2594. if account.Extra != nil {
  2595. expectedEnabled = account.Extra[service.UpstreamBillingProbeEnabledExtraKey]
  2596. }
  2597. expectedEnabledJSON, err := json.Marshal(expectedEnabled)
  2598. if err != nil {
  2599. return err
  2600. }
  2601. var expectedRateSyncEnabled any
  2602. if account.Extra != nil {
  2603. expectedRateSyncEnabled = account.Extra[service.UpstreamBillingRateSyncEnabledExtraKey]
  2604. }
  2605. expectedRateSyncEnabledJSON, err := json.Marshal(expectedRateSyncEnabled)
  2606. if err != nil {
  2607. return err
  2608. }
  2609. client := clientFromContext(ctx, r.client)
  2610. proxyMatches, err := lockAndMatchProbeProxyIdentity(ctx, client, account)
  2611. if err != nil {
  2612. return err
  2613. }
  2614. if !proxyMatches {
  2615. return service.ErrUpstreamBillingProbeIdentityChanged
  2616. }
  2617. var proxyID any
  2618. if account.ProxyID != nil {
  2619. proxyID = *account.ProxyID
  2620. }
  2621. result, err := client.ExecContext(ctx, `
  2622. UPDATE accounts
  2623. SET
  2624. extra = COALESCE(extra, '{}'::jsonb) || $1::jsonb,
  2625. rate_multiplier = CASE
  2626. WHEN $10::numeric IS NOT NULL
  2627. AND extra @> '{"upstream_billing_probe_enabled": true}'::jsonb
  2628. AND extra @> '{"upstream_billing_rate_sync_enabled": true}'::jsonb
  2629. THEN $10::numeric
  2630. ELSE rate_multiplier
  2631. END,
  2632. updated_at = NOW()
  2633. WHERE id = $2
  2634. AND platform = $3
  2635. AND type = $4
  2636. AND credentials = $5::jsonb
  2637. AND proxy_id IS NOT DISTINCT FROM $6
  2638. AND COALESCE(extra -> 'upstream_billing_probe', 'null'::jsonb) = $7::jsonb
  2639. AND COALESCE(extra -> 'upstream_billing_probe_enabled', 'null'::jsonb) = $8::jsonb
  2640. AND COALESCE(extra -> 'upstream_billing_rate_sync_enabled', 'null'::jsonb) = $9::jsonb
  2641. AND deleted_at IS NULL
  2642. `, string(payload), account.ID, account.Platform, account.Type, string(credentials), proxyID, string(expectedSnapshotJSON), string(expectedEnabledJSON), string(expectedRateSyncEnabledJSON), rateMultiplier)
  2643. if err != nil {
  2644. return err
  2645. }
  2646. affected, err := result.RowsAffected()
  2647. if err != nil {
  2648. return err
  2649. }
  2650. if affected == 0 {
  2651. return service.ErrUpstreamBillingProbeIdentityChanged
  2652. }
  2653. return enqueueSchedulerOutbox(ctx, client, service.SchedulerOutboxEventAccountChanged, &account.ID, nil, nil)
  2654. }
  2655. func lockAndMatchProbeProxyIdentity(ctx context.Context, client *dbent.Client, account *service.Account) (bool, error) {
  2656. if account.ProxyID == nil {
  2657. return true, nil
  2658. }
  2659. rows, err := client.QueryContext(ctx, `
  2660. SELECT protocol, host, port, COALESCE(username, ''), COALESCE(password, ''), status
  2661. FROM proxies
  2662. WHERE id = $1 AND deleted_at IS NULL
  2663. FOR SHARE
  2664. `, *account.ProxyID)
  2665. if err != nil {
  2666. return false, err
  2667. }
  2668. defer func() { _ = rows.Close() }()
  2669. if !rows.Next() {
  2670. if err := rows.Err(); err != nil {
  2671. return false, err
  2672. }
  2673. return account.Proxy == nil, nil
  2674. }
  2675. if account.Proxy == nil || account.Proxy.ID != *account.ProxyID {
  2676. return false, nil
  2677. }
  2678. var current proxyProbeIdentity
  2679. if err := rows.Scan(&current.protocol, &current.host, &current.port, &current.username, &current.password, &current.status); err != nil {
  2680. return false, err
  2681. }
  2682. return current == proxyProbeIdentityFromService(account.Proxy), rows.Err()
  2683. }
  2684. func shouldEnqueueSchedulerOutboxForExtraUpdates(updates map[string]any) bool {
  2685. if len(updates) == 0 {
  2686. return false
  2687. }
  2688. for key := range updates {
  2689. if isSchedulerNeutralExtraKey(key) {
  2690. continue
  2691. }
  2692. return true
  2693. }
  2694. return false
  2695. }
  2696. func isSchedulerNeutralExtraKey(key string) bool {
  2697. key = strings.TrimSpace(key)
  2698. if key == "" {
  2699. return false
  2700. }
  2701. if _, ok := schedulerNeutralExtraKeys[key]; ok {
  2702. return true
  2703. }
  2704. for _, prefix := range schedulerNeutralExtraKeyPrefixes {
  2705. if strings.HasPrefix(key, prefix) {
  2706. return true
  2707. }
  2708. }
  2709. return false
  2710. }
  2711. func upstreamBillingProbeExplicitlyDisabled(extra map[string]any) bool {
  2712. enabled, ok := extra[service.UpstreamBillingProbeEnabledExtraKey].(bool)
  2713. return ok && !enabled
  2714. }
  2715. func upstreamBillingProbeSnapshotClearRequested(extra map[string]any) bool {
  2716. value, ok := extra[service.UpstreamBillingProbeExtraKey]
  2717. return ok && value == nil
  2718. }
  2719. func ollamaCloudUsageSnapshotClearRequested(extra map[string]any) bool {
  2720. value, ok := extra[service.OllamaCloudUsageSnapshotExtraKey]
  2721. return ok && value == nil
  2722. }
  2723. func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates service.AccountBulkUpdate) (int64, error) {
  2724. if len(ids) == 0 {
  2725. return 0, nil
  2726. }
  2727. updates.Extra = stripCodexFingerprintSeedFromExtraUpdate(updates.Extra)
  2728. setClauses := make([]string, 0, 8)
  2729. args := make([]any, 0, 8)
  2730. idx := 1
  2731. ollamaProxyIdentityChanged := ""
  2732. concurrencyExpression := ""
  2733. if updates.Name != nil {
  2734. setClauses = append(setClauses, "name = $"+itoa(idx))
  2735. args = append(args, *updates.Name)
  2736. idx++
  2737. }
  2738. if updates.ProxyID != nil {
  2739. // 0 表示清除代理(前端发送 0 而不是 null 来表达清除意图)
  2740. if *updates.ProxyID == 0 {
  2741. setClauses = append(setClauses, "proxy_id = NULL")
  2742. ollamaProxyIdentityChanged = "proxy_id IS NOT NULL"
  2743. } else {
  2744. proxyPlaceholder := "$" + itoa(idx)
  2745. setClauses = append(setClauses, "proxy_id = "+proxyPlaceholder)
  2746. ollamaProxyIdentityChanged = "proxy_id IS DISTINCT FROM " + proxyPlaceholder
  2747. args = append(args, *updates.ProxyID)
  2748. idx++
  2749. }
  2750. }
  2751. if updates.Concurrency != nil {
  2752. concurrencyExpression = protectedConcurrencySQL("$" + itoa(idx) + "::integer")
  2753. setClauses = append(setClauses, "concurrency = "+concurrencyExpression)
  2754. args = append(args, *updates.Concurrency)
  2755. idx++
  2756. }
  2757. if updates.Priority != nil {
  2758. setClauses = append(setClauses, "priority = $"+itoa(idx))
  2759. args = append(args, *updates.Priority)
  2760. idx++
  2761. }
  2762. if updates.RateMultiplier != nil {
  2763. setClauses = append(setClauses, "rate_multiplier = $"+itoa(idx))
  2764. args = append(args, *updates.RateMultiplier)
  2765. idx++
  2766. }
  2767. if updates.LoadFactor != nil {
  2768. if *updates.LoadFactor <= 0 {
  2769. setClauses = append(setClauses, "load_factor = NULL")
  2770. } else {
  2771. setClauses = append(setClauses, "load_factor = $"+itoa(idx))
  2772. args = append(args, *updates.LoadFactor)
  2773. idx++
  2774. }
  2775. }
  2776. if updates.Status != nil {
  2777. setClauses = append(setClauses, "status = $"+itoa(idx))
  2778. args = append(args, *updates.Status)
  2779. idx++
  2780. }
  2781. if updates.Schedulable != nil {
  2782. setClauses = append(setClauses, "schedulable = $"+itoa(idx))
  2783. args = append(args, *updates.Schedulable)
  2784. idx++
  2785. }
  2786. if updates.ProbeEnabled != nil {
  2787. if updates.Extra == nil {
  2788. updates.Extra = make(map[string]any)
  2789. }
  2790. updates.Extra[service.UpstreamBillingProbeEnabledExtraKey] = *updates.ProbeEnabled
  2791. }
  2792. // JSONB 需要合并而非覆盖,使用 raw SQL 保持旧行为。
  2793. credentialPlaceholder := ""
  2794. if len(updates.Credentials) > 0 {
  2795. payload, err := json.Marshal(updates.Credentials)
  2796. if err != nil {
  2797. return 0, err
  2798. }
  2799. credentialPlaceholder = "$" + itoa(idx)
  2800. setClauses = append(setClauses, "credentials = COALESCE(credentials, '{}'::jsonb) || "+credentialPlaceholder+"::jsonb")
  2801. args = append(args, payload)
  2802. idx++
  2803. }
  2804. ollamaGroupIdentityChanges := make([]string, 0, 2)
  2805. if _, ok := updates.Credentials["api_key"]; ok {
  2806. ollamaGroupIdentityChanges = append(ollamaGroupIdentityChanges, "credentials -> 'api_key' IS DISTINCT FROM "+credentialPlaceholder+"::jsonb -> 'api_key'")
  2807. }
  2808. if _, ok := updates.Credentials["base_url"]; ok {
  2809. ollamaGroupIdentityChanges = append(ollamaGroupIdentityChanges,
  2810. "NOT ("+ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'")+
  2811. " AND "+ollamaCloudBaseURLMatchesSQL(credentialPlaceholder+"::jsonb ->> 'base_url'")+")")
  2812. }
  2813. if len(updates.Extra) > 0 || len(ollamaGroupIdentityChanges) > 0 || ollamaProxyIdentityChanged != "" || updates.EnsureCodexFingerprintSeed || updates.Concurrency != nil {
  2814. extraExpression := "COALESCE(extra, '{}'::jsonb)"
  2815. if len(updates.Extra) > 0 {
  2816. payload, err := json.Marshal(updates.Extra)
  2817. if err != nil {
  2818. return 0, err
  2819. }
  2820. extraExpression += " || $" + itoa(idx) + "::jsonb"
  2821. args = append(args, payload)
  2822. idx++
  2823. if upstreamBillingProbeExplicitlyDisabled(updates.Extra) || upstreamBillingProbeSnapshotClearRequested(updates.Extra) {
  2824. extraExpression = "(" + extraExpression + ") - 'upstream_billing_probe'"
  2825. }
  2826. if ollamaCloudUsageSnapshotClearRequested(updates.Extra) {
  2827. extraExpression = "(" + extraExpression + ") - 'ollama_cloud_usage_snapshot'"
  2828. }
  2829. }
  2830. eligibleAccount := "platform IN (" + ollamaCloudUsagePlatformsSQL + ") AND type = 'apikey'"
  2831. groupIdentityChanged := ""
  2832. if len(ollamaGroupIdentityChanges) > 0 {
  2833. groupIdentityChanged = "(" + eligibleAccount + " AND (" + joinClauses(ollamaGroupIdentityChanges, " OR ") + "))"
  2834. }
  2835. snapshotIdentityChanged := groupIdentityChanged
  2836. if ollamaProxyIdentityChanged != "" {
  2837. proxyChanged := "(" + eligibleAccount + " AND " + ollamaProxyIdentityChanged + ")"
  2838. if snapshotIdentityChanged == "" {
  2839. snapshotIdentityChanged = proxyChanged
  2840. } else {
  2841. snapshotIdentityChanged = "(" + snapshotIdentityChanged + " OR " + proxyChanged + ")"
  2842. }
  2843. }
  2844. if groupIdentityChanged != "" {
  2845. extraExpression = "CASE" +
  2846. " WHEN " + groupIdentityChanged + " THEN (" + extraExpression + ") - 'ollama_cloud_usage_session' - 'ollama_cloud_usage_auto_refresh' - 'ollama_cloud_usage_snapshot'" +
  2847. " WHEN " + snapshotIdentityChanged + " THEN (" + extraExpression + ") - 'ollama_cloud_usage_snapshot'" +
  2848. " ELSE " + extraExpression + " END"
  2849. } else if snapshotIdentityChanged != "" {
  2850. extraExpression = "CASE WHEN " + snapshotIdentityChanged + " THEN (" + extraExpression + ") - 'ollama_cloud_usage_snapshot' ELSE " + extraExpression + " END"
  2851. }
  2852. if updates.EnsureCodexFingerprintSeed {
  2853. extraExpression = ensureCodexFingerprintSeedSQL(extraExpression)
  2854. }
  2855. extraExpression = preserveProtectionExtraSQL(ctx, extraExpression)
  2856. if concurrencyExpression != "" {
  2857. extraExpression = "CASE WHEN " + accountProtectionEnabledSQL + " THEN jsonb_set((" + extraExpression + "), '{anti_degrade,max_concurrency}', to_jsonb((" + concurrencyExpression + ")::integer), true) ELSE (" + extraExpression + ") END"
  2858. }
  2859. setClauses = append(setClauses, "extra = "+extraExpression)
  2860. }
  2861. if len(setClauses) == 0 {
  2862. return 0, nil
  2863. }
  2864. setClauses = append(setClauses, "updated_at = NOW()")
  2865. whereClause := " WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL"
  2866. args = append(args, pq.Array(ids))
  2867. idx++
  2868. if updates.ProbeEnabled != nil {
  2869. whereClause += " AND type = $" + itoa(idx)
  2870. args = append(args, service.AccountTypeAPIKey)
  2871. }
  2872. query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + whereClause
  2873. baseCtx := ctx
  2874. contextTx := dbent.TxFromContext(ctx)
  2875. exec := r.sql
  2876. var tx *dbent.Tx
  2877. if contextTx != nil {
  2878. exec = contextTx.Client()
  2879. } else if r.client != nil {
  2880. var txErr error
  2881. tx, txErr = r.client.Tx(ctx)
  2882. if txErr != nil && !errors.Is(txErr, dbent.ErrTxStarted) {
  2883. return 0, txErr
  2884. }
  2885. if tx != nil {
  2886. defer func() { _ = tx.Rollback() }()
  2887. ctx = dbent.NewTxContext(ctx, tx)
  2888. exec = tx.Client()
  2889. }
  2890. }
  2891. if mode, _ := updates.Extra[service.ProxyModeExtraKey].(string); !service.ProtectionManagedWrite(ctx) && strings.EqualFold(strings.TrimSpace(mode), service.ProxyModeRandom) {
  2892. if dbent.TxFromContext(ctx) == nil {
  2893. return 0, errors.New("random proxy bulk changes require a transaction")
  2894. }
  2895. if err := validateLockedBulkProxyMode(ctx, exec, ids, updates.Extra); err != nil {
  2896. return 0, err
  2897. }
  2898. }
  2899. result, err := exec.ExecContext(ctx, query, args...)
  2900. if err != nil {
  2901. return 0, err
  2902. }
  2903. rows, err := result.RowsAffected()
  2904. if err != nil {
  2905. return 0, err
  2906. }
  2907. if updates.ProbeEnabled != nil {
  2908. expectedRows := int64(0)
  2909. seenIDs := make(map[int64]struct{}, len(ids))
  2910. for _, id := range ids {
  2911. if _, seen := seenIDs[id]; seen {
  2912. continue
  2913. }
  2914. seenIDs[id] = struct{}{}
  2915. expectedRows++
  2916. }
  2917. if rows != expectedRows {
  2918. return 0, service.ErrUpstreamBillingProbeAccountInvalid
  2919. }
  2920. }
  2921. if rows > 0 {
  2922. payload := map[string]any{"account_ids": ids}
  2923. if err := enqueueSchedulerOutbox(ctx, exec, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil {
  2924. return 0, err
  2925. }
  2926. }
  2927. if tx != nil {
  2928. if err := tx.Commit(); err != nil {
  2929. return 0, err
  2930. }
  2931. }
  2932. if rows > 0 && contextTx == nil {
  2933. shouldSync := false
  2934. if updates.Status != nil && (*updates.Status == service.StatusError || *updates.Status == service.StatusDisabled) {
  2935. shouldSync = true
  2936. }
  2937. if updates.Schedulable != nil && !*updates.Schedulable {
  2938. shouldSync = true
  2939. }
  2940. if shouldSync {
  2941. r.syncSchedulerAccountSnapshots(baseCtx, ids)
  2942. }
  2943. }
  2944. return rows, nil
  2945. }
  2946. type accountGroupQueryOptions struct {
  2947. status string
  2948. schedulable bool
  2949. ignoreTransientState bool
  2950. platforms []string // 允许的多个平台,空切片表示不进行平台过滤
  2951. }
  2952. func (r *accountRepository) queryAccountsByGroup(ctx context.Context, groupID int64, opts accountGroupQueryOptions) ([]service.Account, error) {
  2953. q := r.client.AccountGroup.Query().
  2954. Where(dbaccountgroup.GroupIDEQ(groupID))
  2955. // 通过 account_groups 中间表查询账号,并按需叠加状态/平台/调度能力过滤。
  2956. preds := make([]dbpredicate.Account, 0, 6)
  2957. preds = append(preds, dbaccount.DeletedAtIsNil())
  2958. if opts.status != "" {
  2959. preds = append(preds, dbaccount.StatusEQ(opts.status))
  2960. }
  2961. if len(opts.platforms) > 0 {
  2962. preds = append(preds, dbaccount.PlatformIn(opts.platforms...))
  2963. }
  2964. if opts.schedulable {
  2965. preds = append(preds, dbaccount.SchedulableEQ(true))
  2966. if !opts.ignoreTransientState {
  2967. now := time.Now()
  2968. preds = append(preds,
  2969. tempUnschedulablePredicate(),
  2970. notExpiredPredicate(now),
  2971. dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
  2972. dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
  2973. )
  2974. }
  2975. }
  2976. if len(preds) > 0 {
  2977. q = q.Where(dbaccountgroup.HasAccountWith(preds...))
  2978. }
  2979. groups, err := q.
  2980. Order(
  2981. dbaccountgroup.ByPriority(),
  2982. dbaccountgroup.ByAccountField(dbaccount.FieldPriority),
  2983. ).
  2984. WithAccount().
  2985. All(ctx)
  2986. if err != nil {
  2987. return nil, err
  2988. }
  2989. orderedIDs := make([]int64, 0, len(groups))
  2990. accountMap := make(map[int64]*dbent.Account, len(groups))
  2991. for _, ag := range groups {
  2992. if ag.Edges.Account == nil {
  2993. continue
  2994. }
  2995. if _, exists := accountMap[ag.AccountID]; exists {
  2996. continue
  2997. }
  2998. accountMap[ag.AccountID] = ag.Edges.Account
  2999. orderedIDs = append(orderedIDs, ag.AccountID)
  3000. }
  3001. accounts := make([]*dbent.Account, 0, len(orderedIDs))
  3002. for _, id := range orderedIDs {
  3003. if acc, ok := accountMap[id]; ok {
  3004. accounts = append(accounts, acc)
  3005. }
  3006. }
  3007. return r.accountsToService(ctx, accounts)
  3008. }
  3009. func (r *accountRepository) accountsToService(ctx context.Context, accounts []*dbent.Account) ([]service.Account, error) {
  3010. if len(accounts) == 0 {
  3011. return []service.Account{}, nil
  3012. }
  3013. accountIDs := make([]int64, 0, len(accounts))
  3014. proxyIDs := make([]int64, 0, len(accounts))
  3015. for _, acc := range accounts {
  3016. accountIDs = append(accountIDs, acc.ID)
  3017. if acc.ProxyID != nil {
  3018. proxyIDs = append(proxyIDs, *acc.ProxyID)
  3019. }
  3020. if acc.ProxyFallbackOriginID != nil {
  3021. proxyIDs = append(proxyIDs, *acc.ProxyFallbackOriginID)
  3022. }
  3023. }
  3024. proxyMap, err := r.loadProxies(ctx, proxyIDs)
  3025. if err != nil {
  3026. return nil, err
  3027. }
  3028. groupsByAccount, groupIDsByAccount, accountGroupsByAccount, err := r.loadAccountGroups(ctx, accountIDs)
  3029. if err != nil {
  3030. return nil, err
  3031. }
  3032. outAccounts := make([]service.Account, 0, len(accounts))
  3033. for _, acc := range accounts {
  3034. out := accountEntityToService(acc)
  3035. if out == nil {
  3036. continue
  3037. }
  3038. if acc.ProxyID != nil {
  3039. if proxy, ok := proxyMap[*acc.ProxyID]; ok {
  3040. out.Proxy = proxy
  3041. }
  3042. }
  3043. out.ProxyFallbackOriginID = acc.ProxyFallbackOriginID
  3044. if acc.ProxyFallbackOriginID != nil {
  3045. if op, ok := proxyMap[*acc.ProxyFallbackOriginID]; ok && op != nil {
  3046. n := op.Name
  3047. out.ProxyFallbackOriginName = &n
  3048. }
  3049. }
  3050. if groups, ok := groupsByAccount[acc.ID]; ok {
  3051. out.Groups = groups
  3052. }
  3053. if groupIDs, ok := groupIDsByAccount[acc.ID]; ok {
  3054. out.GroupIDs = groupIDs
  3055. }
  3056. if ags, ok := accountGroupsByAccount[acc.ID]; ok {
  3057. out.AccountGroups = ags
  3058. }
  3059. outAccounts = append(outAccounts, *out)
  3060. }
  3061. return outAccounts, nil
  3062. }
  3063. func tempUnschedulablePredicate() dbpredicate.Account {
  3064. return dbpredicate.Account(func(s *entsql.Selector) {
  3065. col := s.C("temp_unschedulable_until")
  3066. s.Where(entsql.Or(
  3067. entsql.IsNull(col),
  3068. entsql.LTE(col, entsql.Expr("NOW()")),
  3069. ))
  3070. })
  3071. }
  3072. func notExpiredPredicate(now time.Time) dbpredicate.Account {
  3073. return dbaccount.Or(
  3074. dbaccount.ExpiresAtIsNil(),
  3075. dbaccount.ExpiresAtGT(now),
  3076. dbaccount.AutoPauseOnExpiredEQ(false),
  3077. )
  3078. }
  3079. func (r *accountRepository) loadProxies(ctx context.Context, proxyIDs []int64) (map[int64]*service.Proxy, error) {
  3080. proxyMap := make(map[int64]*service.Proxy)
  3081. proxyIDs = uniquePositiveInt64s(proxyIDs)
  3082. if len(proxyIDs) == 0 {
  3083. return proxyMap, nil
  3084. }
  3085. for start := 0; start < len(proxyIDs); start += postgresParameterBatchSize {
  3086. end := start + postgresParameterBatchSize
  3087. if end > len(proxyIDs) {
  3088. end = len(proxyIDs)
  3089. }
  3090. proxies, err := r.client.Proxy.Query().Where(dbproxy.IDIn(proxyIDs[start:end]...)).All(ctx)
  3091. if err != nil {
  3092. return nil, err
  3093. }
  3094. for _, p := range proxies {
  3095. proxyMap[p.ID] = proxyEntityToService(p)
  3096. }
  3097. }
  3098. return proxyMap, nil
  3099. }
  3100. func (r *accountRepository) loadAccountGroups(ctx context.Context, accountIDs []int64) (map[int64][]*service.Group, map[int64][]int64, map[int64][]service.AccountGroup, error) {
  3101. groupsByAccount := make(map[int64][]*service.Group)
  3102. groupIDsByAccount := make(map[int64][]int64)
  3103. accountGroupsByAccount := make(map[int64][]service.AccountGroup)
  3104. accountIDs = uniquePositiveInt64s(accountIDs)
  3105. if len(accountIDs) == 0 {
  3106. return groupsByAccount, groupIDsByAccount, accountGroupsByAccount, nil
  3107. }
  3108. for start := 0; start < len(accountIDs); start += postgresParameterBatchSize {
  3109. end := start + postgresParameterBatchSize
  3110. if end > len(accountIDs) {
  3111. end = len(accountIDs)
  3112. }
  3113. entries, err := r.client.AccountGroup.Query().
  3114. Where(dbaccountgroup.AccountIDIn(accountIDs[start:end]...)).
  3115. Order(dbaccountgroup.ByAccountID(), dbaccountgroup.ByPriority()).
  3116. All(ctx)
  3117. if err != nil {
  3118. return nil, nil, nil, err
  3119. }
  3120. groupIDs := make([]int64, 0, len(entries))
  3121. for _, ag := range entries {
  3122. groupIDs = append(groupIDs, ag.GroupID)
  3123. }
  3124. groupMap, err := r.loadGroups(ctx, groupIDs)
  3125. if err != nil {
  3126. return nil, nil, nil, err
  3127. }
  3128. for _, ag := range entries {
  3129. groupSvc := groupMap[ag.GroupID]
  3130. agSvc := service.AccountGroup{
  3131. AccountID: ag.AccountID,
  3132. GroupID: ag.GroupID,
  3133. Priority: ag.Priority,
  3134. CreatedAt: ag.CreatedAt,
  3135. Group: groupSvc,
  3136. }
  3137. accountGroupsByAccount[ag.AccountID] = append(accountGroupsByAccount[ag.AccountID], agSvc)
  3138. groupIDsByAccount[ag.AccountID] = append(groupIDsByAccount[ag.AccountID], ag.GroupID)
  3139. if groupSvc != nil {
  3140. groupsByAccount[ag.AccountID] = append(groupsByAccount[ag.AccountID], groupSvc)
  3141. }
  3142. }
  3143. }
  3144. return groupsByAccount, groupIDsByAccount, accountGroupsByAccount, nil
  3145. }
  3146. func (r *accountRepository) loadGroups(ctx context.Context, groupIDs []int64) (map[int64]*service.Group, error) {
  3147. groupMap := make(map[int64]*service.Group)
  3148. groupIDs = uniquePositiveInt64s(groupIDs)
  3149. if len(groupIDs) == 0 {
  3150. return groupMap, nil
  3151. }
  3152. for start := 0; start < len(groupIDs); start += postgresParameterBatchSize {
  3153. end := start + postgresParameterBatchSize
  3154. if end > len(groupIDs) {
  3155. end = len(groupIDs)
  3156. }
  3157. groups, err := r.client.Group.Query().Where(dbgroup.IDIn(groupIDs[start:end]...)).All(ctx)
  3158. if err != nil {
  3159. return nil, err
  3160. }
  3161. for _, g := range groups {
  3162. groupMap[g.ID] = groupEntityToService(g)
  3163. }
  3164. }
  3165. return groupMap, nil
  3166. }
  3167. func uniquePositiveInt64s(ids []int64) []int64 {
  3168. if len(ids) == 0 {
  3169. return nil
  3170. }
  3171. out := make([]int64, 0, len(ids))
  3172. seen := make(map[int64]struct{}, len(ids))
  3173. for _, id := range ids {
  3174. if id <= 0 {
  3175. continue
  3176. }
  3177. if _, ok := seen[id]; ok {
  3178. continue
  3179. }
  3180. seen[id] = struct{}{}
  3181. out = append(out, id)
  3182. }
  3183. return out
  3184. }
  3185. func (r *accountRepository) loadAccountGroupIDs(ctx context.Context, accountID int64) ([]int64, error) {
  3186. entries, err := r.client.AccountGroup.
  3187. Query().
  3188. Where(dbaccountgroup.AccountIDEQ(accountID)).
  3189. All(ctx)
  3190. if err != nil {
  3191. return nil, err
  3192. }
  3193. ids := make([]int64, 0, len(entries))
  3194. for _, entry := range entries {
  3195. ids = append(ids, entry.GroupID)
  3196. }
  3197. return ids, nil
  3198. }
  3199. func mergeGroupIDs(a []int64, b []int64) []int64 {
  3200. seen := make(map[int64]struct{}, len(a)+len(b))
  3201. out := make([]int64, 0, len(a)+len(b))
  3202. for _, id := range a {
  3203. if id <= 0 {
  3204. continue
  3205. }
  3206. if _, ok := seen[id]; ok {
  3207. continue
  3208. }
  3209. seen[id] = struct{}{}
  3210. out = append(out, id)
  3211. }
  3212. for _, id := range b {
  3213. if id <= 0 {
  3214. continue
  3215. }
  3216. if _, ok := seen[id]; ok {
  3217. continue
  3218. }
  3219. seen[id] = struct{}{}
  3220. out = append(out, id)
  3221. }
  3222. return out
  3223. }
  3224. // buildSchedulerGroupPayload 构造 EventAccountChanged / EventAccountGroupsChanged
  3225. // 事件的 payload。空 groupIDs 必须返回 untyped nil(any 而非 map[string]any(nil)),
  3226. // 否则 enqueueSchedulerOutbox 的 "payload != nil" 接口判空会被 typed-nil 欺骗,
  3227. // 把 payload marshal 成 "null" 写入 dedup_key 哈希,破坏与其他 nil-payload 调用的去重一致性。
  3228. func buildSchedulerGroupPayload(groupIDs []int64) any {
  3229. if len(groupIDs) == 0 {
  3230. return nil
  3231. }
  3232. return map[string]any{"group_ids": groupIDs}
  3233. }
  3234. func accountEntityToService(m *dbent.Account) *service.Account {
  3235. if m == nil {
  3236. return nil
  3237. }
  3238. rateMultiplier := m.RateMultiplier
  3239. return &service.Account{
  3240. ID: m.ID,
  3241. Name: m.Name,
  3242. Notes: m.Notes,
  3243. Platform: m.Platform,
  3244. Type: m.Type,
  3245. Credentials: copyJSONMap(m.Credentials),
  3246. Extra: copyJSONMap(m.Extra),
  3247. ProxyID: m.ProxyID,
  3248. ProxyFallbackOriginID: m.ProxyFallbackOriginID,
  3249. Concurrency: m.Concurrency,
  3250. Priority: m.Priority,
  3251. RateMultiplier: &rateMultiplier,
  3252. LoadFactor: m.LoadFactor,
  3253. Status: m.Status,
  3254. ErrorMessage: derefString(m.ErrorMessage),
  3255. LastUsedAt: m.LastUsedAt,
  3256. ExpiresAt: m.ExpiresAt,
  3257. AutoPauseOnExpired: m.AutoPauseOnExpired,
  3258. CreatedAt: m.CreatedAt,
  3259. UpdatedAt: m.UpdatedAt,
  3260. Schedulable: m.Schedulable,
  3261. RateLimitedAt: m.RateLimitedAt,
  3262. RateLimitResetAt: m.RateLimitResetAt,
  3263. OverloadUntil: m.OverloadUntil,
  3264. TempUnschedulableUntil: m.TempUnschedulableUntil,
  3265. TempUnschedulableReason: derefString(m.TempUnschedulableReason),
  3266. SessionWindowStart: m.SessionWindowStart,
  3267. SessionWindowEnd: m.SessionWindowEnd,
  3268. SessionWindowStatus: derefString(m.SessionWindowStatus),
  3269. ParentAccountID: m.ParentAccountID,
  3270. QuotaDimension: string(m.QuotaDimension),
  3271. }
  3272. }
  3273. func normalizeJSONMap(in map[string]any) map[string]any {
  3274. if in == nil {
  3275. return map[string]any{}
  3276. }
  3277. return in
  3278. }
  3279. func copyJSONMap(in map[string]any) map[string]any {
  3280. if in == nil {
  3281. return nil
  3282. }
  3283. out := make(map[string]any, len(in))
  3284. for k, v := range in {
  3285. out[k] = v
  3286. }
  3287. return out
  3288. }
  3289. func joinClauses(clauses []string, sep string) string {
  3290. if len(clauses) == 0 {
  3291. return ""
  3292. }
  3293. out := clauses[0]
  3294. for i := 1; i < len(clauses); i++ {
  3295. out += sep + clauses[i]
  3296. }
  3297. return out
  3298. }
  3299. func itoa(v int) string {
  3300. return strconv.Itoa(v)
  3301. }
  3302. // FindByExtraField 根据 extra 字段中的键值对查找账号。
  3303. // 使用 PostgreSQL JSONB @> 操作符进行高效查询(需要 GIN 索引支持)。
  3304. //
  3305. // FindByExtraField finds accounts by key-value pairs in the extra field.
  3306. // Uses PostgreSQL JSONB @> operator for efficient queries (requires GIN index).
  3307. func (r *accountRepository) FindByExtraField(ctx context.Context, key string, value any) ([]service.Account, error) {
  3308. accounts, err := r.client.Account.Query().
  3309. Where(
  3310. dbaccount.DeletedAtIsNil(),
  3311. func(s *entsql.Selector) {
  3312. path := sqljson.Path(key)
  3313. switch v := value.(type) {
  3314. case string:
  3315. preds := []*entsql.Predicate{sqljson.ValueEQ(dbaccount.FieldExtra, v, path)}
  3316. if parsed, err := strconv.ParseInt(v, 10, 64); err == nil {
  3317. preds = append(preds, sqljson.ValueEQ(dbaccount.FieldExtra, parsed, path))
  3318. }
  3319. if len(preds) == 1 {
  3320. s.Where(preds[0])
  3321. } else {
  3322. s.Where(entsql.Or(preds...))
  3323. }
  3324. case int:
  3325. s.Where(entsql.Or(
  3326. sqljson.ValueEQ(dbaccount.FieldExtra, v, path),
  3327. sqljson.ValueEQ(dbaccount.FieldExtra, strconv.Itoa(v), path),
  3328. ))
  3329. case int64:
  3330. s.Where(entsql.Or(
  3331. sqljson.ValueEQ(dbaccount.FieldExtra, v, path),
  3332. sqljson.ValueEQ(dbaccount.FieldExtra, strconv.FormatInt(v, 10), path),
  3333. ))
  3334. case json.Number:
  3335. if parsed, err := v.Int64(); err == nil {
  3336. s.Where(entsql.Or(
  3337. sqljson.ValueEQ(dbaccount.FieldExtra, parsed, path),
  3338. sqljson.ValueEQ(dbaccount.FieldExtra, v.String(), path),
  3339. ))
  3340. } else {
  3341. s.Where(sqljson.ValueEQ(dbaccount.FieldExtra, v.String(), path))
  3342. }
  3343. default:
  3344. s.Where(sqljson.ValueEQ(dbaccount.FieldExtra, value, path))
  3345. }
  3346. },
  3347. ).
  3348. All(ctx)
  3349. if err != nil {
  3350. return nil, translatePersistenceError(err, service.ErrAccountNotFound, nil)
  3351. }
  3352. return r.accountsToService(ctx, accounts)
  3353. }
  3354. // ListDueUpstreamBillingProbeAccounts bounds result hydration and network work
  3355. // to limit. PostgreSQL must still filter and order all enabled candidates;
  3356. // MATERIALIZED avoids repeating the defensive timestamp parse expression.
  3357. // Go writes next_probe_at via RFC3339Nano (up to 9 fractional digits) while
  3358. // jsonpath datetime() parses at most microseconds, so fractions beyond 6
  3359. // digits are trimmed first — mirroring ListDueOllamaCloudUsageAccounts.
  3360. // Without this, every nanosecond timestamp is treated as malformed and the
  3361. // fail-open ordering pins the cycle to the lowest account IDs, starving the
  3362. // rest of the pool.
  3363. func (r *accountRepository) ListDueUpstreamBillingProbeAccounts(ctx context.Context, now time.Time, limit int) ([]service.Account, error) {
  3364. if limit <= 0 {
  3365. return []service.Account{}, nil
  3366. }
  3367. if r.sql == nil {
  3368. return nil, errors.New("account repository SQL executor not configured")
  3369. }
  3370. rows, err := r.sql.QueryContext(ctx, `
  3371. WITH candidates AS (
  3372. SELECT
  3373. id,
  3374. extra #>> '{upstream_billing_probe,status}' AS probe_status,
  3375. extra #>> '{upstream_billing_probe,next_probe_at}' AS next_probe_at
  3376. FROM accounts
  3377. WHERE deleted_at IS NULL
  3378. AND status = 'active'
  3379. AND type = 'apikey'
  3380. AND extra @> '{"upstream_billing_probe_enabled": true}'::jsonb
  3381. ), parsed AS MATERIALIZED (
  3382. SELECT
  3383. id,
  3384. probe_status,
  3385. next_probe_at,
  3386. next_probe_at ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}(\.[0-9]+)?(Z|[+-][0-9]{2}:[0-9]{2})$' AS rfc3339_shape,
  3387. jsonb_path_query_first_tz(
  3388. jsonb_build_object(
  3389. 'value',
  3390. replace(regexp_replace(regexp_replace(
  3391. next_probe_at,
  3392. '(\.[0-9]{6})[0-9]+(Z|[+-][0-9]{2}:[0-9]{2})$',
  3393. '\1\2'
  3394. ), 'Z$', '+00:00'), 'T', ' ')
  3395. ),
  3396. '$.value.datetime()',
  3397. '{}'::jsonb,
  3398. true
  3399. ) #>> '{}' AS parsed_next_probe_at
  3400. FROM candidates
  3401. ), normalized AS (
  3402. SELECT
  3403. id,
  3404. probe_status,
  3405. next_probe_at,
  3406. parsed_next_probe_at,
  3407. rfc3339_shape AND parsed_next_probe_at IS NOT NULL AS valid_next_probe_at
  3408. FROM parsed
  3409. )
  3410. SELECT id
  3411. FROM normalized
  3412. WHERE probe_status NOT IN ('ok', 'unsupported', 'failed')
  3413. OR probe_status IS NULL
  3414. OR next_probe_at IS NULL
  3415. OR NOT valid_next_probe_at
  3416. OR CASE WHEN valid_next_probe_at THEN parsed_next_probe_at::timestamptz <= $1 ELSE FALSE END
  3417. ORDER BY
  3418. CASE
  3419. WHEN probe_status NOT IN ('ok', 'unsupported', 'failed')
  3420. OR probe_status IS NULL
  3421. OR next_probe_at IS NULL
  3422. OR NOT valid_next_probe_at
  3423. THEN 0
  3424. ELSE 1
  3425. END ASC,
  3426. CASE WHEN valid_next_probe_at THEN parsed_next_probe_at::timestamptz END ASC NULLS FIRST,
  3427. id ASC
  3428. LIMIT $2
  3429. `, now.UTC(), limit)
  3430. if err != nil {
  3431. return nil, err
  3432. }
  3433. defer func() { _ = rows.Close() }()
  3434. ids := make([]int64, 0, limit)
  3435. for rows.Next() {
  3436. var id int64
  3437. if err := rows.Scan(&id); err != nil {
  3438. return nil, err
  3439. }
  3440. ids = append(ids, id)
  3441. }
  3442. if err := rows.Err(); err != nil {
  3443. return nil, err
  3444. }
  3445. if len(ids) == 0 {
  3446. return []service.Account{}, nil
  3447. }
  3448. accounts, err := r.GetByIDs(ctx, ids)
  3449. if err != nil {
  3450. return nil, err
  3451. }
  3452. out := make([]service.Account, 0, len(accounts))
  3453. for _, account := range accounts {
  3454. if account != nil {
  3455. out = append(out, *account)
  3456. }
  3457. }
  3458. return out, nil
  3459. }
  3460. // nowUTC is a SQL expression to generate a UTC RFC3339 timestamp string.
  3461. const nowUTC = `to_char(NOW() AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.US"Z"')`
  3462. // dailyExpiredExpr is a SQL expression that evaluates to TRUE when daily quota period has expired.
  3463. // Supports both rolling (24h from start) and fixed (pre-computed reset_at) modes.
  3464. const dailyExpiredExpr = `(
  3465. CASE WHEN COALESCE(extra->>'quota_daily_reset_mode', 'rolling') = 'fixed'
  3466. THEN NOW() >= COALESCE((extra->>'quota_daily_reset_at')::timestamptz, '1970-01-01'::timestamptz)
  3467. ELSE COALESCE((extra->>'quota_daily_start')::timestamptz, '1970-01-01'::timestamptz)
  3468. + '24 hours'::interval <= NOW()
  3469. END
  3470. )`
  3471. // weeklyExpiredExpr is a SQL expression that evaluates to TRUE when weekly quota period has expired.
  3472. const weeklyExpiredExpr = `(
  3473. CASE WHEN COALESCE(extra->>'quota_weekly_reset_mode', 'rolling') = 'fixed'
  3474. THEN NOW() >= COALESCE((extra->>'quota_weekly_reset_at')::timestamptz, '1970-01-01'::timestamptz)
  3475. ELSE COALESCE((extra->>'quota_weekly_start')::timestamptz, '1970-01-01'::timestamptz)
  3476. + '168 hours'::interval <= NOW()
  3477. END
  3478. )`
  3479. // nextDailyResetAtExpr is a SQL expression to compute the next daily reset_at when a reset occurs.
  3480. // For fixed mode: computes the next future reset time based on NOW(), timezone, and configured hour.
  3481. // This correctly handles long-inactive accounts by jumping directly to the next valid reset point.
  3482. const nextDailyResetAtExpr = `(
  3483. CASE WHEN COALESCE(extra->>'quota_daily_reset_mode', 'rolling') = 'fixed'
  3484. THEN to_char((
  3485. -- Compute today's reset point in the configured timezone, then pick next future one
  3486. CASE WHEN NOW() >= (
  3487. date_trunc('day', NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))
  3488. + (COALESCE((extra->>'quota_daily_reset_hour')::int, 0) || ' hours')::interval
  3489. ) AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC')
  3490. -- NOW() is at or past today's reset point → next reset is tomorrow
  3491. THEN (
  3492. date_trunc('day', NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))
  3493. + (COALESCE((extra->>'quota_daily_reset_hour')::int, 0) || ' hours')::interval
  3494. + '1 day'::interval
  3495. ) AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC')
  3496. -- NOW() is before today's reset point → next reset is today
  3497. ELSE (
  3498. date_trunc('day', NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))
  3499. + (COALESCE((extra->>'quota_daily_reset_hour')::int, 0) || ' hours')::interval
  3500. ) AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC')
  3501. END
  3502. ) AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS"Z"')
  3503. ELSE NULL END
  3504. )`
  3505. // nextWeeklyResetAtExpr is a SQL expression to compute the next weekly reset_at when a reset occurs.
  3506. // For fixed mode: computes the next future reset time based on NOW(), timezone, configured day and hour.
  3507. // This correctly handles long-inactive accounts by jumping directly to the next valid reset point.
  3508. const nextWeeklyResetAtExpr = `(
  3509. CASE WHEN COALESCE(extra->>'quota_weekly_reset_mode', 'rolling') = 'fixed'
  3510. THEN to_char((
  3511. -- Compute this week's reset point in the configured timezone
  3512. -- Step 1: get today's date at reset hour in configured tz
  3513. -- Step 2: compute days forward to target weekday
  3514. -- Step 3: if same day but past reset hour, advance 7 days
  3515. CASE
  3516. WHEN (
  3517. -- days_forward = (target_day - current_day + 7) % 7
  3518. (COALESCE((extra->>'quota_weekly_reset_day')::int, 1)
  3519. - EXTRACT(DOW FROM NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))::int
  3520. + 7) % 7
  3521. ) = 0 AND NOW() >= (
  3522. date_trunc('day', NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))
  3523. + (COALESCE((extra->>'quota_weekly_reset_hour')::int, 0) || ' hours')::interval
  3524. ) AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC')
  3525. -- Same weekday and past reset hour → next week
  3526. THEN (
  3527. date_trunc('day', NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))
  3528. + (COALESCE((extra->>'quota_weekly_reset_hour')::int, 0) || ' hours')::interval
  3529. + '7 days'::interval
  3530. ) AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC')
  3531. ELSE (
  3532. -- Advance to target weekday this week (or next if days_forward > 0)
  3533. date_trunc('day', NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))
  3534. + (COALESCE((extra->>'quota_weekly_reset_hour')::int, 0) || ' hours')::interval
  3535. + ((
  3536. (COALESCE((extra->>'quota_weekly_reset_day')::int, 1)
  3537. - EXTRACT(DOW FROM NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))::int
  3538. + 7) % 7
  3539. ) || ' days')::interval
  3540. ) AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC')
  3541. END
  3542. ) AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS"Z"')
  3543. ELSE NULL END
  3544. )`
  3545. // IncrementQuotaUsed 原子递增账号的配额用量(总/日/周三个维度)
  3546. // 日/周额度在周期过期时自动重置为 0 再递增。
  3547. // 支持滚动窗口(rolling)和固定时间(fixed)两种重置模式。
  3548. func (r *accountRepository) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) error {
  3549. rows, err := r.sql.QueryContext(ctx,
  3550. `UPDATE accounts SET extra = (
  3551. COALESCE(extra, '{}'::jsonb)
  3552. -- 总额度:始终递增
  3553. || jsonb_build_object('quota_used', COALESCE((extra->>'quota_used')::numeric, 0) + $1)
  3554. -- 日额度:仅在 quota_daily_limit > 0 时处理
  3555. || CASE WHEN COALESCE((extra->>'quota_daily_limit')::numeric, 0) > 0 THEN
  3556. jsonb_build_object(
  3557. 'quota_daily_used',
  3558. CASE WHEN `+dailyExpiredExpr+`
  3559. THEN $1
  3560. ELSE COALESCE((extra->>'quota_daily_used')::numeric, 0) + $1 END,
  3561. 'quota_daily_start',
  3562. CASE WHEN `+dailyExpiredExpr+`
  3563. THEN `+nowUTC+`
  3564. ELSE COALESCE(extra->>'quota_daily_start', `+nowUTC+`) END
  3565. )
  3566. -- 固定模式重置时更新下次重置时间
  3567. || CASE WHEN `+dailyExpiredExpr+` AND `+nextDailyResetAtExpr+` IS NOT NULL
  3568. THEN jsonb_build_object('quota_daily_reset_at', `+nextDailyResetAtExpr+`)
  3569. ELSE '{}'::jsonb END
  3570. ELSE '{}'::jsonb END
  3571. -- 周额度:仅在 quota_weekly_limit > 0 时处理
  3572. || CASE WHEN COALESCE((extra->>'quota_weekly_limit')::numeric, 0) > 0 THEN
  3573. jsonb_build_object(
  3574. 'quota_weekly_used',
  3575. CASE WHEN `+weeklyExpiredExpr+`
  3576. THEN $1
  3577. ELSE COALESCE((extra->>'quota_weekly_used')::numeric, 0) + $1 END,
  3578. 'quota_weekly_start',
  3579. CASE WHEN `+weeklyExpiredExpr+`
  3580. THEN `+nowUTC+`
  3581. ELSE COALESCE(extra->>'quota_weekly_start', `+nowUTC+`) END
  3582. )
  3583. -- 固定模式重置时更新下次重置时间
  3584. || CASE WHEN `+weeklyExpiredExpr+` AND `+nextWeeklyResetAtExpr+` IS NOT NULL
  3585. THEN jsonb_build_object('quota_weekly_reset_at', `+nextWeeklyResetAtExpr+`)
  3586. ELSE '{}'::jsonb END
  3587. ELSE '{}'::jsonb END
  3588. ), updated_at = NOW()
  3589. WHERE id = $2 AND deleted_at IS NULL
  3590. RETURNING
  3591. COALESCE((extra->>'quota_used')::numeric, 0),
  3592. COALESCE((extra->>'quota_limit')::numeric, 0)`,
  3593. amount, id)
  3594. if err != nil {
  3595. return err
  3596. }
  3597. defer func() { _ = rows.Close() }()
  3598. var newUsed, limit float64
  3599. if rows.Next() {
  3600. if err := rows.Scan(&newUsed, &limit); err != nil {
  3601. return err
  3602. }
  3603. }
  3604. if err := rows.Err(); err != nil {
  3605. return err
  3606. }
  3607. // 任一维度配额刚超限时触发调度快照刷新
  3608. if limit > 0 && newUsed >= limit && (newUsed-amount) < limit {
  3609. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  3610. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue quota exceeded failed: account=%d err=%v", id, err)
  3611. }
  3612. }
  3613. return nil
  3614. }
  3615. // ResetQuotaUsedAndClearRateLimitCooldown resets all quota dimensions and the
  3616. // account-level cooldown in one statement. Other scheduler blocking state is preserved.
  3617. func (r *accountRepository) ResetQuotaUsedAndClearRateLimitCooldown(ctx context.Context, id int64) error {
  3618. result, err := r.sql.ExecContext(ctx,
  3619. `UPDATE accounts SET extra = (
  3620. COALESCE(extra, '{}'::jsonb)
  3621. || '{"quota_used": 0, "quota_daily_used": 0, "quota_weekly_used": 0}'::jsonb
  3622. ) - 'quota_daily_start' - 'quota_weekly_start' - 'quota_daily_reset_at' - 'quota_weekly_reset_at',
  3623. rate_limited_at = NULL, rate_limit_reset_at = NULL, updated_at = NOW()
  3624. WHERE id = $1 AND deleted_at IS NULL`,
  3625. id)
  3626. if err != nil {
  3627. return err
  3628. }
  3629. affected, err := result.RowsAffected()
  3630. if err != nil {
  3631. return err
  3632. }
  3633. if affected == 0 {
  3634. return service.ErrAccountNotFound
  3635. }
  3636. // 重置配额后触发调度快照刷新,使账号重新参与调度
  3637. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
  3638. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue quota reset failed: account=%d err=%v", id, err)
  3639. }
  3640. r.syncSchedulerAccountSnapshot(ctx, id)
  3641. return nil
  3642. }
  3643. // RevertProxyFallback 将账号的 proxy_id 切回 proxy_fallback_origin_id,并清空 origin 字段。
  3644. // 仅当 proxy_fallback_origin_id IS NOT NULL 时执行更新;
  3645. // 若影响行数为 0,则返回 ErrAccountNotInFallback(账号存在但不在 fallback 状态)。
  3646. func (r *accountRepository) RevertProxyFallback(ctx context.Context, accountID int64) error {
  3647. res, err := r.sql.ExecContext(ctx, `
  3648. UPDATE accounts SET proxy_id=proxy_fallback_origin_id, proxy_fallback_origin_id=NULL, updated_at=NOW()
  3649. WHERE id=$1 AND proxy_fallback_origin_id IS NOT NULL AND deleted_at IS NULL`, accountID)
  3650. if err != nil {
  3651. return err
  3652. }
  3653. n, _ := res.RowsAffected()
  3654. if n == 0 {
  3655. return service.ErrAccountNotInFallback
  3656. }
  3657. if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &accountID, nil, nil); err != nil {
  3658. logger.LegacyPrintf("repository.account", "[SchedulerOutbox] revert fallback enqueue failed: account=%d err=%v", accountID, err)
  3659. }
  3660. return nil
  3661. }
  3662. // ListShadowsByParent 返回指定父账号的影子账号;当前实现仅查 quota_dimension='spark'(唯一预设)。
  3663. // 同时过滤 parent_account_id 和 quota_dimension='spark',防止未来其它 linked 维度被误伤。
  3664. // ⚠️ 新增影子维度时:须更新此函数(或新增维度专用列举),并检查所有调用点(级联删除/一母一影校验/type 守卫),否则会静默漏掉新维度。
  3665. // 软删除行由 SoftDeleteMixin 拦截器自动排除,无需手写 deleted_at IS NULL。
  3666. func (r *accountRepository) ListShadowsByParent(ctx context.Context, parentID int64) ([]*service.Account, error) {
  3667. rows, err := r.client.Account.Query().
  3668. Where(dbaccount.ParentAccountIDEQ(parentID), dbaccount.QuotaDimensionEQ(dbaccount.QuotaDimensionSpark)).
  3669. All(ctx)
  3670. if err != nil {
  3671. return nil, err
  3672. }
  3673. out := make([]*service.Account, 0, len(rows))
  3674. for _, m := range rows {
  3675. out = append(out, accountEntityToService(m))
  3676. }
  3677. return out, nil
  3678. }