| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021302230233024302530263027302830293030303130323033303430353036303730383039304030413042304330443045304630473048304930503051305230533054305530563057305830593060306130623063306430653066306730683069307030713072307330743075307630773078307930803081308230833084308530863087308830893090309130923093309430953096309730983099310031013102310331043105310631073108310931103111311231133114311531163117311831193120312131223123312431253126312731283129313031313132313331343135313631373138313931403141314231433144314531463147314831493150315131523153315431553156315731583159316031613162316331643165316631673168316931703171317231733174317531763177317831793180318131823183318431853186318731883189319031913192319331943195319631973198319932003201320232033204320532063207320832093210321132123213321432153216321732183219322032213222322332243225322632273228322932303231323232333234323532363237323832393240324132423243324432453246324732483249325032513252325332543255325632573258325932603261326232633264326532663267326832693270327132723273327432753276327732783279328032813282328332843285328632873288328932903291329232933294329532963297329832993300330133023303330433053306330733083309331033113312331333143315331633173318331933203321332233233324332533263327332833293330333133323333333433353336333733383339334033413342334333443345334633473348334933503351335233533354335533563357335833593360336133623363336433653366336733683369337033713372337333743375337633773378337933803381338233833384338533863387338833893390339133923393339433953396339733983399340034013402340334043405340634073408340934103411341234133414341534163417341834193420342134223423342434253426342734283429343034313432343334343435343634373438343934403441344234433444344534463447344834493450345134523453345434553456345734583459346034613462346334643465346634673468346934703471347234733474347534763477347834793480348134823483348434853486348734883489349034913492349334943495349634973498349935003501350235033504350535063507350835093510351135123513351435153516351735183519352035213522352335243525352635273528352935303531353235333534353535363537353835393540354135423543354435453546354735483549355035513552355335543555355635573558355935603561356235633564356535663567356835693570357135723573357435753576357735783579358035813582358335843585358635873588358935903591359235933594359535963597359835993600360136023603360436053606360736083609361036113612361336143615361636173618361936203621362236233624362536263627362836293630363136323633363436353636363736383639364036413642364336443645364636473648364936503651365236533654365536563657365836593660366136623663366436653666366736683669367036713672367336743675367636773678367936803681368236833684368536863687368836893690369136923693369436953696369736983699370037013702370337043705370637073708370937103711371237133714371537163717371837193720372137223723372437253726372737283729373037313732373337343735373637373738373937403741374237433744374537463747374837493750375137523753375437553756375737583759376037613762376337643765376637673768376937703771377237733774377537763777377837793780378137823783378437853786378737883789379037913792379337943795379637973798379938003801380238033804380538063807380838093810381138123813381438153816381738183819382038213822382338243825382638273828382938303831383238333834383538363837383838393840384138423843384438453846384738483849385038513852385338543855385638573858385938603861386238633864386538663867386838693870387138723873387438753876387738783879388038813882388338843885388638873888388938903891389238933894389538963897389838993900390139023903390439053906390739083909391039113912391339143915 |
- // Package repository 实现数据访问层(Repository Pattern)。
- //
- // 该包提供了与数据库交互的所有操作,包括 CRUD、复杂查询和批量操作。
- // 采用 Repository 模式将数据访问逻辑与业务逻辑分离,便于测试和维护。
- //
- // 主要特性:
- // - 使用 Ent ORM 进行类型安全的数据库操作
- // - 对于复杂查询(如批量更新、聚合统计)使用原生 SQL
- // - 提供统一的错误翻译机制,将数据库错误转换为业务错误
- // - 支持软删除,所有查询自动过滤已删除记录
- package repository
- import (
- "context"
- "database/sql"
- "encoding/json"
- "errors"
- "math/rand/v2"
- "strconv"
- "strings"
- "time"
- dbent "github.com/Wei-Shaw/sub2api/ent"
- dbaccount "github.com/Wei-Shaw/sub2api/ent/account"
- dbaccountgroup "github.com/Wei-Shaw/sub2api/ent/accountgroup"
- dbgroup "github.com/Wei-Shaw/sub2api/ent/group"
- dbpredicate "github.com/Wei-Shaw/sub2api/ent/predicate"
- dbproxy "github.com/Wei-Shaw/sub2api/ent/proxy"
- "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
- "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
- "github.com/Wei-Shaw/sub2api/internal/service"
- "github.com/lib/pq"
- entsql "entgo.io/ent/dialect/sql"
- "entgo.io/ent/dialect/sql/sqljson"
- )
- // accountRepository 实现 service.AccountRepository 接口。
- // 提供 AI API 账户的完整数据访问功能。
- //
- // 设计说明:
- // - client: Ent 客户端,用于类型安全的 ORM 操作
- // - sql: 原生 SQL 执行器,用于复杂查询和批量操作
- // - schedulerCache: 调度器缓存,用于在账号状态变更时同步快照
- type accountRepository struct {
- client *dbent.Client // Ent ORM 客户端
- sql sqlExecutor // 原生 SQL 执行接口
- // schedulerCache 用于在账号状态变更时主动同步快照到缓存,
- // 确保粘性会话能及时感知账号不可用状态。
- // Used to proactively sync account snapshot to cache when status changes,
- // ensuring sticky sessions can promptly detect unavailable accounts.
- schedulerCache service.SchedulerCache
- }
- var schedulerNeutralExtraKeyPrefixes = []string{
- "codex_primary_",
- "codex_secondary_",
- "codex_5h_",
- "codex_7d_",
- "codex_reset_credit_",
- "passive_usage_",
- "upstream_billing_probe",
- "upstream_billing_rate_sync",
- "ollama_cloud_usage",
- }
- var schedulerNeutralExtraKeys = map[string]struct{}{
- "codex_usage_updated_at": {},
- "grok_billing_snapshot": {},
- "session_window_utilization": {},
- }
- const postgresParameterBatchSize = 50000
- const codexFingerprintSeedCanonicalPattern = "^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
- const codexFingerprintNilSeed = "00000000-0000-0000-0000-000000000000"
- func codexFingerprintSeedValidSQL(extraExpr string) string {
- value := "(" + extraExpr + " ->> 'codex_fingerprint_seed')"
- return "(" + value + " ~ '" + codexFingerprintSeedCanonicalPattern + "' AND " + value + " <> '" + codexFingerprintNilSeed + "')"
- }
- func ensureCodexFingerprintSeedSQL(extraExpr string) string {
- return "CASE WHEN platform = 'openai' AND type = 'oauth' THEN " +
- "jsonb_set(" + extraExpr + ", '{codex_fingerprint_seed}', " +
- "CASE WHEN " + codexFingerprintSeedValidSQL("extra") +
- " THEN to_jsonb(extra ->> 'codex_fingerprint_seed') ELSE to_jsonb(gen_random_uuid()::text) END, true) " +
- "ELSE " + extraExpr + " END"
- }
- func stripCodexFingerprintSeedFromExtraUpdate(extra map[string]any) map[string]any {
- if extra == nil {
- return nil
- }
- if _, exists := extra["codex_fingerprint_seed"]; !exists {
- return extra
- }
- stripped := make(map[string]any, len(extra)-1)
- for key, value := range extra {
- if key == "codex_fingerprint_seed" {
- continue
- }
- stripped[key] = value
- }
- return stripped
- }
- // NewAccountRepository 创建账户仓储实例。
- // 这是对外暴露的构造函数,返回接口类型以便于依赖注入。
- func NewAccountRepository(client *dbent.Client, sqlDB *sql.DB, schedulerCache service.SchedulerCache) service.AccountRepository {
- return newAccountRepositoryWithSQL(client, sqlDB, schedulerCache)
- }
- // NewAdminAccountRepository exposes the account repository's atomic duplication capability
- // as an explicit dependency of the admin service.
- func NewAdminAccountRepository(client *dbent.Client, sqlDB *sql.DB, schedulerCache service.SchedulerCache) service.AdminAccountRepository {
- return newAccountRepositoryWithSQL(client, sqlDB, schedulerCache)
- }
- // newAccountRepositoryWithSQL 是内部构造函数,支持依赖注入 SQL 执行器。
- // 这种设计便于单元测试时注入 mock 对象。
- func newAccountRepositoryWithSQL(client *dbent.Client, sqlq sqlExecutor, schedulerCache service.SchedulerCache) *accountRepository {
- return &accountRepository{client: client, sql: sqlq, schedulerCache: schedulerCache}
- }
- func (r *accountRepository) Create(ctx context.Context, account *service.Account) error {
- if err := createAccountRecord(ctx, r.client, account); err != nil {
- return err
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &account.ID, nil, buildSchedulerGroupPayload(account.GroupIDs)); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue account create failed: account=%d err=%v", account.ID, err)
- }
- return nil
- }
- func createAccountRecord(ctx context.Context, client *dbent.Client, account *service.Account) error {
- if account == nil {
- return service.ErrAccountNilInput
- }
- service.PrepareNewAccountProtection(account)
- builder := client.Account.Create().
- SetName(account.Name).
- SetNillableNotes(account.Notes).
- SetPlatform(account.Platform).
- SetType(account.Type).
- SetCredentials(normalizeJSONMap(account.Credentials)).
- SetExtra(normalizeJSONMap(account.Extra)).
- SetConcurrency(account.Concurrency).
- SetPriority(account.Priority).
- SetStatus(account.Status).
- SetErrorMessage(account.ErrorMessage).
- SetSchedulable(account.Schedulable).
- SetAutoPauseOnExpired(account.AutoPauseOnExpired)
- if account.RateMultiplier != nil {
- builder.SetRateMultiplier(*account.RateMultiplier)
- }
- if account.LoadFactor != nil {
- builder.SetLoadFactor(*account.LoadFactor)
- }
- if account.ProxyID != nil {
- builder.SetProxyID(*account.ProxyID)
- }
- if account.LastUsedAt != nil {
- builder.SetLastUsedAt(*account.LastUsedAt)
- }
- if account.ExpiresAt != nil {
- builder.SetExpiresAt(*account.ExpiresAt)
- }
- if account.RateLimitedAt != nil {
- builder.SetRateLimitedAt(*account.RateLimitedAt)
- }
- if account.RateLimitResetAt != nil {
- builder.SetRateLimitResetAt(*account.RateLimitResetAt)
- }
- if account.OverloadUntil != nil {
- builder.SetOverloadUntil(*account.OverloadUntil)
- }
- if account.SessionWindowStart != nil {
- builder.SetSessionWindowStart(*account.SessionWindowStart)
- }
- if account.SessionWindowEnd != nil {
- builder.SetSessionWindowEnd(*account.SessionWindowEnd)
- }
- if account.SessionWindowStatus != "" {
- builder.SetSessionWindowStatus(account.SessionWindowStatus)
- }
- builder.SetQuotaDimension(dbaccount.QuotaDimension(account.QuotaDimensionOrDefault()))
- if account.ParentAccountID != nil {
- builder.SetParentAccountID(*account.ParentAccountID)
- }
- created, err := builder.Save(ctx)
- if err != nil {
- return translatePersistenceError(err, service.ErrAccountNotFound, nil)
- }
- account.ID = created.ID
- account.CreatedAt = created.CreatedAt
- account.UpdatedAt = created.UpdatedAt
- return nil
- }
- // CreateWithAccountGroups atomically persists an account, its exact per-group priorities,
- // and the scheduler outbox event used to publish the new routing snapshot.
- func (r *accountRepository) CreateWithAccountGroups(ctx context.Context, account *service.Account, groups []service.AccountGroup) error {
- if account == nil {
- return service.ErrAccountNilInput
- }
- tx, err := r.client.Tx(ctx)
- if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
- return err
- }
- var txClient *dbent.Client
- if err == nil {
- defer func() { _ = tx.Rollback() }()
- txClient = tx.Client()
- } else {
- // Reuse a caller-owned transaction when this repository is already transactional.
- txClient = r.client
- }
- groupIDs := make([]int64, 0, len(groups))
- for i := range groups {
- groupIDs = append(groupIDs, groups[i].GroupID)
- }
- if err := lockLiveGroups(ctx, txClient, groupIDs); err != nil {
- return err
- }
- if err := createAccountRecord(ctx, txClient, account); err != nil {
- return err
- }
- if len(groups) > 0 {
- builders := make([]*dbent.AccountGroupCreate, 0, len(groups))
- for i := range groups {
- groups[i].AccountID = account.ID
- builders = append(builders, txClient.AccountGroup.Create().
- SetAccountID(account.ID).
- SetGroupID(groups[i].GroupID).
- SetPriority(groups[i].Priority),
- )
- }
- if _, err := txClient.AccountGroup.CreateBulk(builders...).Save(ctx); err != nil {
- return err
- }
- }
- account.GroupIDs = groupIDs
- account.AccountGroups = append([]service.AccountGroup(nil), groups...)
- if err := enqueueSchedulerOutbox(ctx, txClient, service.SchedulerOutboxEventAccountChanged, &account.ID, nil, buildSchedulerGroupPayload(groupIDs)); err != nil {
- return err
- }
- if tx != nil {
- if err := tx.Commit(); err != nil {
- return err
- }
- }
- return nil
- }
- func (r *accountRepository) GetByID(ctx context.Context, id int64) (*service.Account, error) {
- m, err := r.client.Account.Query().Where(dbaccount.IDEQ(id)).Only(ctx)
- if err != nil {
- return nil, translatePersistenceError(err, service.ErrAccountNotFound, nil)
- }
- accounts, err := r.accountsToService(ctx, []*dbent.Account{m})
- if err != nil {
- return nil, err
- }
- if len(accounts) == 0 {
- return nil, service.ErrAccountNotFound
- }
- return &accounts[0], nil
- }
- // SelectRandomActiveProxy chooses one active, non-expired proxy for an
- // account configured with proxy_mode=random. The query is intentionally
- // performed at request hydration time so the account can rotate across the
- // current pool without persisting a random choice.
- //
- // Sampling uses COUNT + OFFSET instead of ORDER BY RANDOM(): the latter
- // sorts the whole proxy table on every request and degrades as the pool
- // grows, while the count/offset pair is two indexed queries.
- func (r *accountRepository) SelectRandomActiveProxy(ctx context.Context) (*service.Proxy, error) {
- if r == nil || r.client == nil {
- return nil, nil
- }
- eligible := func() *dbent.ProxyQuery {
- return r.client.Proxy.Query().
- Where(
- dbproxy.StatusEQ(service.StatusActive),
- dbproxy.DeletedAtIsNil(),
- dbproxy.Or(
- dbproxy.ExpiresAtIsNil(),
- dbproxy.ExpiresAtGT(time.Now()),
- ),
- )
- }
- total, err := eligible().Count(ctx)
- if err != nil {
- return nil, err
- }
- if total == 0 {
- return nil, nil
- }
- offset := 0
- if total > 1 {
- offset = rand.IntN(total)
- }
- // Deterministic ID order keeps OFFSET sampling well-defined. A row that
- // disappears between COUNT and SELECT surfaces as NotFound and is
- // treated like an empty pool so callers fail closed.
- item, err := eligible().Order(dbproxy.ByID()).Offset(offset).First(ctx)
- if err != nil {
- if dbent.IsNotFound(err) {
- return nil, nil
- }
- return nil, err
- }
- return proxyEntityToService(item), nil
- }
- func (r *accountRepository) GetByIDs(ctx context.Context, ids []int64) ([]*service.Account, error) {
- if len(ids) == 0 {
- return []*service.Account{}, nil
- }
- // De-duplicate while preserving order of first occurrence.
- uniqueIDs := make([]int64, 0, len(ids))
- seen := make(map[int64]struct{}, len(ids))
- for _, id := range ids {
- if id <= 0 {
- continue
- }
- if _, ok := seen[id]; ok {
- continue
- }
- seen[id] = struct{}{}
- uniqueIDs = append(uniqueIDs, id)
- }
- if len(uniqueIDs) == 0 {
- return []*service.Account{}, nil
- }
- entAccounts, err := r.client.Account.
- Query().
- Where(dbaccount.IDIn(uniqueIDs...)).
- WithProxy().
- All(ctx)
- if err != nil {
- return nil, err
- }
- if len(entAccounts) == 0 {
- return []*service.Account{}, nil
- }
- accountIDs := make([]int64, 0, len(entAccounts))
- entByID := make(map[int64]*dbent.Account, len(entAccounts))
- for _, acc := range entAccounts {
- entByID[acc.ID] = acc
- accountIDs = append(accountIDs, acc.ID)
- }
- groupsByAccount, groupIDsByAccount, accountGroupsByAccount, err := r.loadAccountGroups(ctx, accountIDs)
- if err != nil {
- return nil, err
- }
- outByID := make(map[int64]*service.Account, len(entAccounts))
- for _, entAcc := range entAccounts {
- out := accountEntityToService(entAcc)
- if out == nil {
- continue
- }
- // Prefer the preloaded proxy edge when available.
- if entAcc.Edges.Proxy != nil {
- out.Proxy = proxyEntityToService(entAcc.Edges.Proxy)
- }
- if groups, ok := groupsByAccount[entAcc.ID]; ok {
- out.Groups = groups
- }
- if groupIDs, ok := groupIDsByAccount[entAcc.ID]; ok {
- out.GroupIDs = groupIDs
- }
- if ags, ok := accountGroupsByAccount[entAcc.ID]; ok {
- out.AccountGroups = ags
- }
- outByID[entAcc.ID] = out
- }
- // Preserve input order (first occurrence), and ignore missing IDs.
- out := make([]*service.Account, 0, len(uniqueIDs))
- for _, id := range uniqueIDs {
- if _, ok := entByID[id]; !ok {
- continue
- }
- if acc, ok := outByID[id]; ok && acc != nil {
- out = append(out, acc)
- }
- }
- return out, nil
- }
- // ExistsByID 检查指定 ID 的账号是否存在。
- // 相比 GetByID,此方法性能更优,因为:
- // - 使用 Exist() 方法生成 SELECT EXISTS 查询,只返回布尔值
- // - 不加载完整的账号实体及其关联数据(Groups、Proxy 等)
- // - 适用于删除前的存在性检查等只需判断有无的场景
- func (r *accountRepository) ExistsByID(ctx context.Context, id int64) (bool, error) {
- exists, err := r.client.Account.Query().Where(dbaccount.IDEQ(id)).Exist(ctx)
- if err != nil {
- return false, err
- }
- return exists, nil
- }
- func (r *accountRepository) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*service.Account, error) {
- if crsAccountID == "" {
- return nil, nil
- }
- // 使用 sqljson.ValueEQ 生成 JSON 路径过滤,避免手写 SQL 片段导致语法兼容问题。
- // 排除 spark 影子账号(parent_account_id 非空):影子不持凭据,绝不能被 CRS 当作普通账号
- // 更新而覆盖 type/credentials/proxy。即便影子 Extra 被误写入 crs_account_id 也不会命中
- // (外审第7轮 P1)。
- m, err := r.client.Account.Query().
- Where(dbaccount.ParentAccountIDIsNil()).
- Where(func(s *entsql.Selector) {
- s.Where(sqljson.ValueEQ(dbaccount.FieldExtra, crsAccountID, sqljson.Path("crs_account_id")))
- }).
- Only(ctx)
- if err != nil {
- if dbent.IsNotFound(err) {
- return nil, nil
- }
- return nil, err
- }
- accounts, err := r.accountsToService(ctx, []*dbent.Account{m})
- if err != nil {
- return nil, err
- }
- if len(accounts) == 0 {
- return nil, nil
- }
- return &accounts[0], nil
- }
- func (r *accountRepository) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
- // parent_account_id IS NULL 排除 spark 影子账号:影子不是 CRS 账号,绝不能进 CRS 同步映射
- // (否则会被当普通账号更新而覆盖 type/credentials/proxy)(外审第7轮 P1)。
- rows, err := r.sql.QueryContext(ctx, `
- SELECT id, extra->>'crs_account_id'
- FROM accounts
- WHERE deleted_at IS NULL
- AND parent_account_id IS NULL
- AND extra->>'crs_account_id' IS NOT NULL
- AND extra->>'crs_account_id' != ''
- `)
- if err != nil {
- return nil, err
- }
- defer func() { _ = rows.Close() }()
- result := make(map[string]int64)
- for rows.Next() {
- var id int64
- var crsID string
- if err := rows.Scan(&id, &crsID); err != nil {
- return nil, err
- }
- result[crsID] = id
- }
- if err := rows.Err(); err != nil {
- return nil, err
- }
- return result, nil
- }
- func (r *accountRepository) Update(ctx context.Context, account *service.Account) error {
- return r.updateAccount(ctx, account, nil, nil, account.RateMultiplier)
- }
- // UpdateWithAccountBillingSettings applies an admin account edit while
- // preserving a concurrently probe-synchronized rate unless the request
- // explicitly includes a manual rate.
- func (r *accountRepository) UpdateWithAccountBillingSettings(
- ctx context.Context,
- account *service.Account,
- probeEnabled *bool,
- rateSyncEnabled *bool,
- rateMultiplier *float64,
- ) error {
- return r.updateAccount(ctx, account, probeEnabled, rateSyncEnabled, rateMultiplier)
- }
- func (r *accountRepository) updateAccount(
- ctx context.Context,
- account *service.Account,
- explicitProbeEnabled *bool,
- explicitRateSyncEnabled *bool,
- explicitRateMultiplier *float64,
- ) error {
- if account == nil {
- return nil
- }
- baseCtx := ctx
- contextTx := dbent.TxFromContext(ctx)
- client := r.client
- var tx *dbent.Tx
- if contextTx != nil {
- client = contextTx.Client()
- } else {
- var err error
- tx, err = r.client.Tx(ctx)
- if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
- return err
- }
- if tx != nil {
- defer func() { _ = tx.Rollback() }()
- ctx = dbent.NewTxContext(ctx, tx)
- client = tx.Client()
- }
- }
- updated, err := r.updateLockedAccount(
- ctx,
- client,
- account,
- explicitProbeEnabled,
- explicitRateSyncEnabled,
- explicitRateMultiplier,
- )
- if err != nil {
- return translatePersistenceError(err, service.ErrAccountNotFound, nil)
- }
- if err := enqueueSchedulerOutbox(ctx, client, service.SchedulerOutboxEventAccountChanged, &account.ID, nil, buildSchedulerGroupPayload(account.GroupIDs)); err != nil {
- return err
- }
- if tx != nil {
- if err := tx.Commit(); err != nil {
- return err
- }
- }
- account.UpdatedAt = updated.UpdatedAt
- // 普通账号编辑(如 model_mapping / credentials)也需要立即刷新单账号快照,
- // 否则网关在 outbox worker 延迟或异常时仍可能读到旧配置。
- if contextTx == nil {
- r.syncSchedulerAccountSnapshot(baseCtx, account.ID)
- }
- return nil
- }
- func (r *accountRepository) updateLockedAccount(
- ctx context.Context,
- client *dbent.Client,
- account *service.Account,
- explicitProbeEnabled *bool,
- explicitRateSyncEnabled *bool,
- explicitRateMultiplier *float64,
- ) (*dbent.Account, error) {
- extra, err := lockAndMergeAccountProbeExtra(ctx, client, account, explicitProbeEnabled, explicitRateSyncEnabled)
- if err != nil {
- return nil, err
- }
- account.Extra = extra
- if err := preserveLockedAccountProtection(ctx, client, account); err != nil {
- return nil, err
- }
- extra = normalizeJSONMap(account.Extra)
- schedulable := account.Schedulable
- if account.Status == service.StatusError {
- schedulable = false
- }
- builder := client.Account.UpdateOneID(account.ID).
- SetName(account.Name).
- SetNillableNotes(account.Notes).
- SetPlatform(account.Platform).
- SetType(account.Type).
- SetCredentials(normalizeJSONMap(account.Credentials)).
- SetExtra(extra).
- SetConcurrency(account.Concurrency).
- SetPriority(account.Priority).
- SetStatus(account.Status).
- SetErrorMessage(account.ErrorMessage).
- SetSchedulable(schedulable).
- SetAutoPauseOnExpired(account.AutoPauseOnExpired)
- if explicitRateMultiplier != nil {
- builder.SetRateMultiplier(*explicitRateMultiplier)
- }
- if account.LoadFactor != nil {
- builder.SetLoadFactor(*account.LoadFactor)
- } else {
- builder.ClearLoadFactor()
- }
- if account.IsRandomProxy() {
- builder.ClearProxyID()
- } else if account.ProxyID != nil {
- builder.SetProxyID(*account.ProxyID)
- } else {
- builder.ClearProxyID()
- }
- if account.LastUsedAt != nil {
- builder.SetLastUsedAt(*account.LastUsedAt)
- } else {
- builder.ClearLastUsedAt()
- }
- if account.ExpiresAt != nil {
- builder.SetExpiresAt(*account.ExpiresAt)
- } else {
- builder.ClearExpiresAt()
- }
- if account.RateLimitedAt != nil {
- builder.SetRateLimitedAt(*account.RateLimitedAt)
- } else {
- builder.ClearRateLimitedAt()
- }
- if account.RateLimitResetAt != nil {
- builder.SetRateLimitResetAt(*account.RateLimitResetAt)
- } else {
- builder.ClearRateLimitResetAt()
- }
- if account.OverloadUntil != nil {
- builder.SetOverloadUntil(*account.OverloadUntil)
- } else {
- builder.ClearOverloadUntil()
- }
- if account.SessionWindowStart != nil {
- builder.SetSessionWindowStart(*account.SessionWindowStart)
- } else {
- builder.ClearSessionWindowStart()
- }
- if account.SessionWindowEnd != nil {
- builder.SetSessionWindowEnd(*account.SessionWindowEnd)
- } else {
- builder.ClearSessionWindowEnd()
- }
- if account.SessionWindowStatus != "" {
- builder.SetSessionWindowStatus(account.SessionWindowStatus)
- } else {
- builder.ClearSessionWindowStatus()
- }
- if account.Notes == nil {
- builder.ClearNotes()
- }
- builder.SetQuotaDimension(dbaccount.QuotaDimension(account.QuotaDimensionOrDefault()))
- builder.SetNillableParentAccountID(account.ParentAccountID)
- return builder.Save(ctx)
- }
- func lockAndMergeAccountProbeExtra(
- ctx context.Context,
- client *dbent.Client,
- account *service.Account,
- explicitProbeEnabled *bool,
- explicitRateSyncEnabled *bool,
- ) (map[string]any, error) {
- credentials, err := json.Marshal(normalizeJSONMap(account.Credentials))
- if err != nil {
- return nil, err
- }
- var proxyID any
- if account.ProxyID != nil {
- proxyID = *account.ProxyID
- }
- rows, err := client.QueryContext(ctx, `
- SELECT
- platform = $2
- AND type = $3
- AND credentials = $4::jsonb
- AND proxy_id IS NOT DISTINCT FROM $5,
- COALESCE(
- platform IN (`+ollamaCloudUsagePlatformsSQL+`)
- AND $2 IN (`+ollamaCloudUsagePlatformsSQL+`)
- AND type = 'apikey'
- AND $3 = 'apikey'
- AND credentials -> 'api_key' IS NOT DISTINCT FROM $4::jsonb -> 'api_key'
- AND `+ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'")+`
- AND `+ollamaCloudBaseURLMatchesSQL("$4::jsonb ->> 'base_url'")+`,
- false
- ),
- proxy_id IS NOT DISTINCT FROM $5,
- extra -> 'upstream_billing_probe_enabled',
- extra -> 'upstream_billing_rate_sync_enabled',
- extra -> 'upstream_billing_probe',
- extra -> 'ollama_cloud_usage_session',
- extra -> 'ollama_cloud_usage_auto_refresh',
- extra -> 'ollama_cloud_usage_snapshot'
- FROM accounts
- WHERE id = $1 AND deleted_at IS NULL
- FOR NO KEY UPDATE
- `, account.ID, account.Platform, account.Type, string(credentials), proxyID)
- if err != nil {
- return nil, err
- }
- defer func() { _ = rows.Close() }()
- if !rows.Next() {
- if err := rows.Err(); err != nil {
- return nil, err
- }
- return nil, service.ErrAccountNotFound
- }
- var (
- identityUnchanged bool
- ollamaGroupIdentityUnchanged bool
- ollamaProxyIdentityUnchanged bool
- currentEnabled []byte
- currentRateSyncEnabled []byte
- currentSnapshot []byte
- currentOllamaSession []byte
- currentOllamaAutoRefresh []byte
- currentOllamaSnapshot []byte
- )
- if err := rows.Scan(
- &identityUnchanged,
- &ollamaGroupIdentityUnchanged,
- &ollamaProxyIdentityUnchanged,
- ¤tEnabled,
- ¤tRateSyncEnabled,
- ¤tSnapshot,
- ¤tOllamaSession,
- ¤tOllamaAutoRefresh,
- ¤tOllamaSnapshot,
- ); err != nil {
- return nil, err
- }
- if err := rows.Err(); err != nil {
- return nil, err
- }
- extra := copyJSONMap(normalizeJSONMap(account.Extra))
- for _, key := range []string{
- service.UpstreamBillingProbeEnabledExtraKey,
- service.UpstreamBillingRateSyncEnabledExtraKey,
- service.UpstreamBillingProbeExtraKey,
- service.OllamaCloudUsageSessionExtraKey,
- service.OllamaCloudUsageAutoRefreshExtraKey,
- service.OllamaCloudUsageSnapshotExtraKey,
- } {
- delete(extra, key)
- }
- probeAccount := service.IsUpstreamBillingProbeIdentity(account.Platform, account.Type)
- probeEnabled := false
- probeEnabledPresent := false
- if probeAccount {
- if enabled, ok, err := decodeAccountExtraJSON(currentEnabled); err != nil {
- return nil, err
- } else if value, isBool := enabled.(bool); ok && isBool {
- probeEnabled = value
- probeEnabledPresent = true
- }
- if explicitProbeEnabled != nil {
- probeEnabled = *explicitProbeEnabled
- probeEnabledPresent = true
- }
- }
- rateSyncEnabled := false
- rateSyncEnabledPresent := false
- if probeAccount {
- if enabled, ok, err := decodeAccountExtraJSON(currentRateSyncEnabled); err != nil {
- return nil, err
- } else if value, isBool := enabled.(bool); ok && isBool {
- rateSyncEnabled = value
- rateSyncEnabledPresent = true
- }
- if explicitRateSyncEnabled != nil {
- rateSyncEnabled = *explicitRateSyncEnabled
- rateSyncEnabledPresent = true
- }
- if explicitProbeEnabled != nil && !*explicitProbeEnabled {
- rateSyncEnabled = false
- rateSyncEnabledPresent = true
- }
- // 同步依赖探测,方向是单向的:探测关闭(或探测键缺失)一律把同步归零。
- // 不做反向推导——由 rate_sync=true 推出 probe=true 会让一条"同步开、探测键
- // 缺失"的僵尸记录在任意一次无关编辑时静默打开周期性外呼。需要同时打开两个
- // 开关的调用方(管理端编辑)自己显式传 explicitProbeEnabled=true。
- if !probeEnabled {
- rateSyncEnabled = false
- }
- if probeEnabledPresent {
- extra[service.UpstreamBillingProbeEnabledExtraKey] = probeEnabled
- }
- if rateSyncEnabledPresent {
- extra[service.UpstreamBillingRateSyncEnabledExtraKey] = rateSyncEnabled
- }
- }
- probeExplicitlyDisabled := probeEnabledPresent && !probeEnabled
- if identityUnchanged && !probeExplicitlyDisabled {
- if snapshot, ok, err := decodeAccountExtraJSON(currentSnapshot); err != nil {
- return nil, err
- } else if ok {
- extra[service.UpstreamBillingProbeExtraKey] = snapshot
- }
- }
- if service.IsOllamaCloudUsageAccount(account) && ollamaGroupIdentityUnchanged {
- for key, raw := range map[string][]byte{
- service.OllamaCloudUsageSessionExtraKey: currentOllamaSession,
- service.OllamaCloudUsageAutoRefreshExtraKey: currentOllamaAutoRefresh,
- } {
- if value, ok, err := decodeAccountExtraJSON(raw); err != nil {
- return nil, err
- } else if ok {
- extra[key] = value
- }
- }
- if ollamaProxyIdentityUnchanged {
- if snapshot, ok, err := decodeAccountExtraJSON(currentOllamaSnapshot); err != nil {
- return nil, err
- } else if ok {
- extra[service.OllamaCloudUsageSnapshotExtraKey] = snapshot
- }
- }
- }
- return extra, nil
- }
- func decodeAccountExtraJSON(raw []byte) (any, bool, error) {
- if len(raw) == 0 || string(raw) == "null" {
- return nil, false, nil
- }
- var value any
- if err := json.Unmarshal(raw, &value); err != nil {
- return nil, false, err
- }
- return value, true, nil
- }
- func (r *accountRepository) UpdateCredentials(ctx context.Context, id int64, credentials map[string]any) error {
- payload, err := json.Marshal(normalizeJSONMap(credentials))
- if err != nil {
- return err
- }
- baseCtx := ctx
- contextTx := dbent.TxFromContext(ctx)
- client := r.client
- var tx *dbent.Tx
- if contextTx != nil {
- client = contextTx.Client()
- } else if r.client != nil {
- var txErr error
- tx, txErr = r.client.Tx(ctx)
- if txErr != nil && !errors.Is(txErr, dbent.ErrTxStarted) {
- return txErr
- }
- if tx != nil {
- defer func() { _ = tx.Rollback() }()
- ctx = dbent.NewTxContext(ctx, tx)
- client = tx.Client()
- }
- }
- result, err := client.ExecContext(ctx, `
- UPDATE accounts
- SET
- credentials = $1::jsonb,
- extra = CASE
- -- 凭证整体未变化 ⇒ Ollama 组身份必然未变化;顶层 DISTINCT 守卫防止
- -- 非 Ollama 账号的无变化持久化误清探测快照或重写 NULL extra。
- WHEN platform IN (`+ollamaCloudUsagePlatformsSQL+`)
- AND type = 'apikey'
- AND credentials IS DISTINCT FROM $1::jsonb
- AND (
- credentials -> 'api_key' IS DISTINCT FROM $1::jsonb -> 'api_key'
- OR NOT (
- `+ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'")+`
- AND `+ollamaCloudBaseURLMatchesSQL("$1::jsonb ->> 'base_url'")+`
- )
- )
- THEN COALESCE(extra, '{}'::jsonb)
- - 'upstream_billing_probe'
- - 'ollama_cloud_usage_session'
- - 'ollama_cloud_usage_auto_refresh'
- - 'ollama_cloud_usage_snapshot'
- -- 上游倍率探测已放宽到全部 API-key 平台:凭证变化即视为探测
- -- 身份变化,丢弃 stale 快照。
- WHEN type = 'apikey'
- AND credentials IS DISTINCT FROM $1::jsonb
- THEN COALESCE(extra, '{}'::jsonb) - 'upstream_billing_probe'
- ELSE extra
- END,
- updated_at = NOW()
- WHERE id = $2 AND deleted_at IS NULL
- `, string(payload), id)
- if err != nil {
- return err
- }
- affected, err := result.RowsAffected()
- if err != nil {
- return err
- }
- if affected == 0 {
- return service.ErrAccountNotFound
- }
- if err := enqueueSchedulerOutbox(ctx, client, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- return err
- }
- if tx != nil {
- if err := tx.Commit(); err != nil {
- return err
- }
- }
- if contextTx == nil {
- r.syncSchedulerAccountSnapshot(baseCtx, id)
- }
- return nil
- }
- func (r *accountRepository) Delete(ctx context.Context, id int64) error {
- groupIDs, err := r.loadAccountGroupIDs(ctx, id)
- if err != nil {
- return err
- }
- // 使用事务保证账号与关联分组的删除原子性
- tx, err := r.client.Tx(ctx)
- if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
- return err
- }
- var txClient *dbent.Client
- if err == nil {
- defer func() { _ = tx.Rollback() }()
- txClient = tx.Client()
- } else {
- // 已处于外部事务中(ErrTxStarted),复用当前 client
- txClient = r.client
- }
- if _, err := txClient.AccountGroup.Delete().Where(dbaccountgroup.AccountIDEQ(id)).Exec(ctx); err != nil {
- return err
- }
- if _, err := txClient.ExecContext(ctx, "DELETE FROM scheduled_test_plans WHERE account_id = $1", id); err != nil {
- return err
- }
- if _, err := txClient.Account.Delete().Where(dbaccount.IDEQ(id)).Exec(ctx); err != nil {
- return err
- }
- if tx != nil {
- if err := tx.Commit(); err != nil {
- return err
- }
- }
- r.deleteSchedulerAccountSnapshot(ctx, id)
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, buildSchedulerGroupPayload(groupIDs)); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue account delete failed: account=%d err=%v", id, err)
- }
- return nil
- }
- func (r *accountRepository) List(ctx context.Context, params pagination.PaginationParams) ([]service.Account, *pagination.PaginationResult, error) {
- return r.ListWithFilters(ctx, params, "", "", "", "", 0, "")
- }
- func (r *accountRepository) accountListFilteredQuery(platform, accountType, status, search string, groupID int64, privacyMode string) *dbent.AccountQuery {
- q := r.client.Account.Query()
- if platform != "" {
- q = q.Where(dbaccount.PlatformEQ(platform))
- }
- if accountType != "" {
- q = q.Where(dbaccount.TypeEQ(accountType))
- }
- if status != "" {
- switch status {
- case service.StatusActive:
- q = q.Where(
- dbaccount.StatusEQ(status),
- dbaccount.SchedulableEQ(true),
- dbaccount.Or(
- dbaccount.RateLimitResetAtIsNil(),
- dbaccount.RateLimitResetAtLTE(time.Now()),
- ),
- dbpredicate.Account(func(s *entsql.Selector) {
- col := s.C("temp_unschedulable_until")
- s.Where(entsql.Or(
- entsql.IsNull(col),
- entsql.LTE(col, entsql.Expr("NOW()")),
- ))
- }),
- )
- case "rate_limited":
- q = q.Where(
- dbaccount.StatusEQ(service.StatusActive),
- dbaccount.RateLimitResetAtGT(time.Now()),
- dbpredicate.Account(func(s *entsql.Selector) {
- col := s.C("temp_unschedulable_until")
- s.Where(entsql.Or(
- entsql.IsNull(col),
- entsql.LTE(col, entsql.Expr("NOW()")),
- ))
- }),
- )
- case "temp_unschedulable":
- q = q.Where(
- dbaccount.StatusEQ(service.StatusActive),
- dbpredicate.Account(func(s *entsql.Selector) {
- col := s.C("temp_unschedulable_until")
- s.Where(entsql.And(
- entsql.Not(entsql.IsNull(col)),
- entsql.GT(col, entsql.Expr("NOW()")),
- ))
- }),
- )
- case "unschedulable":
- q = q.Where(
- dbaccount.StatusEQ(service.StatusActive),
- dbaccount.SchedulableEQ(false),
- dbaccount.Or(
- dbaccount.RateLimitResetAtIsNil(),
- dbaccount.RateLimitResetAtLTE(time.Now()),
- ),
- dbpredicate.Account(func(s *entsql.Selector) {
- col := s.C("temp_unschedulable_until")
- s.Where(entsql.Or(
- entsql.IsNull(col),
- entsql.LTE(col, entsql.Expr("NOW()")),
- ))
- }),
- )
- default:
- q = q.Where(dbaccount.StatusEQ(status))
- }
- }
- if search != "" {
- q = q.Where(dbaccount.NameContainsFold(search))
- }
- if groupID == service.AccountListGroupUngrouped {
- q = q.Where(dbaccount.Not(dbaccount.HasAccountGroups()))
- } else if groupID > 0 {
- q = q.Where(dbaccount.HasAccountGroupsWith(dbaccountgroup.GroupIDEQ(groupID)))
- }
- if privacyMode != "" {
- q = q.Where(dbpredicate.Account(func(s *entsql.Selector) {
- path := sqljson.Path("privacy_mode")
- switch privacyMode {
- case service.AccountPrivacyModeUnsetFilter:
- s.Where(entsql.Or(
- entsql.Not(sqljson.HasKey(dbaccount.FieldExtra, path)),
- sqljson.ValueEQ(dbaccount.FieldExtra, "", path),
- ))
- default:
- s.Where(sqljson.ValueEQ(dbaccount.FieldExtra, privacyMode, path))
- }
- }))
- }
- return q
- }
- func (r *accountRepository) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]service.Account, *pagination.PaginationResult, error) {
- q := r.accountListFilteredQuery(platform, accountType, status, search, groupID, privacyMode)
- // Clone before Count so interceptor-appended predicates (SoftDeleteMixin's
- // deleted_at IS NULL) don't accumulate on the shared builder and pollute the
- // subsequent list query. Same pattern used in group_repo/promo_code_repo/user_repo
- // (P1-03 audit fix, commit 2588fa6a).
- total, err := q.Clone().Count(ctx)
- if err != nil {
- return nil, nil, err
- }
- accountsQuery := q.
- Offset(params.Offset()).
- Limit(params.Limit())
- for _, order := range accountListOrder(params) {
- accountsQuery = accountsQuery.Order(order)
- }
- accounts, err := accountsQuery.All(ctx)
- if err != nil {
- return nil, nil, err
- }
- outAccounts, err := r.accountsToService(ctx, accounts)
- if err != nil {
- return nil, nil, err
- }
- return outAccounts, paginationResultFromTotal(int64(total), params), nil
- }
- func (r *accountRepository) ListAllWithFilters(ctx context.Context, platform, accountType, status, search string, groupID int64, privacyMode string) ([]service.Account, error) {
- accounts, err := r.accountListFilteredQuery(platform, accountType, status, search, groupID, privacyMode).All(ctx)
- if err != nil {
- return nil, err
- }
- return r.accountsToService(ctx, accounts)
- }
- func (r *accountRepository) ListOpsAccountsForStats(ctx context.Context, platformFilter string, groupIDFilter *int64) ([]service.Account, error) {
- if r == nil || r.client == nil {
- return []service.Account{}, nil
- }
- q := r.client.Account.Query()
- if platformFilter = strings.TrimSpace(platformFilter); platformFilter != "" {
- q = q.Where(dbaccount.PlatformEQ(platformFilter))
- }
- if groupIDFilter != nil && *groupIDFilter > 0 {
- q = q.Where(dbaccount.HasAccountGroupsWith(dbaccountgroup.GroupIDEQ(*groupIDFilter)))
- }
- accounts, err := q.
- Select(
- dbaccount.FieldID,
- dbaccount.FieldName,
- dbaccount.FieldPlatform,
- dbaccount.FieldConcurrency,
- dbaccount.FieldLoadFactor,
- dbaccount.FieldStatus,
- dbaccount.FieldErrorMessage,
- dbaccount.FieldSchedulable,
- dbaccount.FieldRateLimitResetAt,
- dbaccount.FieldOverloadUntil,
- dbaccount.FieldTempUnschedulableUntil,
- ).
- Order(dbent.Asc(dbaccount.FieldID)).
- All(ctx)
- if err != nil {
- return nil, err
- }
- return r.accountsToService(ctx, accounts)
- }
- func accountListOrder(params pagination.PaginationParams) []func(*entsql.Selector) {
- sortBy := strings.ToLower(strings.TrimSpace(params.SortBy))
- sortOrder := params.NormalizedSortOrder(pagination.SortOrderAsc)
- if sortBy == "upstream_billing_rate" {
- direction := "ASC"
- tieOrder := entsql.Asc
- if sortOrder == pagination.SortOrderDesc {
- direction = "DESC"
- tieOrder = entsql.Desc
- }
- return []func(*entsql.Selector){func(s *entsql.Selector) {
- extra := s.C(dbaccount.FieldExtra)
- expression := upstreamBillingRateSortExpression(extra)
- s.OrderExpr(entsql.Expr(expression + " " + direction + " NULLS LAST"))
- s.OrderBy(tieOrder(s.C(dbaccount.FieldID)))
- }}
- }
- field := dbaccount.FieldName
- defaultOrder := true
- switch sortBy {
- case "", "name":
- field = dbaccount.FieldName
- case "id":
- field = dbaccount.FieldID
- defaultOrder = false
- case "status":
- field = dbaccount.FieldStatus
- defaultOrder = false
- case "schedulable":
- field = dbaccount.FieldSchedulable
- defaultOrder = false
- case "priority":
- field = dbaccount.FieldPriority
- defaultOrder = false
- case "rate_multiplier":
- field = dbaccount.FieldRateMultiplier
- defaultOrder = false
- case "last_used_at":
- field = dbaccount.FieldLastUsedAt
- defaultOrder = false
- case "expires_at":
- field = dbaccount.FieldExpiresAt
- defaultOrder = false
- case "created_at":
- field = dbaccount.FieldCreatedAt
- defaultOrder = false
- }
- if sortOrder == pagination.SortOrderDesc {
- return []func(*entsql.Selector){dbent.Desc(field), dbent.Desc(dbaccount.FieldID)}
- }
- if defaultOrder {
- return []func(*entsql.Selector){dbent.Asc(dbaccount.FieldName), dbent.Asc(dbaccount.FieldID)}
- }
- return []func(*entsql.Selector){dbent.Asc(field), dbent.Asc(dbaccount.FieldID)}
- }
- func upstreamBillingRateSortExpression(extra string) string {
- status := extra + " #>> '{upstream_billing_probe,status}'"
- effectiveJSON := extra + " #> '{upstream_billing_probe,data,effective_rate_multiplier}'"
- effective := extra + " #>> '{upstream_billing_probe,data,effective_rate_multiplier}'"
- resolvedJSON := extra + " #> '{upstream_billing_probe,data,resolved_rate_multiplier}'"
- resolved := extra + " #>> '{upstream_billing_probe,data,resolved_rate_multiplier}'"
- peakEnabledJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_enabled}'"
- peakEnabled := extra + " #>> '{upstream_billing_probe,data,peak_rate_enabled}'"
- peakStart := extra + " #>> '{upstream_billing_probe,data,peak_start}'"
- peakEnd := extra + " #>> '{upstream_billing_probe,data,peak_end}'"
- peakMultiplierJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_multiplier}'"
- peakMultiplier := extra + " #>> '{upstream_billing_probe,data,peak_rate_multiplier}'"
- peakMultiplierValue := "(CASE WHEN jsonb_typeof(" + peakMultiplierJSON + ") = 'number' THEN (" + peakMultiplier + ")::numeric END)"
- billingScope := extra + " #>> '{upstream_billing_probe,data,billing_scope}'"
- timezone := extra + " #>> '{upstream_billing_probe,data,timezone}'"
- validClock := "'^([01][0-9]|2[0-3]):[0-5][0-9]$'"
- startMinute := "(CASE WHEN " + peakStart + " ~ " + validClock + " THEN split_part(" + peakStart + ", ':', 1)::numeric * 60 + split_part(" + peakStart + ", ':', 2)::numeric END)"
- endMinute := "(CASE WHEN " + peakEnd + " ~ " + validClock + " THEN split_part(" + peakEnd + ", ':', 1)::numeric * 60 + split_part(" + peakEnd + ", ':', 2)::numeric END)"
- localMinute := "(EXTRACT(HOUR FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))) * 60 + EXTRACT(MINUTE FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))))"
- validPeakWindow := peakStart + " ~ " + validClock + " AND " +
- peakEnd + " ~ " + validClock + " AND " +
- startMinute + " < " + endMinute
- validPeakConfig := validPeakWindow + " AND " + peakMultiplierValue + " >= 0 AND " +
- "EXISTS (SELECT 1 FROM pg_timezone_names WHERE name = " + timezone + ")"
- dynamicRate := "CASE WHEN " + peakEnabled + " = 'false' THEN (" + resolved + ")::numeric WHEN " + peakEnabled + " = 'true' AND " + validPeakConfig +
- " THEN (" + resolved + ")::numeric * CASE WHEN " + localMinute + " >= " + startMinute + " AND " + localMinute + " < " + endMinute +
- " THEN " + peakMultiplierValue + " ELSE 1 END ELSE NULL END"
- legacySnapshot := "jsonb_typeof(" + resolvedJSON + ") IS NULL AND jsonb_typeof(" + peakEnabledJSON + ") IS NULL"
- return "CASE WHEN " + status + " IN ('ok', 'failed') AND (jsonb_typeof(" + resolvedJSON + ") = 'number' OR jsonb_typeof(" + effectiveJSON + ") = 'number') THEN CASE WHEN jsonb_typeof(" +
- resolvedJSON + ") = 'number' AND jsonb_typeof(" + peakEnabledJSON + ") = 'boolean' THEN CASE WHEN " + billingScope + " = 'token' THEN " + dynamicRate + " ELSE NULL END WHEN " + legacySnapshot +
- " AND jsonb_typeof(" + effectiveJSON + ") = 'number' THEN (" + effective + ")::numeric END END"
- }
- func (r *accountRepository) ListByGroup(ctx context.Context, groupID int64) ([]service.Account, error) {
- accounts, err := r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{
- status: service.StatusActive,
- })
- if err != nil {
- return nil, err
- }
- return accounts, nil
- }
- func (r *accountRepository) ListActive(ctx context.Context) ([]service.Account, error) {
- accounts, err := r.client.Account.Query().
- Where(dbaccount.StatusEQ(service.StatusActive)).
- Order(dbent.Asc(dbaccount.FieldPriority)).
- All(ctx)
- if err != nil {
- return nil, err
- }
- return r.accountsToService(ctx, accounts)
- }
- func (r *accountRepository) ListOAuthRefreshCandidatePage(ctx context.Context, options service.OAuthRefreshPageOptions) (*service.OAuthRefreshCandidatePage, error) {
- if r.sql == nil {
- return nil, errors.New("account repository SQL executor not configured")
- }
- if len(options.Platforms) == 0 {
- return nil, errors.New("oauth refresh candidate platforms cannot be empty")
- }
- if options.Limit <= 0 || options.Limit > 1000 {
- return nil, errors.New("oauth refresh candidate page limit must be between 1 and 1000")
- }
- // (cond) IS NOT TRUE 把 NULL 和 FALSE 都视为"可被刷新"。直接写
- // NOT (a AND b) 在 PG 三值逻辑下会把 a 或 b 为 NULL 的行(即绝大多数
- // 健康账号:temp_unschedulable_until=NULL)也排除,导致后台 token
- // 刷新工作器漏掉所有正常账号 → access_token 到期后请求开始 401。
- query := `
- SELECT id
- FROM accounts
- WHERE deleted_at IS NULL
- AND schedulable = TRUE
- AND platform = ANY($1)
- AND id > $2`
- if options.ActiveOnly {
- query += `
- AND status = 'active'`
- }
- if options.IncludeSetupToken {
- query += `
- AND type IN ('oauth', 'setup-token')`
- } else {
- query += `
- AND type = 'oauth'`
- }
- if options.RequireRefreshToken {
- query += `
- AND credentials ? 'refresh_token'
- AND btrim(credentials->>'refresh_token') <> ''`
- }
- if options.ExcludeRetryCooldown {
- query += `
- AND (
- temp_unschedulable_until > NOW()
- AND temp_unschedulable_reason LIKE 'token refresh retry exhausted:%'
- ) IS NOT TRUE`
- }
- query += `
- ORDER BY id ASC
- LIMIT $3`
- rows, err := r.sql.QueryContext(ctx, query, pq.Array(options.Platforms), options.AfterID, options.Limit)
- if err != nil {
- return nil, err
- }
- defer func() { _ = rows.Close() }()
- var ids []int64
- for rows.Next() {
- var id int64
- if err := rows.Scan(&id); err != nil {
- return nil, err
- }
- ids = append(ids, id)
- }
- if err := rows.Err(); err != nil {
- return nil, err
- }
- if len(ids) == 0 {
- return &service.OAuthRefreshCandidatePage{Accounts: []service.Account{}}, nil
- }
- accounts, err := r.GetByIDs(ctx, ids)
- if err != nil {
- return nil, err
- }
- accountsByID := make(map[int64]*service.Account, len(accounts))
- for _, account := range accounts {
- if account != nil {
- accountsByID[account.ID] = account
- }
- }
- out := make([]service.Account, 0, len(accounts))
- for _, id := range ids {
- if account := accountsByID[id]; account != nil {
- out = append(out, *account)
- }
- }
- page := &service.OAuthRefreshCandidatePage{
- Accounts: out,
- HasMore: len(ids) == options.Limit,
- }
- if len(ids) > 0 {
- page.NextAfterID = ids[len(ids)-1]
- }
- return page, nil
- }
- func (r *accountRepository) ListByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
- accounts, err := r.client.Account.Query().
- Where(
- dbaccount.PlatformEQ(platform),
- dbaccount.StatusEQ(service.StatusActive),
- ).
- Order(dbent.Asc(dbaccount.FieldPriority)).
- All(ctx)
- if err != nil {
- return nil, err
- }
- return r.accountsToService(ctx, accounts)
- }
- func (r *accountRepository) UpdateLastUsed(ctx context.Context, id int64) error {
- now := time.Now()
- _, err := r.client.Account.Update().
- Where(dbaccount.IDEQ(id)).
- SetLastUsedAt(now).
- Save(ctx)
- if err != nil {
- return err
- }
- payload := map[string]any{
- "last_used": map[string]int64{
- strconv.FormatInt(id, 10): now.Unix(),
- },
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountLastUsed, &id, nil, payload); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue last used failed: account=%d err=%v", id, err)
- }
- return nil
- }
- func (r *accountRepository) BatchUpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
- if len(updates) == 0 {
- return nil
- }
- ids := make([]int64, 0, len(updates))
- args := make([]any, 0, len(updates)*2+1)
- caseSQL := "UPDATE accounts SET last_used_at = CASE id"
- idx := 1
- for id, ts := range updates {
- caseSQL += " WHEN $" + itoa(idx) + " THEN $" + itoa(idx+1) + "::timestamptz"
- args = append(args, id, ts)
- ids = append(ids, id)
- idx += 2
- }
- caseSQL += " END, updated_at = NOW() WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL"
- args = append(args, pq.Array(ids))
- _, err := r.sql.ExecContext(ctx, caseSQL, args...)
- if err != nil {
- return err
- }
- lastUsedPayload := make(map[string]int64, len(updates))
- for id, ts := range updates {
- lastUsedPayload[strconv.FormatInt(id, 10)] = ts.Unix()
- }
- payload := map[string]any{"last_used": lastUsedPayload}
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountLastUsed, nil, nil, payload); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue batch last used failed: err=%v", err)
- }
- return nil
- }
- func (r *accountRepository) SetError(ctx context.Context, id int64, errorMsg string) error {
- _, err := r.client.Account.Update().
- Where(dbaccount.IDEQ(id)).
- SetStatus(service.StatusError).
- SetErrorMessage(errorMsg).
- SetSchedulable(false).
- Save(ctx)
- if err != nil {
- return err
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue set error failed: account=%d err=%v", id, err)
- }
- r.syncSchedulerAccountSnapshot(ctx, id)
- return nil
- }
- func (r *accountRepository) SetGrokCredentialErrorIfMatch(
- ctx context.Context,
- id int64,
- snapshot service.GrokCredentialMutationSnapshot,
- errorMsg string,
- ) (bool, error) {
- result, err := r.sql.ExecContext(ctx, `
- WITH updated AS (
- UPDATE accounts AS a
- SET status = $1,
- error_message = $2,
- schedulable = false,
- updated_at = NOW()
- WHERE a.id = $3
- AND a.deleted_at IS NULL
- AND a.status = $4
- AND a.platform = $5
- AND a.type = $6
- AND a.schedulable IS TRUE
- AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW())
- AND (a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW())
- AND (a.overload_until IS NULL OR a.overload_until <= NOW())
- AND (a.auto_pause_on_expired IS NOT TRUE OR a.expires_at IS NULL OR a.expires_at > NOW())
- AND a.credentials = $7::jsonb
- AND a.proxy_id IS NOT DISTINCT FROM $8
- AND ($2 <> $9 OR (
- a.proxy_id IS NOT NULL AND NOT EXISTS (
- SELECT 1 FROM proxies p WHERE p.id = a.proxy_id AND p.deleted_at IS NULL
- )
- ))
- RETURNING a.id
- )
- INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
- SELECT $10, updated.id, NULL, NULL FROM updated
- `, service.StatusError, errorMsg, id, service.StatusActive, service.PlatformGrok, service.AccountTypeOAuth,
- snapshot.CredentialsJSON, snapshot.ProxyID, string(service.GrokCredentialReasonProxyInvalid),
- service.SchedulerOutboxEventAccountChanged)
- if err != nil {
- return false, err
- }
- affected, err := result.RowsAffected()
- if err != nil || affected == 0 {
- return false, err
- }
- r.syncSchedulerAccountSnapshotDetached(ctx, id)
- return true, nil
- }
- // SetGrokOAuthErrorIfCredentialsUnchanged atomically quarantines a structurally
- // invalid Grok OAuth account only if it is still active and its complete JSONB
- // credential document matches the state observed by reconciliation. Exact
- // JSONB equality includes _token_version when present and prevents a concurrent
- // reauthorization from being overwritten by a stale check-then-mutate path.
- func (r *accountRepository) SetGrokOAuthErrorIfCredentialsUnchanged(
- ctx context.Context,
- id int64,
- expectedCredentials map[string]any,
- errorMsg string,
- ) (bool, error) {
- if r == nil || r.sql == nil {
- return false, errors.New("account repository SQL executor is not configured")
- }
- expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials))
- if err != nil {
- return false, err
- }
- result, err := r.sql.ExecContext(ctx, `
- WITH updated AS (
- UPDATE accounts AS a
- SET status = $1,
- error_message = $2,
- schedulable = FALSE,
- updated_at = NOW()
- WHERE a.id = $3
- AND a.deleted_at IS NULL
- AND a.platform = $4
- AND a.type = $5
- AND a.status = $6
- AND a.credentials = $7::jsonb
- AND NULLIF(BTRIM(a.credentials->>'refresh_token'), '') IS NULL
- RETURNING a.id
- )
- INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
- SELECT $8, updated.id, NULL, NULL FROM updated
- `,
- service.StatusError,
- errorMsg,
- id,
- service.PlatformGrok,
- service.AccountTypeOAuth,
- service.StatusActive,
- string(expectedJSON),
- service.SchedulerOutboxEventAccountChanged,
- )
- if err != nil {
- return false, err
- }
- rowsAffected, err := result.RowsAffected()
- if err != nil {
- return false, err
- }
- if rowsAffected == 0 {
- return false, nil
- }
- r.syncSchedulerAccountSnapshotDetached(ctx, id)
- return true, nil
- }
- // UpdateGrokOAuthCredentialsIfUnchanged persists provider-issued replacement
- // credentials only while the complete Grok OAuth credential document and
- // proxy still match the fresh snapshot used by the upstream refresh call. The
- // scheduler outbox insert is part of the same PostgreSQL statement, so a
- // durable invalidation failure rolls the credential update back as well.
- func (r *accountRepository) UpdateGrokOAuthCredentialsIfUnchanged(
- ctx context.Context,
- id int64,
- expectedCredentials map[string]any,
- expectedProxyID *int64,
- credentials map[string]any,
- ) (bool, error) {
- if r == nil || r.sql == nil {
- return false, errors.New("account repository SQL executor is not configured")
- }
- expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials))
- if err != nil {
- return false, err
- }
- credentialsJSON, err := json.Marshal(normalizeJSONMap(credentials))
- if err != nil {
- return false, err
- }
- result, err := r.sql.ExecContext(ctx, `
- WITH updated AS (
- UPDATE accounts AS a
- SET credentials = $1::jsonb,
- updated_at = NOW()
- WHERE a.id = $2
- AND a.deleted_at IS NULL
- AND a.platform = $3
- AND a.type = $4
- AND a.credentials = $5::jsonb
- AND a.proxy_id IS NOT DISTINCT FROM $6
- RETURNING a.id
- )
- INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
- SELECT $7, updated.id, NULL, NULL FROM updated
- `,
- string(credentialsJSON),
- id,
- service.PlatformGrok,
- service.AccountTypeOAuth,
- string(expectedJSON),
- expectedProxyID,
- service.SchedulerOutboxEventAccountChanged,
- )
- if err != nil {
- return false, err
- }
- rowsAffected, err := result.RowsAffected()
- if err != nil {
- return false, err
- }
- if rowsAffected == 0 {
- return false, nil
- }
- r.syncSchedulerAccountSnapshotDetached(ctx, id)
- return true, nil
- }
- // SetGrokOAuthRefreshErrorIfCredentialsUnchanged is the background-refresh
- // counterpart to reconciliation's stricter missing-refresh-token mutation. It
- // matches the complete credential document used by the failed upstream attempt
- // but deliberately does not require the refresh token to be absent.
- func (r *accountRepository) SetGrokOAuthRefreshErrorIfCredentialsUnchanged(
- ctx context.Context,
- id int64,
- expectedCredentials map[string]any,
- expectedProxyID *int64,
- errorMsg string,
- ) (bool, error) {
- if r == nil || r.sql == nil {
- return false, errors.New("account repository SQL executor is not configured")
- }
- expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials))
- if err != nil {
- return false, err
- }
- result, err := r.sql.ExecContext(ctx, `
- WITH updated AS (
- UPDATE accounts AS a
- SET status = $1,
- error_message = $2,
- schedulable = FALSE,
- updated_at = NOW()
- WHERE a.id = $3
- AND a.deleted_at IS NULL
- AND a.platform = $4
- AND a.type = $5
- AND a.status = $6
- AND a.credentials = $7::jsonb
- AND a.proxy_id IS NOT DISTINCT FROM $8
- RETURNING a.id
- )
- INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
- SELECT $9, updated.id, NULL, NULL FROM updated
- `,
- service.StatusError,
- errorMsg,
- id,
- service.PlatformGrok,
- service.AccountTypeOAuth,
- service.StatusActive,
- string(expectedJSON),
- expectedProxyID,
- service.SchedulerOutboxEventAccountChanged,
- )
- if err != nil {
- return false, err
- }
- rowsAffected, err := result.RowsAffected()
- if err != nil {
- return false, err
- }
- if rowsAffected == 0 {
- return false, nil
- }
- r.syncSchedulerAccountSnapshotDetached(ctx, id)
- return true, nil
- }
- // SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged applies a bounded
- // transient refresh quarantine only while the active Grok OAuth credential
- // document still matches the exact upstream attempt.
- func (r *accountRepository) SetGrokOAuthRefreshTempUnschedulableIfCredentialsUnchanged(
- ctx context.Context,
- id int64,
- expectedCredentials map[string]any,
- expectedProxyID *int64,
- until time.Time,
- reason string,
- ) (bool, error) {
- if r == nil || r.sql == nil {
- return false, errors.New("account repository SQL executor is not configured")
- }
- expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials))
- if err != nil {
- return false, err
- }
- result, err := r.sql.ExecContext(ctx, `
- WITH updated AS (
- UPDATE accounts AS a
- SET temp_unschedulable_until = $1,
- temp_unschedulable_reason = $2,
- updated_at = NOW()
- WHERE a.id = $3
- AND a.deleted_at IS NULL
- AND a.platform = $4
- AND a.type = $5
- AND a.status = $6
- AND a.credentials = $7::jsonb
- AND a.proxy_id IS NOT DISTINCT FROM $8
- AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until < $1)
- RETURNING a.id
- )
- INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
- SELECT $9, updated.id, NULL, NULL FROM updated
- `,
- until,
- reason,
- id,
- service.PlatformGrok,
- service.AccountTypeOAuth,
- service.StatusActive,
- string(expectedJSON),
- expectedProxyID,
- service.SchedulerOutboxEventAccountChanged,
- )
- if err != nil {
- return false, err
- }
- rowsAffected, err := result.RowsAffected()
- if err != nil {
- return false, err
- }
- if rowsAffected == 0 {
- return false, nil
- }
- r.syncSchedulerAccountSnapshotDetached(ctx, id)
- return true, nil
- }
- // syncSchedulerAccountSnapshot 在账号状态变更时主动同步快照到调度器缓存。
- // 当账号被设置为错误、禁用、不可调度或临时不可调度时调用,
- // 确保调度器和粘性会话逻辑能及时感知账号的最新状态,避免继续使用不可用账号。
- //
- // syncSchedulerAccountSnapshot proactively syncs account snapshot to scheduler cache
- // when account status changes. Called when account is set to error, disabled,
- // unschedulable, or temporarily unschedulable, ensuring scheduler and sticky session
- // logic can promptly detect the latest account state and avoid using unavailable accounts.
- func (r *accountRepository) syncSchedulerAccountSnapshot(ctx context.Context, accountID int64) {
- if r == nil || r.schedulerCache == nil || accountID <= 0 {
- return
- }
- account, err := r.GetByID(ctx, accountID)
- if err != nil {
- logger.LegacyPrintf("repository.account", "[Scheduler] sync account snapshot read failed: id=%d err=%v", accountID, err)
- return
- }
- if err := r.schedulerCache.SetAccount(ctx, account); err != nil {
- logger.LegacyPrintf("repository.account", "[Scheduler] sync account snapshot write failed: id=%d err=%v", accountID, err)
- }
- }
- func (r *accountRepository) syncSchedulerAccountSnapshotDetached(ctx context.Context, accountID int64) {
- base := context.Background()
- if ctx != nil {
- base = context.WithoutCancel(ctx)
- }
- propagationCtx, cancel := context.WithTimeout(base, 2*time.Second)
- defer cancel()
- r.syncSchedulerAccountSnapshot(propagationCtx, accountID)
- }
- func (r *accountRepository) deleteSchedulerAccountSnapshot(ctx context.Context, accountID int64) {
- if r == nil || r.schedulerCache == nil || accountID <= 0 {
- return
- }
- if err := r.schedulerCache.DeleteAccount(ctx, accountID); err != nil {
- logger.LegacyPrintf("repository.account", "[Scheduler] delete account snapshot failed: id=%d err=%v", accountID, err)
- }
- }
- func (r *accountRepository) syncSchedulerAccountSnapshots(ctx context.Context, accountIDs []int64) {
- if r == nil || r.schedulerCache == nil || len(accountIDs) == 0 {
- return
- }
- uniqueIDs := make([]int64, 0, len(accountIDs))
- seen := make(map[int64]struct{}, len(accountIDs))
- for _, id := range accountIDs {
- if id <= 0 {
- continue
- }
- if _, exists := seen[id]; exists {
- continue
- }
- seen[id] = struct{}{}
- uniqueIDs = append(uniqueIDs, id)
- }
- if len(uniqueIDs) == 0 {
- return
- }
- accounts, err := r.GetByIDs(ctx, uniqueIDs)
- if err != nil {
- logger.LegacyPrintf("repository.account", "[Scheduler] batch sync account snapshot read failed: count=%d err=%v", len(uniqueIDs), err)
- return
- }
- for _, account := range accounts {
- if account == nil {
- continue
- }
- if err := r.schedulerCache.SetAccount(ctx, account); err != nil {
- logger.LegacyPrintf("repository.account", "[Scheduler] batch sync account snapshot write failed: id=%d err=%v", account.ID, err)
- }
- }
- }
- func (r *accountRepository) ClearError(ctx context.Context, id int64) error {
- _, err := r.client.Account.Update().
- Where(dbaccount.IDEQ(id)).
- SetStatus(service.StatusActive).
- SetErrorMessage("").
- Save(ctx)
- if err != nil {
- return err
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue clear error failed: account=%d err=%v", id, err)
- }
- r.syncSchedulerAccountSnapshot(ctx, id)
- return nil
- }
- func (r *accountRepository) AddToGroup(ctx context.Context, accountID, groupID int64, priority int) error {
- tx, err := r.client.Tx(ctx)
- if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
- return err
- }
- client := r.client
- if tx != nil {
- defer func() { _ = tx.Rollback() }()
- client = tx.Client()
- }
- if err := lockLiveGroups(ctx, client, []int64{groupID}); err != nil {
- return err
- }
- _, err = client.AccountGroup.Create().
- SetAccountID(accountID).
- SetGroupID(groupID).
- SetPriority(priority).
- Save(ctx)
- if err != nil {
- return err
- }
- if tx != nil {
- if err := tx.Commit(); err != nil {
- return err
- }
- }
- payload := buildSchedulerGroupPayload([]int64{groupID})
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountGroupsChanged, &accountID, nil, payload); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue add to group failed: account=%d group=%d err=%v", accountID, groupID, err)
- }
- return nil
- }
- func (r *accountRepository) RemoveFromGroup(ctx context.Context, accountID, groupID int64) error {
- _, err := r.client.AccountGroup.Delete().
- Where(
- dbaccountgroup.AccountIDEQ(accountID),
- dbaccountgroup.GroupIDEQ(groupID),
- ).
- Exec(ctx)
- if err != nil {
- return err
- }
- payload := buildSchedulerGroupPayload([]int64{groupID})
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountGroupsChanged, &accountID, nil, payload); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue remove from group failed: account=%d group=%d err=%v", accountID, groupID, err)
- }
- return nil
- }
- func (r *accountRepository) GetGroups(ctx context.Context, accountID int64) ([]service.Group, error) {
- groups, err := r.client.Group.Query().
- Where(
- dbgroup.HasAccountsWith(dbaccount.IDEQ(accountID)),
- ).
- All(ctx)
- if err != nil {
- return nil, err
- }
- outGroups := make([]service.Group, 0, len(groups))
- for i := range groups {
- outGroups = append(outGroups, *groupEntityToService(groups[i]))
- }
- return outGroups, nil
- }
- func (r *accountRepository) BindGroups(ctx context.Context, accountID int64, groupIDs []int64) error {
- existingGroupIDs, err := r.loadAccountGroupIDs(ctx, accountID)
- if err != nil {
- return err
- }
- // 使用事务保证删除旧绑定与创建新绑定的原子性
- tx, err := r.client.Tx(ctx)
- if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
- return err
- }
- var txClient *dbent.Client
- if err == nil {
- defer func() { _ = tx.Rollback() }()
- txClient = tx.Client()
- } else {
- // 已处于外部事务中(ErrTxStarted),复用当前 client
- txClient = r.client
- }
- if err := lockLiveGroups(ctx, txClient, groupIDs); err != nil {
- return err
- }
- if _, err := txClient.AccountGroup.Delete().Where(dbaccountgroup.AccountIDEQ(accountID)).Exec(ctx); err != nil {
- return err
- }
- if len(groupIDs) == 0 {
- if tx != nil {
- return tx.Commit()
- }
- return nil
- }
- builders := make([]*dbent.AccountGroupCreate, 0, len(groupIDs))
- for i, groupID := range groupIDs {
- builders = append(builders, txClient.AccountGroup.Create().
- SetAccountID(accountID).
- SetGroupID(groupID).
- SetPriority(i+1),
- )
- }
- if _, err := txClient.AccountGroup.CreateBulk(builders...).Save(ctx); err != nil {
- return err
- }
- if tx != nil {
- if err := tx.Commit(); err != nil {
- return err
- }
- }
- payload := buildSchedulerGroupPayload(mergeGroupIDs(existingGroupIDs, groupIDs))
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountGroupsChanged, &accountID, nil, payload); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue bind groups failed: account=%d err=%v", accountID, err)
- }
- return nil
- }
- func (r *accountRepository) ListSchedulable(ctx context.Context) ([]service.Account, error) {
- accounts, err := r.schedulableAccountsQuery(time.Now()).All(ctx)
- if err != nil {
- return nil, err
- }
- return r.accountsToService(ctx, accounts)
- }
- func (r *accountRepository) ListSchedulableAccountLoads(ctx context.Context) ([]service.AccountWithConcurrency, error) {
- accounts, err := r.schedulableAccountsQuery(time.Now()).
- Select(
- dbaccount.FieldID,
- dbaccount.FieldConcurrency,
- dbaccount.FieldLoadFactor,
- ).
- All(ctx)
- if err != nil {
- return nil, err
- }
- loads := make([]service.AccountWithConcurrency, 0, len(accounts))
- for _, account := range accounts {
- projection := service.Account{
- ID: account.ID,
- Concurrency: account.Concurrency,
- LoadFactor: account.LoadFactor,
- }
- loads = append(loads, service.AccountWithConcurrency{
- ID: account.ID,
- MaxConcurrency: projection.EffectiveLoadFactor(),
- })
- }
- return loads, nil
- }
- func (r *accountRepository) schedulableAccountsQuery(now time.Time) *dbent.AccountQuery {
- return r.client.Account.Query().
- Where(
- dbaccount.StatusEQ(service.StatusActive),
- dbaccount.SchedulableEQ(true),
- tempUnschedulablePredicate(),
- notExpiredPredicate(now),
- dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
- dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
- ).
- Order(dbent.Asc(dbaccount.FieldPriority))
- }
- func (r *accountRepository) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]service.Account, error) {
- return r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{
- status: service.StatusActive,
- schedulable: true,
- })
- }
- func (r *accountRepository) ListSchedulableCapacityByGroupIDs(ctx context.Context, groupIDs []int64) ([]service.GroupAccountCapacityRow, error) {
- groupIDs = uniquePositiveInt64s(groupIDs)
- if len(groupIDs) == 0 {
- return []service.GroupAccountCapacityRow{}, nil
- }
- if r.sql == nil {
- rows := make([]service.GroupAccountCapacityRow, 0)
- for _, groupID := range groupIDs {
- accounts, err := r.ListSchedulableByGroupID(ctx, groupID)
- if err != nil {
- return nil, err
- }
- for i := range accounts {
- acc := &accounts[i]
- rows = append(rows, service.GroupAccountCapacityRow{
- GroupID: groupID,
- AccountID: acc.ID,
- Concurrency: acc.Concurrency,
- Extra: copyJSONMap(acc.Extra),
- SessionWindowStart: acc.SessionWindowStart,
- SessionWindowEnd: acc.SessionWindowEnd,
- SessionWindowStatus: acc.SessionWindowStatus,
- })
- }
- }
- return rows, nil
- }
- rows, err := r.sql.QueryContext(ctx, `
- SELECT
- ag.group_id,
- a.id AS account_id,
- a.concurrency,
- COALESCE(a.extra, '{}'::jsonb)::text AS extra,
- a.session_window_start,
- a.session_window_end,
- COALESCE(a.session_window_status, '') AS session_window_status
- FROM account_groups ag
- JOIN accounts a ON a.id = ag.account_id
- WHERE ag.group_id = ANY($1)
- AND a.deleted_at IS NULL
- AND a.status = $2
- AND a.schedulable = TRUE
- AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= $3)
- AND (a.expires_at IS NULL OR a.expires_at > $3 OR a.auto_pause_on_expired = FALSE)
- AND (a.overload_until IS NULL OR a.overload_until <= $3)
- AND (a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= $3)
- ORDER BY ag.group_id ASC, ag.priority ASC, a.priority ASC, a.id ASC
- `, pq.Array(groupIDs), service.StatusActive, time.Now())
- if err != nil {
- return nil, err
- }
- defer func() { _ = rows.Close() }()
- out := make([]service.GroupAccountCapacityRow, 0)
- for rows.Next() {
- var row service.GroupAccountCapacityRow
- var extraRaw string
- if err := rows.Scan(
- &row.GroupID,
- &row.AccountID,
- &row.Concurrency,
- &extraRaw,
- &row.SessionWindowStart,
- &row.SessionWindowEnd,
- &row.SessionWindowStatus,
- ); err != nil {
- return nil, err
- }
- if extraRaw != "" && extraRaw != "null" {
- var extra map[string]any
- if err := json.Unmarshal([]byte(extraRaw), &extra); err != nil {
- return nil, err
- }
- row.Extra = extra
- }
- out = append(out, row)
- }
- if err := rows.Err(); err != nil {
- return nil, err
- }
- return out, nil
- }
- func (r *accountRepository) ListSchedulableByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
- now := time.Now()
- accounts, err := r.client.Account.Query().
- Where(
- dbaccount.PlatformEQ(platform),
- dbaccount.StatusEQ(service.StatusActive),
- dbaccount.SchedulableEQ(true),
- tempUnschedulablePredicate(),
- notExpiredPredicate(now),
- dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
- dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
- ).
- Order(dbent.Asc(dbaccount.FieldPriority)).
- All(ctx)
- if err != nil {
- return nil, err
- }
- return r.accountsToService(ctx, accounts)
- }
- func (r *accountRepository) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]service.Account, error) {
- // 单平台查询复用多平台逻辑,保持过滤条件与排序策略一致。
- return r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{
- status: service.StatusActive,
- schedulable: true,
- platforms: []string{platform},
- })
- }
- func (r *accountRepository) ListSchedulableByPlatforms(ctx context.Context, platforms []string) ([]service.Account, error) {
- if len(platforms) == 0 {
- return nil, nil
- }
- // 仅返回可调度的活跃账号,并过滤处于过载/限流窗口的账号。
- // 代理与分组信息统一在 accountsToService 中批量加载,避免 N+1 查询。
- now := time.Now()
- accounts, err := r.client.Account.Query().
- Where(
- dbaccount.PlatformIn(platforms...),
- dbaccount.StatusEQ(service.StatusActive),
- dbaccount.SchedulableEQ(true),
- tempUnschedulablePredicate(),
- notExpiredPredicate(now),
- dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
- dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
- ).
- Order(dbent.Asc(dbaccount.FieldPriority)).
- All(ctx)
- if err != nil {
- return nil, err
- }
- return r.accountsToService(ctx, accounts)
- }
- func (r *accountRepository) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
- now := time.Now()
- accounts, err := r.client.Account.Query().
- Where(
- dbaccount.PlatformEQ(platform),
- dbaccount.StatusEQ(service.StatusActive),
- dbaccount.SchedulableEQ(true),
- dbaccount.Not(dbaccount.HasAccountGroups()),
- tempUnschedulablePredicate(),
- notExpiredPredicate(now),
- dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
- dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
- ).
- Order(dbent.Asc(dbaccount.FieldPriority)).
- All(ctx)
- if err != nil {
- return nil, err
- }
- return r.accountsToService(ctx, accounts)
- }
- func (r *accountRepository) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]service.Account, error) {
- if len(platforms) == 0 {
- return nil, nil
- }
- now := time.Now()
- accounts, err := r.client.Account.Query().
- Where(
- dbaccount.PlatformIn(platforms...),
- dbaccount.StatusEQ(service.StatusActive),
- dbaccount.SchedulableEQ(true),
- dbaccount.Not(dbaccount.HasAccountGroups()),
- tempUnschedulablePredicate(),
- notExpiredPredicate(now),
- dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
- dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
- ).
- Order(dbent.Asc(dbaccount.FieldPriority)).
- All(ctx)
- if err != nil {
- return nil, err
- }
- return r.accountsToService(ctx, accounts)
- }
- func (r *accountRepository) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]service.Account, error) {
- if len(platforms) == 0 {
- return nil, nil
- }
- // 复用按分组查询逻辑,保证分组优先级 + 账号优先级的排序与筛选一致。
- return r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{
- status: service.StatusActive,
- schedulable: true,
- platforms: platforms,
- })
- }
- // ListModelAvailabilityCandidates returns the persistently configured account
- // pool used to decide whether a model is supported. Unlike scheduling queries,
- // it intentionally ignores transient runtime state (rate limits, overload,
- // temporary unschedulability, and expiry windows).
- func (r *accountRepository) ListModelAvailabilityCandidates(
- ctx context.Context,
- groupID *int64,
- platforms []string,
- includeGrouped bool,
- ) ([]service.Account, error) {
- if len(platforms) == 0 {
- return []service.Account{}, nil
- }
- if groupID != nil {
- return r.queryAccountsByGroup(ctx, *groupID, accountGroupQueryOptions{
- status: service.StatusActive,
- schedulable: true,
- ignoreTransientState: true,
- platforms: platforms,
- })
- }
- preds := []dbpredicate.Account{
- dbaccount.StatusEQ(service.StatusActive),
- dbaccount.SchedulableEQ(true),
- dbaccount.PlatformIn(platforms...),
- }
- if !includeGrouped {
- preds = append(preds, dbaccount.Not(dbaccount.HasAccountGroups()))
- }
- accounts, err := r.client.Account.Query().
- Where(preds...).
- Order(dbent.Asc(dbaccount.FieldPriority)).
- All(ctx)
- if err != nil {
- return nil, err
- }
- return r.accountsToService(ctx, accounts)
- }
- func (r *accountRepository) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
- // Every upstream observation extends a limit. Only an explicit recovery
- // operation may clear it; late short responses must not undo a longer wait.
- return r.SetRateLimitedIfLater(ctx, id, resetAt)
- }
- // SetRateLimitedIfLater atomically extends an account-level rate limit. Grok
- // requests may finish concurrently, so an older response must not overwrite a
- // later reset boundary observed by another request or instance.
- func (r *accountRepository) SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error {
- now := time.Now()
- updated, err := r.client.Account.Update().
- Where(
- dbaccount.IDEQ(id),
- dbaccount.Or(
- dbaccount.RateLimitResetAtIsNil(),
- dbaccount.RateLimitResetAtLT(resetAt),
- ),
- ).
- SetRateLimitedAt(now).
- SetRateLimitResetAt(resetAt).
- Save(ctx)
- if err != nil {
- return err
- }
- if updated == 0 {
- // This instance may not have observed the later value written elsewhere.
- // Refresh its local scheduler snapshot even though no outbox event is needed.
- r.syncSchedulerAccountSnapshot(ctx, id)
- return nil
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue extended rate limit failed: account=%d err=%v", id, err)
- }
- r.syncSchedulerAccountSnapshot(ctx, id)
- return nil
- }
- // ClearRateLimitIfObserved clears exactly the Grok rate-limit generation seen
- // by a successful request. Matching both timestamps prevents a stale success
- // from erasing a later clear/re-arm generation with an equal or shorter reset.
- func (r *accountRepository) ClearRateLimitIfObserved(ctx context.Context, id int64, observedLimitedAt, observedResetAt time.Time) (bool, error) {
- updated, err := r.client.Account.Update().
- Where(
- dbaccount.IDEQ(id),
- dbaccount.PlatformEQ(service.PlatformGrok),
- dbaccount.TypeEQ(service.AccountTypeOAuth),
- dbaccount.RateLimitedAtEQ(observedLimitedAt),
- dbaccount.RateLimitResetAtEQ(observedResetAt),
- ).
- ClearRateLimitedAt().
- ClearRateLimitResetAt().
- Save(ctx)
- if err != nil {
- return false, err
- }
- if updated == 0 {
- r.syncSchedulerAccountSnapshot(ctx, id)
- return false, nil
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue observed rate-limit clear failed: account=%d err=%v", id, err)
- }
- r.syncSchedulerAccountSnapshot(ctx, id)
- return true, nil
- }
- func (r *accountRepository) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error {
- if scope == "" {
- return nil
- }
- now := time.Now().UTC()
- payload := map[string]string{
- "rate_limited_at": now.Format(time.RFC3339),
- "rate_limit_reset_at": resetAt.UTC().Format(time.RFC3339),
- }
- if len(reason) > 0 {
- if value := strings.TrimSpace(reason[0]); value != "" {
- payload["reason"] = value
- }
- }
- raw, err := json.Marshal(payload)
- if err != nil {
- return err
- }
- client := clientFromContext(ctx, r.client)
- result, err := client.ExecContext(
- ctx,
- `UPDATE accounts SET
- extra = jsonb_set(
- jsonb_set(COALESCE(extra, '{}'::jsonb), '{model_rate_limits}'::text[], COALESCE(extra->'model_rate_limits', '{}'::jsonb), true),
- ARRAY['model_rate_limits', $1]::text[],
- $2::jsonb,
- true
- ),
- updated_at = NOW()
- WHERE id = $3 AND deleted_at IS NULL`,
- scope,
- raw,
- id,
- )
- if err != nil {
- return err
- }
- affected, err := result.RowsAffected()
- if err != nil {
- return err
- }
- if affected == 0 {
- return service.ErrAccountNotFound
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue model rate limit failed: account=%d err=%v", id, err)
- }
- r.syncSchedulerAccountSnapshot(ctx, id)
- return nil
- }
- func (r *accountRepository) SetOverloaded(ctx context.Context, id int64, until time.Time) error {
- _, err := r.client.Account.Update().
- Where(dbaccount.IDEQ(id)).
- SetOverloadUntil(until).
- Save(ctx)
- if err != nil {
- return err
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue overload failed: account=%d err=%v", id, err)
- }
- r.syncSchedulerAccountSnapshot(ctx, id)
- return nil
- }
- func (r *accountRepository) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
- result, err := r.sql.ExecContext(ctx, `
- UPDATE accounts
- SET temp_unschedulable_until = $1,
- temp_unschedulable_reason = $2,
- updated_at = NOW()
- WHERE id = $3
- AND deleted_at IS NULL
- AND (temp_unschedulable_until IS NULL OR temp_unschedulable_until < $1)
- `, until, reason, id)
- if err != nil {
- return err
- }
- affected, err := result.RowsAffected()
- if err != nil {
- return err
- }
- if affected <= 0 {
- return nil
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue temp unschedulable failed: account=%d err=%v", id, err)
- }
- r.syncSchedulerAccountSnapshot(ctx, id)
- return nil
- }
- func (r *accountRepository) SetGrokCredentialTempUnschedulableIfMatch(
- ctx context.Context,
- id int64,
- snapshot service.GrokCredentialMutationSnapshot,
- until time.Time,
- reason string,
- ) (bool, error) {
- result, err := r.sql.ExecContext(ctx, `
- WITH updated AS (
- UPDATE accounts AS a
- SET temp_unschedulable_until = CASE
- WHEN a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until < $1 THEN $1
- ELSE a.temp_unschedulable_until
- END,
- temp_unschedulable_reason = $2,
- updated_at = NOW()
- WHERE a.id = $3
- AND a.deleted_at IS NULL
- AND a.status = $4
- AND a.platform = $5
- AND a.type = $6
- AND a.schedulable IS TRUE
- AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW())
- AND (a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW())
- AND (a.overload_until IS NULL OR a.overload_until <= NOW())
- AND (a.auto_pause_on_expired IS NOT TRUE OR a.expires_at IS NULL OR a.expires_at > NOW())
- AND a.credentials = $7::jsonb
- AND a.proxy_id IS NOT DISTINCT FROM $8
- RETURNING a.id
- )
- INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
- SELECT $9, updated.id, NULL, NULL FROM updated
- `, until, reason, id, service.StatusActive, service.PlatformGrok, service.AccountTypeOAuth,
- snapshot.CredentialsJSON, snapshot.ProxyID, service.SchedulerOutboxEventAccountChanged)
- if err != nil {
- return false, err
- }
- affected, err := result.RowsAffected()
- if err != nil || affected == 0 {
- return false, err
- }
- r.syncSchedulerAccountSnapshotDetached(ctx, id)
- return true, nil
- }
- func (r *accountRepository) ClearTempUnschedulable(ctx context.Context, id int64) error {
- _, err := r.sql.ExecContext(ctx, `
- UPDATE accounts
- SET temp_unschedulable_until = NULL,
- temp_unschedulable_reason = NULL,
- updated_at = NOW()
- WHERE id = $1
- AND deleted_at IS NULL
- `, id)
- if err != nil {
- return err
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue clear temp unschedulable failed: account=%d err=%v", id, err)
- }
- r.syncSchedulerAccountSnapshot(ctx, id)
- return nil
- }
- func (r *accountRepository) ClearRateLimit(ctx context.Context, id int64) error {
- _, err := r.client.Account.Update().
- Where(dbaccount.IDEQ(id)).
- ClearRateLimitedAt().
- ClearRateLimitResetAt().
- ClearOverloadUntil().
- Save(ctx)
- if err != nil {
- return err
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue clear rate limit failed: account=%d err=%v", id, err)
- }
- r.syncSchedulerAccountSnapshot(ctx, id)
- return nil
- }
- func (r *accountRepository) ClearAntigravityQuotaScopes(ctx context.Context, id int64) error {
- client := clientFromContext(ctx, r.client)
- result, err := client.ExecContext(
- ctx,
- "UPDATE accounts SET extra = COALESCE(extra, '{}'::jsonb) - 'antigravity_quota_scopes', updated_at = NOW() WHERE id = $1 AND deleted_at IS NULL",
- id,
- )
- if err != nil {
- return err
- }
- affected, err := result.RowsAffected()
- if err != nil {
- return err
- }
- if affected == 0 {
- return service.ErrAccountNotFound
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue clear quota scopes failed: account=%d err=%v", id, err)
- }
- return nil
- }
- func (r *accountRepository) ClearModelRateLimits(ctx context.Context, id int64) error {
- client := clientFromContext(ctx, r.client)
- result, err := client.ExecContext(
- ctx,
- "UPDATE accounts SET extra = COALESCE(extra, '{}'::jsonb) - 'model_rate_limits', updated_at = NOW() WHERE id = $1 AND deleted_at IS NULL",
- id,
- )
- if err != nil {
- return err
- }
- affected, err := result.RowsAffected()
- if err != nil {
- return err
- }
- if affected == 0 {
- return service.ErrAccountNotFound
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue clear model rate limit failed: account=%d err=%v", id, err)
- }
- r.syncSchedulerAccountSnapshot(ctx, id)
- return nil
- }
- func (r *accountRepository) UpdateSessionWindow(ctx context.Context, id int64, start, end *time.Time, status string) error {
- builder := r.client.Account.Update().
- Where(dbaccount.IDEQ(id)).
- SetSessionWindowStatus(status)
- if start != nil {
- builder.SetSessionWindowStart(*start)
- }
- if end != nil {
- builder.SetSessionWindowEnd(*end)
- }
- _, err := builder.Save(ctx)
- if err != nil {
- return err
- }
- // 触发调度器缓存更新(仅当窗口时间有变化时)
- if start != nil || end != nil {
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue session window update failed: account=%d err=%v", id, err)
- }
- }
- return nil
- }
- func (r *accountRepository) UpdateSessionWindowEnd(ctx context.Context, id int64, end time.Time) error {
- _, err := r.client.Account.Update().
- Where(dbaccount.IDEQ(id)).
- SetSessionWindowEnd(end).
- Save(ctx)
- if err != nil {
- return err
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue session window end update failed: account=%d err=%v", id, err)
- }
- return nil
- }
- func (r *accountRepository) SetSchedulable(ctx context.Context, id int64, schedulable bool) error {
- _, err := r.client.Account.Update().
- Where(dbaccount.IDEQ(id)).
- SetSchedulable(schedulable).
- Save(ctx)
- if err != nil {
- return err
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue schedulable change failed: account=%d err=%v", id, err)
- }
- if !schedulable {
- r.syncSchedulerAccountSnapshot(ctx, id)
- }
- return nil
- }
- func (r *accountRepository) AutoPauseExpiredAccounts(ctx context.Context, now time.Time) (int64, error) {
- rows, err := r.sql.QueryContext(ctx, `
- UPDATE accounts
- SET schedulable = FALSE,
- updated_at = NOW()
- WHERE deleted_at IS NULL
- AND schedulable = TRUE
- AND auto_pause_on_expired = TRUE
- AND expires_at IS NOT NULL
- AND expires_at <= $1
- RETURNING id
- `, now)
- if err != nil {
- return 0, err
- }
- defer func() {
- _ = rows.Close()
- }()
- accountIDs := make([]int64, 0)
- for rows.Next() {
- var accountID int64
- if err := rows.Scan(&accountID); err != nil {
- return 0, err
- }
- accountIDs = append(accountIDs, accountID)
- }
- if err := rows.Err(); err != nil {
- return 0, err
- }
- if len(accountIDs) > 0 {
- // 只刷新本次暂停的账号及其所属分组,避免少量账号到期触发所有调度桶重建。
- payload := map[string]any{"account_ids": accountIDs}
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue auto pause account changes failed: err=%v", err)
- }
- }
- return int64(len(accountIDs)), nil
- }
- func (r *accountRepository) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
- updates = stripCodexFingerprintSeedFromExtraUpdate(updates)
- if len(updates) == 0 {
- return nil
- }
- // 使用 JSONB 合并操作实现原子更新,避免读-改-写的并发丢失更新问题
- payload, err := json.Marshal(updates)
- if err != nil {
- return err
- }
- clearProbeSnapshot := upstreamBillingProbeExplicitlyDisabled(updates) || upstreamBillingProbeSnapshotClearRequested(updates)
- durableSchedulerChange := shouldEnqueueSchedulerOutboxForExtraUpdates(updates) || clearProbeSnapshot
- baseCtx := ctx
- contextTx := dbent.TxFromContext(ctx)
- client := clientFromContext(ctx, r.client)
- var tx *dbent.Tx
- if durableSchedulerChange && contextTx == nil {
- var txErr error
- tx, txErr = r.client.Tx(ctx)
- if txErr != nil && !errors.Is(txErr, dbent.ErrTxStarted) {
- return txErr
- }
- if tx != nil {
- defer func() { _ = tx.Rollback() }()
- ctx = dbent.NewTxContext(ctx, tx)
- client = tx.Client()
- }
- }
- extraExpression := "COALESCE(extra, '{}'::jsonb) || $1::jsonb"
- if clearProbeSnapshot {
- extraExpression = "(" + extraExpression + ") - 'upstream_billing_probe'"
- }
- if service.ShouldEnsureCodexFingerprintSeedForExtraUpdates(updates) {
- extraExpression = ensureCodexFingerprintSeedSQL(extraExpression)
- }
- extraExpression = preserveProtectionExtraSQL(ctx, extraExpression)
- result, err := client.ExecContext(
- ctx,
- "UPDATE accounts SET extra = "+extraExpression+", updated_at = NOW() WHERE id = $2 AND deleted_at IS NULL",
- string(payload), id,
- )
- if err != nil {
- return err
- }
- affected, err := result.RowsAffected()
- if err != nil {
- return err
- }
- if affected == 0 {
- return service.ErrAccountNotFound
- }
- if durableSchedulerChange {
- if err := enqueueSchedulerOutbox(ctx, client, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- return err
- }
- if tx != nil {
- if err := tx.Commit(); err != nil {
- return err
- }
- }
- if contextTx == nil {
- r.syncSchedulerAccountSnapshot(baseCtx, id)
- }
- } else {
- // 观测型 extra 字段不需要触发 bucket 重建,但仍同步单账号快照,
- // 让 sticky session / GetAccount 命中缓存时也能读到最新数据,
- // 同时避免缓存局部 patch 覆盖掉并发写入的其它账号字段。
- if dbent.TxFromContext(ctx) == nil {
- r.syncSchedulerAccountSnapshot(ctx, id)
- }
- }
- return nil
- }
- // UpdateUpstreamBillingProbeSnapshot stores a probe result only while the
- // network identity used by that probe is still current.
- func (r *accountRepository) UpdateUpstreamBillingProbeSnapshot(
- ctx context.Context,
- account *service.Account,
- snapshot *service.UpstreamBillingProbeSnapshot,
- rateMultiplier *float64,
- ) error {
- if account == nil || snapshot == nil {
- return service.ErrAccountNilInput
- }
- if snapshot.Status != service.UpstreamBillingProbeStatusOK {
- rateMultiplier = nil
- }
- if dbent.TxFromContext(ctx) == nil {
- tx, err := r.client.Tx(ctx)
- if errors.Is(err, dbent.ErrTxStarted) {
- return r.updateUpstreamBillingProbeSnapshotInTx(ctx, account, snapshot, rateMultiplier)
- }
- if err != nil {
- return err
- }
- defer func() { _ = tx.Rollback() }()
- if err := r.updateUpstreamBillingProbeSnapshotInTx(dbent.NewTxContext(ctx, tx), account, snapshot, rateMultiplier); err != nil {
- return err
- }
- if err := tx.Commit(); err != nil {
- return err
- }
- // The durable outbox event is committed with the snapshot. This direct
- // cache write only reduces visibility latency on the current instance.
- r.syncSchedulerAccountSnapshot(ctx, account.ID)
- return nil
- }
- return r.updateUpstreamBillingProbeSnapshotInTx(ctx, account, snapshot, rateMultiplier)
- }
- func (r *accountRepository) updateUpstreamBillingProbeSnapshotInTx(
- ctx context.Context,
- account *service.Account,
- snapshot *service.UpstreamBillingProbeSnapshot,
- rateMultiplier *float64,
- ) error {
- payload, err := json.Marshal(map[string]any{service.UpstreamBillingProbeExtraKey: snapshot})
- if err != nil {
- return err
- }
- credentials, err := json.Marshal(account.Credentials)
- if err != nil {
- return err
- }
- var expectedSnapshot any
- if account.Extra != nil {
- expectedSnapshot = account.Extra[service.UpstreamBillingProbeExtraKey]
- }
- expectedSnapshotJSON, err := json.Marshal(expectedSnapshot)
- if err != nil {
- return err
- }
- var expectedEnabled any
- if account.Extra != nil {
- expectedEnabled = account.Extra[service.UpstreamBillingProbeEnabledExtraKey]
- }
- expectedEnabledJSON, err := json.Marshal(expectedEnabled)
- if err != nil {
- return err
- }
- var expectedRateSyncEnabled any
- if account.Extra != nil {
- expectedRateSyncEnabled = account.Extra[service.UpstreamBillingRateSyncEnabledExtraKey]
- }
- expectedRateSyncEnabledJSON, err := json.Marshal(expectedRateSyncEnabled)
- if err != nil {
- return err
- }
- client := clientFromContext(ctx, r.client)
- proxyMatches, err := lockAndMatchProbeProxyIdentity(ctx, client, account)
- if err != nil {
- return err
- }
- if !proxyMatches {
- return service.ErrUpstreamBillingProbeIdentityChanged
- }
- var proxyID any
- if account.ProxyID != nil {
- proxyID = *account.ProxyID
- }
- result, err := client.ExecContext(ctx, `
- UPDATE accounts
- SET
- extra = COALESCE(extra, '{}'::jsonb) || $1::jsonb,
- rate_multiplier = CASE
- WHEN $10::numeric IS NOT NULL
- AND extra @> '{"upstream_billing_probe_enabled": true}'::jsonb
- AND extra @> '{"upstream_billing_rate_sync_enabled": true}'::jsonb
- THEN $10::numeric
- ELSE rate_multiplier
- END,
- updated_at = NOW()
- WHERE id = $2
- AND platform = $3
- AND type = $4
- AND credentials = $5::jsonb
- AND proxy_id IS NOT DISTINCT FROM $6
- AND COALESCE(extra -> 'upstream_billing_probe', 'null'::jsonb) = $7::jsonb
- AND COALESCE(extra -> 'upstream_billing_probe_enabled', 'null'::jsonb) = $8::jsonb
- AND COALESCE(extra -> 'upstream_billing_rate_sync_enabled', 'null'::jsonb) = $9::jsonb
- AND deleted_at IS NULL
- `, string(payload), account.ID, account.Platform, account.Type, string(credentials), proxyID, string(expectedSnapshotJSON), string(expectedEnabledJSON), string(expectedRateSyncEnabledJSON), rateMultiplier)
- if err != nil {
- return err
- }
- affected, err := result.RowsAffected()
- if err != nil {
- return err
- }
- if affected == 0 {
- return service.ErrUpstreamBillingProbeIdentityChanged
- }
- return enqueueSchedulerOutbox(ctx, client, service.SchedulerOutboxEventAccountChanged, &account.ID, nil, nil)
- }
- func lockAndMatchProbeProxyIdentity(ctx context.Context, client *dbent.Client, account *service.Account) (bool, error) {
- if account.ProxyID == nil {
- return true, nil
- }
- rows, err := client.QueryContext(ctx, `
- SELECT protocol, host, port, COALESCE(username, ''), COALESCE(password, ''), status
- FROM proxies
- WHERE id = $1 AND deleted_at IS NULL
- FOR SHARE
- `, *account.ProxyID)
- if err != nil {
- return false, err
- }
- defer func() { _ = rows.Close() }()
- if !rows.Next() {
- if err := rows.Err(); err != nil {
- return false, err
- }
- return account.Proxy == nil, nil
- }
- if account.Proxy == nil || account.Proxy.ID != *account.ProxyID {
- return false, nil
- }
- var current proxyProbeIdentity
- if err := rows.Scan(¤t.protocol, ¤t.host, ¤t.port, ¤t.username, ¤t.password, ¤t.status); err != nil {
- return false, err
- }
- return current == proxyProbeIdentityFromService(account.Proxy), rows.Err()
- }
- func shouldEnqueueSchedulerOutboxForExtraUpdates(updates map[string]any) bool {
- if len(updates) == 0 {
- return false
- }
- for key := range updates {
- if isSchedulerNeutralExtraKey(key) {
- continue
- }
- return true
- }
- return false
- }
- func isSchedulerNeutralExtraKey(key string) bool {
- key = strings.TrimSpace(key)
- if key == "" {
- return false
- }
- if _, ok := schedulerNeutralExtraKeys[key]; ok {
- return true
- }
- for _, prefix := range schedulerNeutralExtraKeyPrefixes {
- if strings.HasPrefix(key, prefix) {
- return true
- }
- }
- return false
- }
- func upstreamBillingProbeExplicitlyDisabled(extra map[string]any) bool {
- enabled, ok := extra[service.UpstreamBillingProbeEnabledExtraKey].(bool)
- return ok && !enabled
- }
- func upstreamBillingProbeSnapshotClearRequested(extra map[string]any) bool {
- value, ok := extra[service.UpstreamBillingProbeExtraKey]
- return ok && value == nil
- }
- func ollamaCloudUsageSnapshotClearRequested(extra map[string]any) bool {
- value, ok := extra[service.OllamaCloudUsageSnapshotExtraKey]
- return ok && value == nil
- }
- func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates service.AccountBulkUpdate) (int64, error) {
- if len(ids) == 0 {
- return 0, nil
- }
- updates.Extra = stripCodexFingerprintSeedFromExtraUpdate(updates.Extra)
- setClauses := make([]string, 0, 8)
- args := make([]any, 0, 8)
- idx := 1
- ollamaProxyIdentityChanged := ""
- concurrencyExpression := ""
- if updates.Name != nil {
- setClauses = append(setClauses, "name = $"+itoa(idx))
- args = append(args, *updates.Name)
- idx++
- }
- if updates.ProxyID != nil {
- // 0 表示清除代理(前端发送 0 而不是 null 来表达清除意图)
- if *updates.ProxyID == 0 {
- setClauses = append(setClauses, "proxy_id = NULL")
- ollamaProxyIdentityChanged = "proxy_id IS NOT NULL"
- } else {
- proxyPlaceholder := "$" + itoa(idx)
- setClauses = append(setClauses, "proxy_id = "+proxyPlaceholder)
- ollamaProxyIdentityChanged = "proxy_id IS DISTINCT FROM " + proxyPlaceholder
- args = append(args, *updates.ProxyID)
- idx++
- }
- }
- if updates.Concurrency != nil {
- concurrencyExpression = protectedConcurrencySQL("$" + itoa(idx) + "::integer")
- setClauses = append(setClauses, "concurrency = "+concurrencyExpression)
- args = append(args, *updates.Concurrency)
- idx++
- }
- if updates.Priority != nil {
- setClauses = append(setClauses, "priority = $"+itoa(idx))
- args = append(args, *updates.Priority)
- idx++
- }
- if updates.RateMultiplier != nil {
- setClauses = append(setClauses, "rate_multiplier = $"+itoa(idx))
- args = append(args, *updates.RateMultiplier)
- idx++
- }
- if updates.LoadFactor != nil {
- if *updates.LoadFactor <= 0 {
- setClauses = append(setClauses, "load_factor = NULL")
- } else {
- setClauses = append(setClauses, "load_factor = $"+itoa(idx))
- args = append(args, *updates.LoadFactor)
- idx++
- }
- }
- if updates.Status != nil {
- setClauses = append(setClauses, "status = $"+itoa(idx))
- args = append(args, *updates.Status)
- idx++
- }
- if updates.Schedulable != nil {
- setClauses = append(setClauses, "schedulable = $"+itoa(idx))
- args = append(args, *updates.Schedulable)
- idx++
- }
- if updates.ProbeEnabled != nil {
- if updates.Extra == nil {
- updates.Extra = make(map[string]any)
- }
- updates.Extra[service.UpstreamBillingProbeEnabledExtraKey] = *updates.ProbeEnabled
- }
- // JSONB 需要合并而非覆盖,使用 raw SQL 保持旧行为。
- credentialPlaceholder := ""
- if len(updates.Credentials) > 0 {
- payload, err := json.Marshal(updates.Credentials)
- if err != nil {
- return 0, err
- }
- credentialPlaceholder = "$" + itoa(idx)
- setClauses = append(setClauses, "credentials = COALESCE(credentials, '{}'::jsonb) || "+credentialPlaceholder+"::jsonb")
- args = append(args, payload)
- idx++
- }
- ollamaGroupIdentityChanges := make([]string, 0, 2)
- if _, ok := updates.Credentials["api_key"]; ok {
- ollamaGroupIdentityChanges = append(ollamaGroupIdentityChanges, "credentials -> 'api_key' IS DISTINCT FROM "+credentialPlaceholder+"::jsonb -> 'api_key'")
- }
- if _, ok := updates.Credentials["base_url"]; ok {
- ollamaGroupIdentityChanges = append(ollamaGroupIdentityChanges,
- "NOT ("+ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'")+
- " AND "+ollamaCloudBaseURLMatchesSQL(credentialPlaceholder+"::jsonb ->> 'base_url'")+")")
- }
- if len(updates.Extra) > 0 || len(ollamaGroupIdentityChanges) > 0 || ollamaProxyIdentityChanged != "" || updates.EnsureCodexFingerprintSeed || updates.Concurrency != nil {
- extraExpression := "COALESCE(extra, '{}'::jsonb)"
- if len(updates.Extra) > 0 {
- payload, err := json.Marshal(updates.Extra)
- if err != nil {
- return 0, err
- }
- extraExpression += " || $" + itoa(idx) + "::jsonb"
- args = append(args, payload)
- idx++
- if upstreamBillingProbeExplicitlyDisabled(updates.Extra) || upstreamBillingProbeSnapshotClearRequested(updates.Extra) {
- extraExpression = "(" + extraExpression + ") - 'upstream_billing_probe'"
- }
- if ollamaCloudUsageSnapshotClearRequested(updates.Extra) {
- extraExpression = "(" + extraExpression + ") - 'ollama_cloud_usage_snapshot'"
- }
- }
- eligibleAccount := "platform IN (" + ollamaCloudUsagePlatformsSQL + ") AND type = 'apikey'"
- groupIdentityChanged := ""
- if len(ollamaGroupIdentityChanges) > 0 {
- groupIdentityChanged = "(" + eligibleAccount + " AND (" + joinClauses(ollamaGroupIdentityChanges, " OR ") + "))"
- }
- snapshotIdentityChanged := groupIdentityChanged
- if ollamaProxyIdentityChanged != "" {
- proxyChanged := "(" + eligibleAccount + " AND " + ollamaProxyIdentityChanged + ")"
- if snapshotIdentityChanged == "" {
- snapshotIdentityChanged = proxyChanged
- } else {
- snapshotIdentityChanged = "(" + snapshotIdentityChanged + " OR " + proxyChanged + ")"
- }
- }
- if groupIdentityChanged != "" {
- extraExpression = "CASE" +
- " WHEN " + groupIdentityChanged + " THEN (" + extraExpression + ") - 'ollama_cloud_usage_session' - 'ollama_cloud_usage_auto_refresh' - 'ollama_cloud_usage_snapshot'" +
- " WHEN " + snapshotIdentityChanged + " THEN (" + extraExpression + ") - 'ollama_cloud_usage_snapshot'" +
- " ELSE " + extraExpression + " END"
- } else if snapshotIdentityChanged != "" {
- extraExpression = "CASE WHEN " + snapshotIdentityChanged + " THEN (" + extraExpression + ") - 'ollama_cloud_usage_snapshot' ELSE " + extraExpression + " END"
- }
- if updates.EnsureCodexFingerprintSeed {
- extraExpression = ensureCodexFingerprintSeedSQL(extraExpression)
- }
- extraExpression = preserveProtectionExtraSQL(ctx, extraExpression)
- if concurrencyExpression != "" {
- extraExpression = "CASE WHEN " + accountProtectionEnabledSQL + " THEN jsonb_set((" + extraExpression + "), '{anti_degrade,max_concurrency}', to_jsonb((" + concurrencyExpression + ")::integer), true) ELSE (" + extraExpression + ") END"
- }
- setClauses = append(setClauses, "extra = "+extraExpression)
- }
- if len(setClauses) == 0 {
- return 0, nil
- }
- setClauses = append(setClauses, "updated_at = NOW()")
- whereClause := " WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL"
- args = append(args, pq.Array(ids))
- idx++
- if updates.ProbeEnabled != nil {
- whereClause += " AND type = $" + itoa(idx)
- args = append(args, service.AccountTypeAPIKey)
- }
- query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + whereClause
- baseCtx := ctx
- contextTx := dbent.TxFromContext(ctx)
- exec := r.sql
- var tx *dbent.Tx
- if contextTx != nil {
- exec = contextTx.Client()
- } else if r.client != nil {
- var txErr error
- tx, txErr = r.client.Tx(ctx)
- if txErr != nil && !errors.Is(txErr, dbent.ErrTxStarted) {
- return 0, txErr
- }
- if tx != nil {
- defer func() { _ = tx.Rollback() }()
- ctx = dbent.NewTxContext(ctx, tx)
- exec = tx.Client()
- }
- }
- if mode, _ := updates.Extra[service.ProxyModeExtraKey].(string); !service.ProtectionManagedWrite(ctx) && strings.EqualFold(strings.TrimSpace(mode), service.ProxyModeRandom) {
- if dbent.TxFromContext(ctx) == nil {
- return 0, errors.New("random proxy bulk changes require a transaction")
- }
- if err := validateLockedBulkProxyMode(ctx, exec, ids, updates.Extra); err != nil {
- return 0, err
- }
- }
- result, err := exec.ExecContext(ctx, query, args...)
- if err != nil {
- return 0, err
- }
- rows, err := result.RowsAffected()
- if err != nil {
- return 0, err
- }
- if updates.ProbeEnabled != nil {
- expectedRows := int64(0)
- seenIDs := make(map[int64]struct{}, len(ids))
- for _, id := range ids {
- if _, seen := seenIDs[id]; seen {
- continue
- }
- seenIDs[id] = struct{}{}
- expectedRows++
- }
- if rows != expectedRows {
- return 0, service.ErrUpstreamBillingProbeAccountInvalid
- }
- }
- if rows > 0 {
- payload := map[string]any{"account_ids": ids}
- if err := enqueueSchedulerOutbox(ctx, exec, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil {
- return 0, err
- }
- }
- if tx != nil {
- if err := tx.Commit(); err != nil {
- return 0, err
- }
- }
- if rows > 0 && contextTx == nil {
- shouldSync := false
- if updates.Status != nil && (*updates.Status == service.StatusError || *updates.Status == service.StatusDisabled) {
- shouldSync = true
- }
- if updates.Schedulable != nil && !*updates.Schedulable {
- shouldSync = true
- }
- if shouldSync {
- r.syncSchedulerAccountSnapshots(baseCtx, ids)
- }
- }
- return rows, nil
- }
- type accountGroupQueryOptions struct {
- status string
- schedulable bool
- ignoreTransientState bool
- platforms []string // 允许的多个平台,空切片表示不进行平台过滤
- }
- func (r *accountRepository) queryAccountsByGroup(ctx context.Context, groupID int64, opts accountGroupQueryOptions) ([]service.Account, error) {
- q := r.client.AccountGroup.Query().
- Where(dbaccountgroup.GroupIDEQ(groupID))
- // 通过 account_groups 中间表查询账号,并按需叠加状态/平台/调度能力过滤。
- preds := make([]dbpredicate.Account, 0, 6)
- preds = append(preds, dbaccount.DeletedAtIsNil())
- if opts.status != "" {
- preds = append(preds, dbaccount.StatusEQ(opts.status))
- }
- if len(opts.platforms) > 0 {
- preds = append(preds, dbaccount.PlatformIn(opts.platforms...))
- }
- if opts.schedulable {
- preds = append(preds, dbaccount.SchedulableEQ(true))
- if !opts.ignoreTransientState {
- now := time.Now()
- preds = append(preds,
- tempUnschedulablePredicate(),
- notExpiredPredicate(now),
- dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
- dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
- )
- }
- }
- if len(preds) > 0 {
- q = q.Where(dbaccountgroup.HasAccountWith(preds...))
- }
- groups, err := q.
- Order(
- dbaccountgroup.ByPriority(),
- dbaccountgroup.ByAccountField(dbaccount.FieldPriority),
- ).
- WithAccount().
- All(ctx)
- if err != nil {
- return nil, err
- }
- orderedIDs := make([]int64, 0, len(groups))
- accountMap := make(map[int64]*dbent.Account, len(groups))
- for _, ag := range groups {
- if ag.Edges.Account == nil {
- continue
- }
- if _, exists := accountMap[ag.AccountID]; exists {
- continue
- }
- accountMap[ag.AccountID] = ag.Edges.Account
- orderedIDs = append(orderedIDs, ag.AccountID)
- }
- accounts := make([]*dbent.Account, 0, len(orderedIDs))
- for _, id := range orderedIDs {
- if acc, ok := accountMap[id]; ok {
- accounts = append(accounts, acc)
- }
- }
- return r.accountsToService(ctx, accounts)
- }
- func (r *accountRepository) accountsToService(ctx context.Context, accounts []*dbent.Account) ([]service.Account, error) {
- if len(accounts) == 0 {
- return []service.Account{}, nil
- }
- accountIDs := make([]int64, 0, len(accounts))
- proxyIDs := make([]int64, 0, len(accounts))
- for _, acc := range accounts {
- accountIDs = append(accountIDs, acc.ID)
- if acc.ProxyID != nil {
- proxyIDs = append(proxyIDs, *acc.ProxyID)
- }
- if acc.ProxyFallbackOriginID != nil {
- proxyIDs = append(proxyIDs, *acc.ProxyFallbackOriginID)
- }
- }
- proxyMap, err := r.loadProxies(ctx, proxyIDs)
- if err != nil {
- return nil, err
- }
- groupsByAccount, groupIDsByAccount, accountGroupsByAccount, err := r.loadAccountGroups(ctx, accountIDs)
- if err != nil {
- return nil, err
- }
- outAccounts := make([]service.Account, 0, len(accounts))
- for _, acc := range accounts {
- out := accountEntityToService(acc)
- if out == nil {
- continue
- }
- if acc.ProxyID != nil {
- if proxy, ok := proxyMap[*acc.ProxyID]; ok {
- out.Proxy = proxy
- }
- }
- out.ProxyFallbackOriginID = acc.ProxyFallbackOriginID
- if acc.ProxyFallbackOriginID != nil {
- if op, ok := proxyMap[*acc.ProxyFallbackOriginID]; ok && op != nil {
- n := op.Name
- out.ProxyFallbackOriginName = &n
- }
- }
- if groups, ok := groupsByAccount[acc.ID]; ok {
- out.Groups = groups
- }
- if groupIDs, ok := groupIDsByAccount[acc.ID]; ok {
- out.GroupIDs = groupIDs
- }
- if ags, ok := accountGroupsByAccount[acc.ID]; ok {
- out.AccountGroups = ags
- }
- outAccounts = append(outAccounts, *out)
- }
- return outAccounts, nil
- }
- func tempUnschedulablePredicate() dbpredicate.Account {
- return dbpredicate.Account(func(s *entsql.Selector) {
- col := s.C("temp_unschedulable_until")
- s.Where(entsql.Or(
- entsql.IsNull(col),
- entsql.LTE(col, entsql.Expr("NOW()")),
- ))
- })
- }
- func notExpiredPredicate(now time.Time) dbpredicate.Account {
- return dbaccount.Or(
- dbaccount.ExpiresAtIsNil(),
- dbaccount.ExpiresAtGT(now),
- dbaccount.AutoPauseOnExpiredEQ(false),
- )
- }
- func (r *accountRepository) loadProxies(ctx context.Context, proxyIDs []int64) (map[int64]*service.Proxy, error) {
- proxyMap := make(map[int64]*service.Proxy)
- proxyIDs = uniquePositiveInt64s(proxyIDs)
- if len(proxyIDs) == 0 {
- return proxyMap, nil
- }
- for start := 0; start < len(proxyIDs); start += postgresParameterBatchSize {
- end := start + postgresParameterBatchSize
- if end > len(proxyIDs) {
- end = len(proxyIDs)
- }
- proxies, err := r.client.Proxy.Query().Where(dbproxy.IDIn(proxyIDs[start:end]...)).All(ctx)
- if err != nil {
- return nil, err
- }
- for _, p := range proxies {
- proxyMap[p.ID] = proxyEntityToService(p)
- }
- }
- return proxyMap, nil
- }
- func (r *accountRepository) loadAccountGroups(ctx context.Context, accountIDs []int64) (map[int64][]*service.Group, map[int64][]int64, map[int64][]service.AccountGroup, error) {
- groupsByAccount := make(map[int64][]*service.Group)
- groupIDsByAccount := make(map[int64][]int64)
- accountGroupsByAccount := make(map[int64][]service.AccountGroup)
- accountIDs = uniquePositiveInt64s(accountIDs)
- if len(accountIDs) == 0 {
- return groupsByAccount, groupIDsByAccount, accountGroupsByAccount, nil
- }
- for start := 0; start < len(accountIDs); start += postgresParameterBatchSize {
- end := start + postgresParameterBatchSize
- if end > len(accountIDs) {
- end = len(accountIDs)
- }
- entries, err := r.client.AccountGroup.Query().
- Where(dbaccountgroup.AccountIDIn(accountIDs[start:end]...)).
- Order(dbaccountgroup.ByAccountID(), dbaccountgroup.ByPriority()).
- All(ctx)
- if err != nil {
- return nil, nil, nil, err
- }
- groupIDs := make([]int64, 0, len(entries))
- for _, ag := range entries {
- groupIDs = append(groupIDs, ag.GroupID)
- }
- groupMap, err := r.loadGroups(ctx, groupIDs)
- if err != nil {
- return nil, nil, nil, err
- }
- for _, ag := range entries {
- groupSvc := groupMap[ag.GroupID]
- agSvc := service.AccountGroup{
- AccountID: ag.AccountID,
- GroupID: ag.GroupID,
- Priority: ag.Priority,
- CreatedAt: ag.CreatedAt,
- Group: groupSvc,
- }
- accountGroupsByAccount[ag.AccountID] = append(accountGroupsByAccount[ag.AccountID], agSvc)
- groupIDsByAccount[ag.AccountID] = append(groupIDsByAccount[ag.AccountID], ag.GroupID)
- if groupSvc != nil {
- groupsByAccount[ag.AccountID] = append(groupsByAccount[ag.AccountID], groupSvc)
- }
- }
- }
- return groupsByAccount, groupIDsByAccount, accountGroupsByAccount, nil
- }
- func (r *accountRepository) loadGroups(ctx context.Context, groupIDs []int64) (map[int64]*service.Group, error) {
- groupMap := make(map[int64]*service.Group)
- groupIDs = uniquePositiveInt64s(groupIDs)
- if len(groupIDs) == 0 {
- return groupMap, nil
- }
- for start := 0; start < len(groupIDs); start += postgresParameterBatchSize {
- end := start + postgresParameterBatchSize
- if end > len(groupIDs) {
- end = len(groupIDs)
- }
- groups, err := r.client.Group.Query().Where(dbgroup.IDIn(groupIDs[start:end]...)).All(ctx)
- if err != nil {
- return nil, err
- }
- for _, g := range groups {
- groupMap[g.ID] = groupEntityToService(g)
- }
- }
- return groupMap, nil
- }
- func uniquePositiveInt64s(ids []int64) []int64 {
- if len(ids) == 0 {
- return nil
- }
- out := make([]int64, 0, len(ids))
- seen := make(map[int64]struct{}, len(ids))
- for _, id := range ids {
- if id <= 0 {
- continue
- }
- if _, ok := seen[id]; ok {
- continue
- }
- seen[id] = struct{}{}
- out = append(out, id)
- }
- return out
- }
- func (r *accountRepository) loadAccountGroupIDs(ctx context.Context, accountID int64) ([]int64, error) {
- entries, err := r.client.AccountGroup.
- Query().
- Where(dbaccountgroup.AccountIDEQ(accountID)).
- All(ctx)
- if err != nil {
- return nil, err
- }
- ids := make([]int64, 0, len(entries))
- for _, entry := range entries {
- ids = append(ids, entry.GroupID)
- }
- return ids, nil
- }
- func mergeGroupIDs(a []int64, b []int64) []int64 {
- seen := make(map[int64]struct{}, len(a)+len(b))
- out := make([]int64, 0, len(a)+len(b))
- for _, id := range a {
- if id <= 0 {
- continue
- }
- if _, ok := seen[id]; ok {
- continue
- }
- seen[id] = struct{}{}
- out = append(out, id)
- }
- for _, id := range b {
- if id <= 0 {
- continue
- }
- if _, ok := seen[id]; ok {
- continue
- }
- seen[id] = struct{}{}
- out = append(out, id)
- }
- return out
- }
- // buildSchedulerGroupPayload 构造 EventAccountChanged / EventAccountGroupsChanged
- // 事件的 payload。空 groupIDs 必须返回 untyped nil(any 而非 map[string]any(nil)),
- // 否则 enqueueSchedulerOutbox 的 "payload != nil" 接口判空会被 typed-nil 欺骗,
- // 把 payload marshal 成 "null" 写入 dedup_key 哈希,破坏与其他 nil-payload 调用的去重一致性。
- func buildSchedulerGroupPayload(groupIDs []int64) any {
- if len(groupIDs) == 0 {
- return nil
- }
- return map[string]any{"group_ids": groupIDs}
- }
- func accountEntityToService(m *dbent.Account) *service.Account {
- if m == nil {
- return nil
- }
- rateMultiplier := m.RateMultiplier
- return &service.Account{
- ID: m.ID,
- Name: m.Name,
- Notes: m.Notes,
- Platform: m.Platform,
- Type: m.Type,
- Credentials: copyJSONMap(m.Credentials),
- Extra: copyJSONMap(m.Extra),
- ProxyID: m.ProxyID,
- ProxyFallbackOriginID: m.ProxyFallbackOriginID,
- Concurrency: m.Concurrency,
- Priority: m.Priority,
- RateMultiplier: &rateMultiplier,
- LoadFactor: m.LoadFactor,
- Status: m.Status,
- ErrorMessage: derefString(m.ErrorMessage),
- LastUsedAt: m.LastUsedAt,
- ExpiresAt: m.ExpiresAt,
- AutoPauseOnExpired: m.AutoPauseOnExpired,
- CreatedAt: m.CreatedAt,
- UpdatedAt: m.UpdatedAt,
- Schedulable: m.Schedulable,
- RateLimitedAt: m.RateLimitedAt,
- RateLimitResetAt: m.RateLimitResetAt,
- OverloadUntil: m.OverloadUntil,
- TempUnschedulableUntil: m.TempUnschedulableUntil,
- TempUnschedulableReason: derefString(m.TempUnschedulableReason),
- SessionWindowStart: m.SessionWindowStart,
- SessionWindowEnd: m.SessionWindowEnd,
- SessionWindowStatus: derefString(m.SessionWindowStatus),
- ParentAccountID: m.ParentAccountID,
- QuotaDimension: string(m.QuotaDimension),
- }
- }
- func normalizeJSONMap(in map[string]any) map[string]any {
- if in == nil {
- return map[string]any{}
- }
- return in
- }
- func copyJSONMap(in map[string]any) map[string]any {
- if in == nil {
- return nil
- }
- out := make(map[string]any, len(in))
- for k, v := range in {
- out[k] = v
- }
- return out
- }
- func joinClauses(clauses []string, sep string) string {
- if len(clauses) == 0 {
- return ""
- }
- out := clauses[0]
- for i := 1; i < len(clauses); i++ {
- out += sep + clauses[i]
- }
- return out
- }
- func itoa(v int) string {
- return strconv.Itoa(v)
- }
- // FindByExtraField 根据 extra 字段中的键值对查找账号。
- // 使用 PostgreSQL JSONB @> 操作符进行高效查询(需要 GIN 索引支持)。
- //
- // FindByExtraField finds accounts by key-value pairs in the extra field.
- // Uses PostgreSQL JSONB @> operator for efficient queries (requires GIN index).
- func (r *accountRepository) FindByExtraField(ctx context.Context, key string, value any) ([]service.Account, error) {
- accounts, err := r.client.Account.Query().
- Where(
- dbaccount.DeletedAtIsNil(),
- func(s *entsql.Selector) {
- path := sqljson.Path(key)
- switch v := value.(type) {
- case string:
- preds := []*entsql.Predicate{sqljson.ValueEQ(dbaccount.FieldExtra, v, path)}
- if parsed, err := strconv.ParseInt(v, 10, 64); err == nil {
- preds = append(preds, sqljson.ValueEQ(dbaccount.FieldExtra, parsed, path))
- }
- if len(preds) == 1 {
- s.Where(preds[0])
- } else {
- s.Where(entsql.Or(preds...))
- }
- case int:
- s.Where(entsql.Or(
- sqljson.ValueEQ(dbaccount.FieldExtra, v, path),
- sqljson.ValueEQ(dbaccount.FieldExtra, strconv.Itoa(v), path),
- ))
- case int64:
- s.Where(entsql.Or(
- sqljson.ValueEQ(dbaccount.FieldExtra, v, path),
- sqljson.ValueEQ(dbaccount.FieldExtra, strconv.FormatInt(v, 10), path),
- ))
- case json.Number:
- if parsed, err := v.Int64(); err == nil {
- s.Where(entsql.Or(
- sqljson.ValueEQ(dbaccount.FieldExtra, parsed, path),
- sqljson.ValueEQ(dbaccount.FieldExtra, v.String(), path),
- ))
- } else {
- s.Where(sqljson.ValueEQ(dbaccount.FieldExtra, v.String(), path))
- }
- default:
- s.Where(sqljson.ValueEQ(dbaccount.FieldExtra, value, path))
- }
- },
- ).
- All(ctx)
- if err != nil {
- return nil, translatePersistenceError(err, service.ErrAccountNotFound, nil)
- }
- return r.accountsToService(ctx, accounts)
- }
- // ListDueUpstreamBillingProbeAccounts bounds result hydration and network work
- // to limit. PostgreSQL must still filter and order all enabled candidates;
- // MATERIALIZED avoids repeating the defensive timestamp parse expression.
- // Go writes next_probe_at via RFC3339Nano (up to 9 fractional digits) while
- // jsonpath datetime() parses at most microseconds, so fractions beyond 6
- // digits are trimmed first — mirroring ListDueOllamaCloudUsageAccounts.
- // Without this, every nanosecond timestamp is treated as malformed and the
- // fail-open ordering pins the cycle to the lowest account IDs, starving the
- // rest of the pool.
- func (r *accountRepository) ListDueUpstreamBillingProbeAccounts(ctx context.Context, now time.Time, limit int) ([]service.Account, error) {
- if limit <= 0 {
- return []service.Account{}, nil
- }
- if r.sql == nil {
- return nil, errors.New("account repository SQL executor not configured")
- }
- rows, err := r.sql.QueryContext(ctx, `
- WITH candidates AS (
- SELECT
- id,
- extra #>> '{upstream_billing_probe,status}' AS probe_status,
- extra #>> '{upstream_billing_probe,next_probe_at}' AS next_probe_at
- FROM accounts
- WHERE deleted_at IS NULL
- AND status = 'active'
- AND type = 'apikey'
- AND extra @> '{"upstream_billing_probe_enabled": true}'::jsonb
- ), parsed AS MATERIALIZED (
- SELECT
- id,
- probe_status,
- next_probe_at,
- 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,
- jsonb_path_query_first_tz(
- jsonb_build_object(
- 'value',
- replace(regexp_replace(regexp_replace(
- next_probe_at,
- '(\.[0-9]{6})[0-9]+(Z|[+-][0-9]{2}:[0-9]{2})$',
- '\1\2'
- ), 'Z$', '+00:00'), 'T', ' ')
- ),
- '$.value.datetime()',
- '{}'::jsonb,
- true
- ) #>> '{}' AS parsed_next_probe_at
- FROM candidates
- ), normalized AS (
- SELECT
- id,
- probe_status,
- next_probe_at,
- parsed_next_probe_at,
- rfc3339_shape AND parsed_next_probe_at IS NOT NULL AS valid_next_probe_at
- FROM parsed
- )
- SELECT id
- FROM normalized
- WHERE probe_status NOT IN ('ok', 'unsupported', 'failed')
- OR probe_status IS NULL
- OR next_probe_at IS NULL
- OR NOT valid_next_probe_at
- OR CASE WHEN valid_next_probe_at THEN parsed_next_probe_at::timestamptz <= $1 ELSE FALSE END
- ORDER BY
- CASE
- WHEN probe_status NOT IN ('ok', 'unsupported', 'failed')
- OR probe_status IS NULL
- OR next_probe_at IS NULL
- OR NOT valid_next_probe_at
- THEN 0
- ELSE 1
- END ASC,
- CASE WHEN valid_next_probe_at THEN parsed_next_probe_at::timestamptz END ASC NULLS FIRST,
- id ASC
- LIMIT $2
- `, now.UTC(), limit)
- if err != nil {
- return nil, err
- }
- defer func() { _ = rows.Close() }()
- ids := make([]int64, 0, limit)
- for rows.Next() {
- var id int64
- if err := rows.Scan(&id); err != nil {
- return nil, err
- }
- ids = append(ids, id)
- }
- if err := rows.Err(); err != nil {
- return nil, err
- }
- if len(ids) == 0 {
- return []service.Account{}, nil
- }
- accounts, err := r.GetByIDs(ctx, ids)
- if err != nil {
- return nil, err
- }
- out := make([]service.Account, 0, len(accounts))
- for _, account := range accounts {
- if account != nil {
- out = append(out, *account)
- }
- }
- return out, nil
- }
- // nowUTC is a SQL expression to generate a UTC RFC3339 timestamp string.
- const nowUTC = `to_char(NOW() AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.US"Z"')`
- // dailyExpiredExpr is a SQL expression that evaluates to TRUE when daily quota period has expired.
- // Supports both rolling (24h from start) and fixed (pre-computed reset_at) modes.
- const dailyExpiredExpr = `(
- CASE WHEN COALESCE(extra->>'quota_daily_reset_mode', 'rolling') = 'fixed'
- THEN NOW() >= COALESCE((extra->>'quota_daily_reset_at')::timestamptz, '1970-01-01'::timestamptz)
- ELSE COALESCE((extra->>'quota_daily_start')::timestamptz, '1970-01-01'::timestamptz)
- + '24 hours'::interval <= NOW()
- END
- )`
- // weeklyExpiredExpr is a SQL expression that evaluates to TRUE when weekly quota period has expired.
- const weeklyExpiredExpr = `(
- CASE WHEN COALESCE(extra->>'quota_weekly_reset_mode', 'rolling') = 'fixed'
- THEN NOW() >= COALESCE((extra->>'quota_weekly_reset_at')::timestamptz, '1970-01-01'::timestamptz)
- ELSE COALESCE((extra->>'quota_weekly_start')::timestamptz, '1970-01-01'::timestamptz)
- + '168 hours'::interval <= NOW()
- END
- )`
- // nextDailyResetAtExpr is a SQL expression to compute the next daily reset_at when a reset occurs.
- // For fixed mode: computes the next future reset time based on NOW(), timezone, and configured hour.
- // This correctly handles long-inactive accounts by jumping directly to the next valid reset point.
- const nextDailyResetAtExpr = `(
- CASE WHEN COALESCE(extra->>'quota_daily_reset_mode', 'rolling') = 'fixed'
- THEN to_char((
- -- Compute today's reset point in the configured timezone, then pick next future one
- CASE WHEN NOW() >= (
- date_trunc('day', NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))
- + (COALESCE((extra->>'quota_daily_reset_hour')::int, 0) || ' hours')::interval
- ) AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC')
- -- NOW() is at or past today's reset point → next reset is tomorrow
- THEN (
- date_trunc('day', NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))
- + (COALESCE((extra->>'quota_daily_reset_hour')::int, 0) || ' hours')::interval
- + '1 day'::interval
- ) AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC')
- -- NOW() is before today's reset point → next reset is today
- ELSE (
- date_trunc('day', NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))
- + (COALESCE((extra->>'quota_daily_reset_hour')::int, 0) || ' hours')::interval
- ) AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC')
- END
- ) AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS"Z"')
- ELSE NULL END
- )`
- // nextWeeklyResetAtExpr is a SQL expression to compute the next weekly reset_at when a reset occurs.
- // For fixed mode: computes the next future reset time based on NOW(), timezone, configured day and hour.
- // This correctly handles long-inactive accounts by jumping directly to the next valid reset point.
- const nextWeeklyResetAtExpr = `(
- CASE WHEN COALESCE(extra->>'quota_weekly_reset_mode', 'rolling') = 'fixed'
- THEN to_char((
- -- Compute this week's reset point in the configured timezone
- -- Step 1: get today's date at reset hour in configured tz
- -- Step 2: compute days forward to target weekday
- -- Step 3: if same day but past reset hour, advance 7 days
- CASE
- WHEN (
- -- days_forward = (target_day - current_day + 7) % 7
- (COALESCE((extra->>'quota_weekly_reset_day')::int, 1)
- - EXTRACT(DOW FROM NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))::int
- + 7) % 7
- ) = 0 AND NOW() >= (
- date_trunc('day', NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))
- + (COALESCE((extra->>'quota_weekly_reset_hour')::int, 0) || ' hours')::interval
- ) AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC')
- -- Same weekday and past reset hour → next week
- THEN (
- date_trunc('day', NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))
- + (COALESCE((extra->>'quota_weekly_reset_hour')::int, 0) || ' hours')::interval
- + '7 days'::interval
- ) AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC')
- ELSE (
- -- Advance to target weekday this week (or next if days_forward > 0)
- date_trunc('day', NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))
- + (COALESCE((extra->>'quota_weekly_reset_hour')::int, 0) || ' hours')::interval
- + ((
- (COALESCE((extra->>'quota_weekly_reset_day')::int, 1)
- - EXTRACT(DOW FROM NOW() AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC'))::int
- + 7) % 7
- ) || ' days')::interval
- ) AT TIME ZONE COALESCE(extra->>'quota_reset_timezone', 'UTC')
- END
- ) AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS"Z"')
- ELSE NULL END
- )`
- // IncrementQuotaUsed 原子递增账号的配额用量(总/日/周三个维度)
- // 日/周额度在周期过期时自动重置为 0 再递增。
- // 支持滚动窗口(rolling)和固定时间(fixed)两种重置模式。
- func (r *accountRepository) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) error {
- rows, err := r.sql.QueryContext(ctx,
- `UPDATE accounts SET extra = (
- COALESCE(extra, '{}'::jsonb)
- -- 总额度:始终递增
- || jsonb_build_object('quota_used', COALESCE((extra->>'quota_used')::numeric, 0) + $1)
- -- 日额度:仅在 quota_daily_limit > 0 时处理
- || CASE WHEN COALESCE((extra->>'quota_daily_limit')::numeric, 0) > 0 THEN
- jsonb_build_object(
- 'quota_daily_used',
- CASE WHEN `+dailyExpiredExpr+`
- THEN $1
- ELSE COALESCE((extra->>'quota_daily_used')::numeric, 0) + $1 END,
- 'quota_daily_start',
- CASE WHEN `+dailyExpiredExpr+`
- THEN `+nowUTC+`
- ELSE COALESCE(extra->>'quota_daily_start', `+nowUTC+`) END
- )
- -- 固定模式重置时更新下次重置时间
- || CASE WHEN `+dailyExpiredExpr+` AND `+nextDailyResetAtExpr+` IS NOT NULL
- THEN jsonb_build_object('quota_daily_reset_at', `+nextDailyResetAtExpr+`)
- ELSE '{}'::jsonb END
- ELSE '{}'::jsonb END
- -- 周额度:仅在 quota_weekly_limit > 0 时处理
- || CASE WHEN COALESCE((extra->>'quota_weekly_limit')::numeric, 0) > 0 THEN
- jsonb_build_object(
- 'quota_weekly_used',
- CASE WHEN `+weeklyExpiredExpr+`
- THEN $1
- ELSE COALESCE((extra->>'quota_weekly_used')::numeric, 0) + $1 END,
- 'quota_weekly_start',
- CASE WHEN `+weeklyExpiredExpr+`
- THEN `+nowUTC+`
- ELSE COALESCE(extra->>'quota_weekly_start', `+nowUTC+`) END
- )
- -- 固定模式重置时更新下次重置时间
- || CASE WHEN `+weeklyExpiredExpr+` AND `+nextWeeklyResetAtExpr+` IS NOT NULL
- THEN jsonb_build_object('quota_weekly_reset_at', `+nextWeeklyResetAtExpr+`)
- ELSE '{}'::jsonb END
- ELSE '{}'::jsonb END
- ), updated_at = NOW()
- WHERE id = $2 AND deleted_at IS NULL
- RETURNING
- COALESCE((extra->>'quota_used')::numeric, 0),
- COALESCE((extra->>'quota_limit')::numeric, 0)`,
- amount, id)
- if err != nil {
- return err
- }
- defer func() { _ = rows.Close() }()
- var newUsed, limit float64
- if rows.Next() {
- if err := rows.Scan(&newUsed, &limit); err != nil {
- return err
- }
- }
- if err := rows.Err(); err != nil {
- return err
- }
- // 任一维度配额刚超限时触发调度快照刷新
- if limit > 0 && newUsed >= limit && (newUsed-amount) < limit {
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue quota exceeded failed: account=%d err=%v", id, err)
- }
- }
- return nil
- }
- // ResetQuotaUsedAndClearRateLimitCooldown resets all quota dimensions and the
- // account-level cooldown in one statement. Other scheduler blocking state is preserved.
- func (r *accountRepository) ResetQuotaUsedAndClearRateLimitCooldown(ctx context.Context, id int64) error {
- result, err := r.sql.ExecContext(ctx,
- `UPDATE accounts SET extra = (
- COALESCE(extra, '{}'::jsonb)
- || '{"quota_used": 0, "quota_daily_used": 0, "quota_weekly_used": 0}'::jsonb
- ) - 'quota_daily_start' - 'quota_weekly_start' - 'quota_daily_reset_at' - 'quota_weekly_reset_at',
- rate_limited_at = NULL, rate_limit_reset_at = NULL, updated_at = NOW()
- WHERE id = $1 AND deleted_at IS NULL`,
- id)
- if err != nil {
- return err
- }
- affected, err := result.RowsAffected()
- if err != nil {
- return err
- }
- if affected == 0 {
- return service.ErrAccountNotFound
- }
- // 重置配额后触发调度快照刷新,使账号重新参与调度
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue quota reset failed: account=%d err=%v", id, err)
- }
- r.syncSchedulerAccountSnapshot(ctx, id)
- return nil
- }
- // RevertProxyFallback 将账号的 proxy_id 切回 proxy_fallback_origin_id,并清空 origin 字段。
- // 仅当 proxy_fallback_origin_id IS NOT NULL 时执行更新;
- // 若影响行数为 0,则返回 ErrAccountNotInFallback(账号存在但不在 fallback 状态)。
- func (r *accountRepository) RevertProxyFallback(ctx context.Context, accountID int64) error {
- res, err := r.sql.ExecContext(ctx, `
- UPDATE accounts SET proxy_id=proxy_fallback_origin_id, proxy_fallback_origin_id=NULL, updated_at=NOW()
- WHERE id=$1 AND proxy_fallback_origin_id IS NOT NULL AND deleted_at IS NULL`, accountID)
- if err != nil {
- return err
- }
- n, _ := res.RowsAffected()
- if n == 0 {
- return service.ErrAccountNotInFallback
- }
- if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &accountID, nil, nil); err != nil {
- logger.LegacyPrintf("repository.account", "[SchedulerOutbox] revert fallback enqueue failed: account=%d err=%v", accountID, err)
- }
- return nil
- }
- // ListShadowsByParent 返回指定父账号的影子账号;当前实现仅查 quota_dimension='spark'(唯一预设)。
- // 同时过滤 parent_account_id 和 quota_dimension='spark',防止未来其它 linked 维度被误伤。
- // ⚠️ 新增影子维度时:须更新此函数(或新增维度专用列举),并检查所有调用点(级联删除/一母一影校验/type 守卫),否则会静默漏掉新维度。
- // 软删除行由 SoftDeleteMixin 拦截器自动排除,无需手写 deleted_at IS NULL。
- func (r *accountRepository) ListShadowsByParent(ctx context.Context, parentID int64) ([]*service.Account, error) {
- rows, err := r.client.Account.Query().
- Where(dbaccount.ParentAccountIDEQ(parentID), dbaccount.QuotaDimensionEQ(dbaccount.QuotaDimensionSpark)).
- All(ctx)
- if err != nil {
- return nil, err
- }
- out := make([]*service.Account, 0, len(rows))
- for _, m := range rows {
- out = append(out, accountEntityToService(m))
- }
- return out, nil
- }
|