user_repo.go 48 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615
  1. package repository
  2. import (
  3. "context"
  4. "database/sql"
  5. "errors"
  6. "fmt"
  7. "sort"
  8. "strings"
  9. "time"
  10. dbent "github.com/Wei-Shaw/sub2api/ent"
  11. "github.com/Wei-Shaw/sub2api/ent/apikey"
  12. "github.com/Wei-Shaw/sub2api/ent/authidentity"
  13. "github.com/Wei-Shaw/sub2api/ent/authidentitychannel"
  14. dbgroup "github.com/Wei-Shaw/sub2api/ent/group"
  15. "github.com/Wei-Shaw/sub2api/ent/identityadoptiondecision"
  16. "github.com/Wei-Shaw/sub2api/ent/predicate"
  17. "github.com/Wei-Shaw/sub2api/ent/schema/mixins"
  18. dbuser "github.com/Wei-Shaw/sub2api/ent/user"
  19. "github.com/Wei-Shaw/sub2api/ent/userallowedgroup"
  20. "github.com/Wei-Shaw/sub2api/ent/usersubscription"
  21. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  22. "github.com/Wei-Shaw/sub2api/internal/service"
  23. "github.com/lib/pq"
  24. entsql "entgo.io/ent/dialect/sql"
  25. )
  26. type userRepository struct {
  27. client *dbent.Client
  28. sql sqlExecutor
  29. }
  30. var _ service.RedeemUserAdjustmentRepository = (*userRepository)(nil)
  31. func NewUserRepository(client *dbent.Client, sqlDB *sql.DB) service.UserRepository {
  32. return newUserRepositoryWithSQL(client, sqlDB)
  33. }
  34. func newUserRepositoryWithSQL(client *dbent.Client, sqlq sqlExecutor) *userRepository {
  35. return &userRepository{client: client, sql: sqlq}
  36. }
  37. func (r *userRepository) Create(ctx context.Context, userIn *service.User) error {
  38. return r.create(ctx, userIn, false, "")
  39. }
  40. // CreateWithEmailAliasGuard 见 service.UserRepository:在邮箱唯一性锁内复查收件箱身份,
  41. // 供注册路径使用。
  42. func (r *userRepository) CreateWithEmailAliasGuard(ctx context.Context, userIn *service.User) error {
  43. return r.create(ctx, userIn, true, "")
  44. }
  45. // CountUsersByEmailDomain 统计指定可注册主域名及其子域名下的未删除用户。
  46. func (r *userRepository) CountUsersByEmailDomain(ctx context.Context, domain string) (int, error) {
  47. return countUsersByEmailDomainWithClient(ctx, clientFromContext(ctx, r.client), domain)
  48. }
  49. // CreateWithEmailAliasGuardAndDomainLimit 串行化非白名单域名的注册请求,
  50. // 并在用户写入的同一事务内复查域名额度。
  51. func (r *userRepository) CreateWithEmailAliasGuardAndDomainLimit(ctx context.Context, userIn *service.User, domain string) error {
  52. return r.create(ctx, userIn, true, normalizeEmailDomain(domain))
  53. }
  54. func (r *userRepository) create(ctx context.Context, userIn *service.User, guardEmailAlias bool, domainLimit string) error {
  55. if userIn == nil {
  56. return nil
  57. }
  58. // 统一使用 ent 的事务:保证用户与允许分组的更新原子化,
  59. // 并避免基于 *sql.Tx 手动构造 ent client 导致的 ExecQuerier 断言错误。
  60. //
  61. // 注意:ent 的 Client.Tx 不感知上下文中的事务(只检查 driver 类型),
  62. // 因此必须显式检查 TxFromContext:当调用方已开启外部事务(如注册时的
  63. // “建用户 + 占用邀请码”原子事务),直接复用其 client,由调用方统一提交/回滚,
  64. // 否则用户写入会落入独立事务并自行提交,导致外层事务无法回滚(孤儿用户)。
  65. var txClient *dbent.Client
  66. txCtx := ctx
  67. var ownedTx *dbent.Tx
  68. if existingTx := dbent.TxFromContext(ctx); existingTx != nil {
  69. txClient = existingTx.Client()
  70. } else {
  71. tx, err := r.client.Tx(ctx)
  72. switch {
  73. case errors.Is(err, dbent.ErrTxStarted):
  74. // r.client 本身已是事务绑定 client(client 注入式事务,如集成测试
  75. // 夹具 tx.Client()):直接复用,提交/回滚由 client 的持有方负责。
  76. txClient = r.client
  77. case err != nil:
  78. return err
  79. default:
  80. ownedTx = tx
  81. defer func() { _ = ownedTx.Rollback() }()
  82. txClient = tx.Client()
  83. txCtx = dbent.NewTxContext(ctx, tx)
  84. }
  85. }
  86. releaseHierarchy, err := guardUserHierarchyCreate(txCtx, txClient, userIn)
  87. if err != nil {
  88. return err
  89. }
  90. defer releaseHierarchy()
  91. lockKeys := []string{normalizedEmailUniquenessLockKey(userIn.Email)}
  92. if guardEmailAlias {
  93. // 别名变体的字面量不同,唯一索引无法兜底;用收件箱身份锁把同一收件箱的并发注册串行化。
  94. lockKeys = append(lockKeys, emailAliasUniquenessLockKey(userIn.Email))
  95. }
  96. if domainLimit != "" {
  97. lockKeys = append(lockKeys, registrationEmailDomainLockKey(domainLimit))
  98. }
  99. releaseEmailLock, err := lockRepositoryScopedKeys(
  100. txCtx,
  101. txClient,
  102. txAwareSQLExecutor(txCtx, r.sql, r.client),
  103. lockKeys...,
  104. )
  105. if err != nil {
  106. return err
  107. }
  108. defer releaseEmailLock()
  109. if domainLimit != "" {
  110. count, err := countUsersByEmailDomainWithClient(txCtx, txClient, domainLimit)
  111. if err != nil {
  112. return err
  113. }
  114. if count > 0 {
  115. return service.ErrEmailDomainRegistrationLimit
  116. }
  117. }
  118. if err := ensureNormalizedEmailAvailableWithClient(txCtx, txClient, 0, userIn.Email); err != nil {
  119. return err
  120. }
  121. if guardEmailAlias {
  122. aliasExists, err := existsByEmailAliasWithClient(txCtx, txClient, userIn.Email)
  123. if err != nil {
  124. return err
  125. }
  126. if aliasExists {
  127. return service.ErrEmailExists
  128. }
  129. }
  130. created, err := txClient.User.Create().
  131. SetEmail(userIn.Email).
  132. SetUsername(userIn.Username).
  133. SetNotes(userIn.Notes).
  134. SetPasswordHash(userIn.PasswordHash).
  135. SetRole(userIn.Role).
  136. SetBalance(userIn.Balance).
  137. SetConcurrency(userIn.Concurrency).
  138. SetStatus(userIn.Status).
  139. SetSignupSource(userSignupSourceOrDefault(userIn.SignupSource)).
  140. SetNillableLastLoginAt(userIn.LastLoginAt).
  141. SetNillableLastActiveAt(userIn.LastActiveAt).
  142. SetRpmLimit(userIn.RPMLimit).
  143. SetRestrictPublicGroups(userIn.RestrictPublicGroups).
  144. Save(txCtx)
  145. if err != nil {
  146. return translatePersistenceError(err, nil, service.ErrEmailExists)
  147. }
  148. if err := r.syncUserAllowedGroupsWithClient(txCtx, txClient, created.ID, userIn.AllowedGroups); err != nil {
  149. return err
  150. }
  151. if err := ensureEmailAuthIdentityWithClient(txCtx, txClient, created.ID, created.Email, "user_repo_create"); err != nil {
  152. return err
  153. }
  154. if ownedTx != nil {
  155. if err := ownedTx.Commit(); err != nil {
  156. return err
  157. }
  158. }
  159. applyUserEntityToService(userIn, created)
  160. return nil
  161. }
  162. func (r *userRepository) GetByID(ctx context.Context, id int64) (*service.User, error) {
  163. m, err := r.client.User.Query().Where(dbuser.IDEQ(id)).Only(ctx)
  164. if err != nil {
  165. return nil, translatePersistenceError(err, service.ErrUserNotFound, nil)
  166. }
  167. out := userEntityToService(m)
  168. groups, err := r.loadAllowedGroups(ctx, []int64{id})
  169. if err != nil {
  170. return nil, err
  171. }
  172. if v, ok := groups[id]; ok {
  173. out.AllowedGroups = v
  174. }
  175. return out, nil
  176. }
  177. func (r *userRepository) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.User, error) {
  178. ctx = mixins.SkipSoftDelete(ctx)
  179. m, err := r.client.User.Query().Where(dbuser.IDEQ(id)).Only(ctx)
  180. if err != nil {
  181. return nil, translatePersistenceError(err, service.ErrUserNotFound, nil)
  182. }
  183. out := userEntityToService(m)
  184. groups, err := r.loadAllowedGroups(ctx, []int64{id})
  185. if err != nil {
  186. return nil, err
  187. }
  188. if v, ok := groups[id]; ok {
  189. out.AllowedGroups = v
  190. }
  191. return out, nil
  192. }
  193. func (r *userRepository) GetByEmail(ctx context.Context, email string) (*service.User, error) {
  194. matches, err := r.client.User.Query().
  195. Where(userEmailLookupPredicate(email)).
  196. Order(dbent.Asc(dbuser.FieldID)).
  197. All(ctx)
  198. if err != nil {
  199. return nil, err
  200. }
  201. if len(matches) == 0 {
  202. return nil, service.ErrUserNotFound
  203. }
  204. if len(matches) > 1 {
  205. return nil, fmt.Errorf("normalized email lookup matched multiple users for %q", strings.TrimSpace(email))
  206. }
  207. m := matches[0]
  208. out := userEntityToService(m)
  209. groups, err := r.loadAllowedGroups(ctx, []int64{m.ID})
  210. if err != nil {
  211. return nil, err
  212. }
  213. if v, ok := groups[m.ID]; ok {
  214. out.AllowedGroups = v
  215. }
  216. return out, nil
  217. }
  218. func (r *userRepository) Update(ctx context.Context, userIn *service.User, fields service.UserUpdateFields) error {
  219. if userIn == nil {
  220. return nil
  221. }
  222. // 空掩码代表调用方不改任何列,直接返回,避免产生一次无意义的整行写。
  223. if fields.IsEmpty() {
  224. return nil
  225. }
  226. // 使用 ent 事务包裹用户更新与 allowed_groups 同步,避免跨层事务不一致。
  227. var tx *dbent.Tx
  228. var txClient *dbent.Client
  229. txCtx := ctx
  230. if existingTx := dbent.TxFromContext(ctx); existingTx != nil {
  231. txClient = existingTx.Client()
  232. } else {
  233. var err error
  234. tx, err = r.client.Tx(ctx)
  235. switch {
  236. case errors.Is(err, dbent.ErrTxStarted):
  237. txClient = r.client
  238. case err != nil:
  239. return err
  240. default:
  241. defer func() { _ = tx.Rollback() }()
  242. txClient = tx.Client()
  243. txCtx = dbent.NewTxContext(ctx, tx)
  244. }
  245. }
  246. releaseHierarchy, err := guardUserHierarchyUpdate(txCtx, txClient, userIn, fields)
  247. if err != nil {
  248. return err
  249. }
  250. defer releaseHierarchy()
  251. // 邮箱唯一性锁与查重只在本次确实要改邮箱时才做:不改邮箱的更新既不需要
  252. // 串行化,也不该因为快照里的旧邮箱已被他人占用而报 ErrEmailExists。
  253. if fields.Email {
  254. releaseEmailLock, err := lockRepositoryScopedKeys(
  255. txCtx,
  256. txClient,
  257. txAwareSQLExecutor(txCtx, r.sql, r.client),
  258. normalizedEmailUniquenessLockKey(userIn.Email),
  259. )
  260. if err != nil {
  261. return err
  262. }
  263. defer releaseEmailLock()
  264. if err := ensureNormalizedEmailAvailableWithClient(txCtx, txClient, userIn.ID, userIn.Email); err != nil {
  265. return err
  266. }
  267. }
  268. existing, err := clientFromContext(txCtx, txClient).User.Get(txCtx, userIn.ID)
  269. if err != nil {
  270. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  271. }
  272. oldEmail := existing.Email
  273. updateOp := txClient.User.UpdateOneID(userIn.ID)
  274. if fields.Email {
  275. updateOp = updateOp.SetEmail(userIn.Email)
  276. }
  277. if fields.Username {
  278. updateOp = updateOp.SetUsername(userIn.Username)
  279. }
  280. if fields.Notes {
  281. updateOp = updateOp.SetNotes(userIn.Notes)
  282. }
  283. if fields.PasswordHash {
  284. updateOp = updateOp.SetPasswordHash(userIn.PasswordHash)
  285. }
  286. if fields.Role {
  287. updateOp = updateOp.SetRole(userIn.Role)
  288. }
  289. if fields.Concurrency {
  290. updateOp = updateOp.SetConcurrency(userIn.Concurrency)
  291. }
  292. if fields.RPMLimit {
  293. updateOp = updateOp.SetRpmLimit(userIn.RPMLimit)
  294. }
  295. if fields.Status {
  296. updateOp = updateOp.SetStatus(userIn.Status)
  297. }
  298. if fields.RestrictPublicGroups {
  299. updateOp = updateOp.SetRestrictPublicGroups(userIn.RestrictPublicGroups)
  300. }
  301. if fields.BalanceNotifySettings {
  302. updateOp = updateOp.
  303. SetBalanceNotifyEnabled(userIn.BalanceNotifyEnabled).
  304. SetBalanceNotifyThresholdType(userIn.BalanceNotifyThresholdType).
  305. SetNillableBalanceNotifyThreshold(userIn.BalanceNotifyThreshold)
  306. if userIn.BalanceNotifyThreshold == nil {
  307. updateOp = updateOp.ClearBalanceNotifyThreshold()
  308. }
  309. }
  310. if fields.BalanceNotifyExtraEmails {
  311. updateOp = updateOp.SetBalanceNotifyExtraEmails(marshalExtraEmails(userIn.BalanceNotifyExtraEmails))
  312. }
  313. if fields.SignupSource && userIn.SignupSource != "" {
  314. updateOp = updateOp.SetSignupSource(userIn.SignupSource)
  315. }
  316. if fields.LastLoginAt && userIn.LastLoginAt != nil {
  317. updateOp = updateOp.SetLastLoginAt(*userIn.LastLoginAt)
  318. }
  319. if fields.LastActiveAt && userIn.LastActiveAt != nil {
  320. updateOp = updateOp.SetLastActiveAt(*userIn.LastActiveAt)
  321. }
  322. updated, err := updateOp.Save(txCtx)
  323. if err != nil {
  324. return translatePersistenceError(err, service.ErrUserNotFound, service.ErrEmailExists)
  325. }
  326. if fields.AllowedGroups {
  327. if err := r.syncUserAllowedGroupsWithClient(txCtx, txClient, updated.ID, userIn.AllowedGroups); err != nil {
  328. return err
  329. }
  330. }
  331. // 始终以库中的邮箱为准补齐 email 身份:未改邮箱时 updated.Email == oldEmail,
  332. // 这里退化为幂等的身份补写,与改邮箱前的行为一致。
  333. if err := replaceEmailAuthIdentityWithClient(txCtx, txClient, updated.ID, oldEmail, updated.Email, "user_repo_update"); err != nil {
  334. return err
  335. }
  336. if tx != nil {
  337. if err := tx.Commit(); err != nil {
  338. return err
  339. }
  340. }
  341. userIn.UpdatedAt = updated.UpdatedAt
  342. return nil
  343. }
  344. func ensureEmailAuthIdentityWithClient(ctx context.Context, client *dbent.Client, userID int64, email string, source string) error {
  345. client = clientFromContext(ctx, client)
  346. if client == nil || userID <= 0 {
  347. return nil
  348. }
  349. subject := normalizeEmailAuthIdentitySubject(email)
  350. if subject == "" {
  351. return nil
  352. }
  353. if err := client.AuthIdentity.Create().
  354. SetUserID(userID).
  355. SetProviderType("email").
  356. SetProviderKey("email").
  357. SetProviderSubject(subject).
  358. SetVerifiedAt(time.Now().UTC()).
  359. SetMetadata(map[string]any{"source": source}).
  360. OnConflictColumns(
  361. authidentity.FieldProviderType,
  362. authidentity.FieldProviderKey,
  363. authidentity.FieldProviderSubject,
  364. ).
  365. DoNothing().
  366. Exec(ctx); err != nil {
  367. if !isSQLNoRowsError(err) {
  368. return err
  369. }
  370. }
  371. identity, err := client.AuthIdentity.Query().
  372. Where(
  373. authidentity.ProviderTypeEQ("email"),
  374. authidentity.ProviderKeyEQ("email"),
  375. authidentity.ProviderSubjectEQ(subject),
  376. ).
  377. Only(ctx)
  378. if err != nil {
  379. if dbent.IsNotFound(err) {
  380. return nil
  381. }
  382. return err
  383. }
  384. if identity.UserID != userID {
  385. return ErrAuthIdentityOwnershipConflict
  386. }
  387. return nil
  388. }
  389. func replaceEmailAuthIdentityWithClient(ctx context.Context, client *dbent.Client, userID int64, oldEmail, newEmail string, source string) error {
  390. newSubject := normalizeEmailAuthIdentitySubject(newEmail)
  391. if err := ensureEmailAuthIdentityWithClient(ctx, client, userID, newEmail, source); err != nil {
  392. return err
  393. }
  394. oldSubject := normalizeEmailAuthIdentitySubject(oldEmail)
  395. if oldSubject == "" || oldSubject == newSubject {
  396. return nil
  397. }
  398. _, err := clientFromContext(ctx, client).AuthIdentity.Delete().
  399. Where(
  400. authidentity.UserIDEQ(userID),
  401. authidentity.ProviderTypeEQ("email"),
  402. authidentity.ProviderKeyEQ("email"),
  403. authidentity.ProviderSubjectEQ(oldSubject),
  404. ).
  405. Exec(ctx)
  406. return err
  407. }
  408. func normalizeEmailAuthIdentitySubject(email string) string {
  409. normalized := strings.ToLower(strings.TrimSpace(email))
  410. if normalized == "" {
  411. return ""
  412. }
  413. if strings.HasSuffix(normalized, service.LinuxDoConnectSyntheticEmailDomain) ||
  414. strings.HasSuffix(normalized, service.OIDCConnectSyntheticEmailDomain) ||
  415. strings.HasSuffix(normalized, service.WeChatConnectSyntheticEmailDomain) ||
  416. strings.HasSuffix(normalized, service.DingTalkConnectSyntheticEmailDomain) {
  417. return ""
  418. }
  419. return normalized
  420. }
  421. func (r *userRepository) Delete(ctx context.Context, id int64) error {
  422. // 复用 context 中已存在的事务(如 AdminService.DeleteUser 把删 Key 与删 User 包在同一事务中),
  423. // 由调用方负责提交/回滚,保证两者的原子性。
  424. if existingTx := dbent.TxFromContext(ctx); existingTx != nil {
  425. return r.deleteUser(ctx, existingTx.Client(), id)
  426. }
  427. tx, err := r.client.Tx(ctx)
  428. if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
  429. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  430. }
  431. exec := r.client
  432. if err == nil {
  433. defer func() { _ = tx.Rollback() }()
  434. exec = tx.Client()
  435. }
  436. // err == dbent.ErrTxStarted 时复用当前事务(exec = r.client)。
  437. if err := r.deleteUser(ctx, exec, id); err != nil {
  438. return err
  439. }
  440. if tx != nil {
  441. if err := tx.Commit(); err != nil {
  442. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  443. }
  444. }
  445. return nil
  446. }
  447. // deleteUser 在给定 client(可能是外部事务 client)上删除用户及其身份关联记录,自身不开启/提交事务。
  448. func (r *userRepository) deleteUser(ctx context.Context, exec *dbent.Client, id int64) error {
  449. releaseHierarchy, err := guardUserHierarchyDelete(ctx, exec, id)
  450. if err != nil {
  451. return err
  452. }
  453. defer releaseHierarchy()
  454. identityIDs, err := exec.AuthIdentity.Query().
  455. Where(authidentity.UserIDEQ(id)).
  456. IDs(ctx)
  457. if err != nil {
  458. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  459. }
  460. if len(identityIDs) > 0 {
  461. if _, err := exec.IdentityAdoptionDecision.Update().
  462. Where(identityadoptiondecision.IdentityIDIn(identityIDs...)).
  463. ClearIdentityID().
  464. Save(ctx); err != nil {
  465. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  466. }
  467. if _, err := exec.AuthIdentityChannel.Delete().
  468. Where(authidentitychannel.IdentityIDIn(identityIDs...)).
  469. Exec(ctx); err != nil {
  470. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  471. }
  472. if _, err := exec.AuthIdentity.Delete().
  473. Where(authidentity.UserIDEQ(id)).
  474. Exec(ctx); err != nil {
  475. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  476. }
  477. }
  478. affected, err := exec.User.Delete().Where(dbuser.IDEQ(id)).Exec(ctx)
  479. if err != nil {
  480. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  481. }
  482. if affected == 0 {
  483. return service.ErrUserNotFound
  484. }
  485. return nil
  486. }
  487. func (r *userRepository) List(ctx context.Context, params pagination.PaginationParams) ([]service.User, *pagination.PaginationResult, error) {
  488. return r.ListWithFilters(ctx, params, service.UserListFilters{})
  489. }
  490. func (r *userRepository) ListWithFilters(ctx context.Context, params pagination.PaginationParams, filters service.UserListFilters) ([]service.User, *pagination.PaginationResult, error) {
  491. // SkipSoftDelete 仅作用于 User 身份解析(下方 Count/All);订阅、分组等关联实体沿用原始 ctx,避免穿透到这些同样带软删除的实体而带出已删除行。
  492. userCtx := ctx
  493. if filters.IncludeDeleted {
  494. userCtx = mixins.SkipSoftDelete(ctx)
  495. }
  496. q := r.client.User.Query()
  497. if filters.Status != "" {
  498. q = q.Where(dbuser.StatusEQ(filters.Status))
  499. }
  500. if filters.Role != "" {
  501. q = q.Where(dbuser.RoleEQ(filters.Role))
  502. }
  503. if filters.Search != "" {
  504. q = q.Where(
  505. dbuser.Or(
  506. dbuser.EmailContainsFold(filters.Search),
  507. dbuser.UsernameContainsFold(filters.Search),
  508. dbuser.NotesContainsFold(filters.Search),
  509. dbuser.HasAPIKeysWith(apikey.KeyContainsFold(filters.Search)),
  510. ),
  511. )
  512. }
  513. if filters.GroupName != "" {
  514. q = q.Where(dbuser.HasAllowedGroupsWith(
  515. dbgroup.NameContainsFold(filters.GroupName),
  516. ))
  517. }
  518. if filters.APIKeyGroupID > 0 {
  519. // 按"API Key 实际绑定的分组"过滤:用户只要有任意一个未软删除的 API Key
  520. // 绑定到该分组即命中(EXISTS 语义)。
  521. // 注意:SoftDeleteMixin 的拦截器不会自动下沉到 HasAPIKeysWith 子查询,
  522. // 必须显式加 apikey.DeletedAtIsNil(),否则已软删除的 key 会污染过滤结果。
  523. q = q.Where(dbuser.HasAPIKeysWith(
  524. apikey.GroupIDEQ(filters.APIKeyGroupID),
  525. apikey.DeletedAtIsNil(),
  526. ))
  527. }
  528. // If attribute filters are specified, we need to filter by user IDs first
  529. var allowedUserIDs []int64
  530. if len(filters.Attributes) > 0 {
  531. var attrErr error
  532. allowedUserIDs, attrErr = r.filterUsersByAttributes(ctx, filters.Attributes)
  533. if attrErr != nil {
  534. return nil, nil, attrErr
  535. }
  536. if len(allowedUserIDs) == 0 {
  537. // No users match the attribute filters
  538. return []service.User{}, paginationResultFromTotal(0, params), nil
  539. }
  540. q = q.Where(dbuser.IDIn(allowedUserIDs...))
  541. }
  542. total, err := q.Clone().Count(userCtx)
  543. if err != nil {
  544. return nil, nil, err
  545. }
  546. usersQuery := q.
  547. Offset(params.Offset()).
  548. Limit(params.Limit())
  549. for _, order := range userListOrder(params) {
  550. usersQuery = usersQuery.Order(order)
  551. }
  552. users, err := usersQuery.All(userCtx)
  553. if err != nil {
  554. return nil, nil, err
  555. }
  556. outUsers := make([]service.User, 0, len(users))
  557. if len(users) == 0 {
  558. return outUsers, paginationResultFromTotal(int64(total), params), nil
  559. }
  560. userIDs := make([]int64, 0, len(users))
  561. userMap := make(map[int64]*service.User, len(users))
  562. for i := range users {
  563. userIDs = append(userIDs, users[i].ID)
  564. u := userEntityToService(users[i])
  565. outUsers = append(outUsers, *u)
  566. userMap[u.ID] = &outUsers[len(outUsers)-1]
  567. }
  568. shouldLoadSubscriptions := filters.IncludeSubscriptions == nil || *filters.IncludeSubscriptions
  569. if shouldLoadSubscriptions {
  570. // Batch load active subscriptions with groups to avoid N+1.
  571. subs, err := r.client.UserSubscription.Query().
  572. Where(
  573. usersubscription.UserIDIn(userIDs...),
  574. usersubscription.StatusEQ(service.SubscriptionStatusActive),
  575. ).
  576. WithGroup().
  577. All(ctx)
  578. if err != nil {
  579. return nil, nil, err
  580. }
  581. for i := range subs {
  582. if u, ok := userMap[subs[i].UserID]; ok {
  583. u.Subscriptions = append(u.Subscriptions, *userSubscriptionEntityToService(subs[i]))
  584. }
  585. }
  586. }
  587. allowedGroupsByUser, err := r.loadAllowedGroups(ctx, userIDs)
  588. if err != nil {
  589. return nil, nil, err
  590. }
  591. for id, u := range userMap {
  592. if groups, ok := allowedGroupsByUser[id]; ok {
  593. u.AllowedGroups = groups
  594. }
  595. }
  596. return outUsers, paginationResultFromTotal(int64(total), params), nil
  597. }
  598. func userListOrder(params pagination.PaginationParams) []func(*entsql.Selector) {
  599. sortBy := strings.ToLower(strings.TrimSpace(params.SortBy))
  600. sortOrder := params.NormalizedSortOrder(pagination.SortOrderDesc)
  601. if sortBy == "last_used_at" {
  602. return userLastUsedAtOrder(sortOrder)
  603. }
  604. var field string
  605. defaultField := true
  606. nullsLastField := false
  607. switch sortBy {
  608. case "email":
  609. field = dbuser.FieldEmail
  610. defaultField = false
  611. case "username":
  612. field = dbuser.FieldUsername
  613. defaultField = false
  614. case "role":
  615. field = dbuser.FieldRole
  616. defaultField = false
  617. case "balance":
  618. field = dbuser.FieldBalance
  619. defaultField = false
  620. case "concurrency":
  621. field = dbuser.FieldConcurrency
  622. defaultField = false
  623. case "status":
  624. field = dbuser.FieldStatus
  625. defaultField = false
  626. case "created_at":
  627. field = dbuser.FieldCreatedAt
  628. defaultField = false
  629. case "last_active_at":
  630. field = dbuser.FieldLastActiveAt
  631. defaultField = false
  632. nullsLastField = true
  633. default:
  634. field = dbuser.FieldID
  635. }
  636. if sortOrder == pagination.SortOrderAsc {
  637. if defaultField && field == dbuser.FieldID {
  638. return []func(*entsql.Selector){dbent.Asc(dbuser.FieldID)}
  639. }
  640. if nullsLastField {
  641. return []func(*entsql.Selector){
  642. entsql.OrderByField(field, entsql.OrderNullsLast()).ToFunc(),
  643. dbent.Asc(dbuser.FieldID),
  644. }
  645. }
  646. return []func(*entsql.Selector){dbent.Asc(field), dbent.Asc(dbuser.FieldID)}
  647. }
  648. if defaultField && field == dbuser.FieldID {
  649. return []func(*entsql.Selector){dbent.Desc(dbuser.FieldID)}
  650. }
  651. if nullsLastField {
  652. return []func(*entsql.Selector){
  653. entsql.OrderByField(field, entsql.OrderDesc(), entsql.OrderNullsLast()).ToFunc(),
  654. dbent.Desc(dbuser.FieldID),
  655. }
  656. }
  657. return []func(*entsql.Selector){dbent.Desc(field), dbent.Desc(dbuser.FieldID)}
  658. }
  659. func (r *userRepository) GetLatestUsedAtByUserIDs(ctx context.Context, userIDs []int64) (map[int64]*time.Time, error) {
  660. result := make(map[int64]*time.Time, len(userIDs))
  661. if len(userIDs) == 0 {
  662. return result, nil
  663. }
  664. if r.sql == nil {
  665. return nil, fmt.Errorf("sql executor is not configured")
  666. }
  667. const query = `
  668. SELECT user_id, MAX(created_at) AS last_used_at
  669. FROM usage_logs
  670. WHERE user_id = ANY($1)
  671. GROUP BY user_id
  672. `
  673. rows, err := r.sql.QueryContext(ctx, query, pq.Array(userIDs))
  674. if err != nil {
  675. return nil, err
  676. }
  677. defer func() { _ = rows.Close() }()
  678. for rows.Next() {
  679. var (
  680. userID int64
  681. lastUsedAt time.Time
  682. )
  683. if scanErr := rows.Scan(&userID, &lastUsedAt); scanErr != nil {
  684. return nil, scanErr
  685. }
  686. ts := lastUsedAt.UTC()
  687. result[userID] = &ts
  688. }
  689. if err := rows.Err(); err != nil {
  690. return nil, err
  691. }
  692. return result, nil
  693. }
  694. func (r *userRepository) GetLatestUsedAtByUserID(ctx context.Context, userID int64) (*time.Time, error) {
  695. latestByUserID, err := r.GetLatestUsedAtByUserIDs(ctx, []int64{userID})
  696. if err != nil {
  697. return nil, err
  698. }
  699. return latestByUserID[userID], nil
  700. }
  701. func userLastUsedAtOrder(sortOrder string) []func(*entsql.Selector) {
  702. orderExpr := func(direction, nulls string, tieOrder func(string) string) func(*entsql.Selector) {
  703. return func(s *entsql.Selector) {
  704. subquery := fmt.Sprintf("(SELECT MAX(created_at) FROM usage_logs WHERE user_id = %s)", s.C(dbuser.FieldID))
  705. s.OrderExpr(entsql.Expr(subquery + " " + direction + " NULLS " + nulls))
  706. s.OrderBy(tieOrder(s.C(dbuser.FieldID)))
  707. }
  708. }
  709. if sortOrder == pagination.SortOrderAsc {
  710. return []func(*entsql.Selector){
  711. orderExpr("ASC", "FIRST", entsql.Asc),
  712. }
  713. }
  714. return []func(*entsql.Selector){
  715. orderExpr("DESC", "LAST", entsql.Desc),
  716. }
  717. }
  718. // filterUsersByAttributes returns user IDs that match ALL the given attribute filters
  719. func (r *userRepository) filterUsersByAttributes(ctx context.Context, attrs map[int64]string) ([]int64, error) {
  720. if len(attrs) == 0 {
  721. return nil, nil
  722. }
  723. if r.sql == nil {
  724. return nil, fmt.Errorf("sql executor is not configured")
  725. }
  726. clauses := make([]string, 0, len(attrs))
  727. args := make([]any, 0, len(attrs)*2+1)
  728. argIndex := 1
  729. for attrID, value := range attrs {
  730. clauses = append(clauses, fmt.Sprintf("(attribute_id = $%d AND value ILIKE $%d)", argIndex, argIndex+1))
  731. args = append(args, attrID, "%"+value+"%")
  732. argIndex += 2
  733. }
  734. query := fmt.Sprintf(
  735. `SELECT user_id
  736. FROM user_attribute_values
  737. WHERE %s
  738. GROUP BY user_id
  739. HAVING COUNT(DISTINCT attribute_id) = $%d`,
  740. strings.Join(clauses, " OR "),
  741. argIndex,
  742. )
  743. args = append(args, len(attrs))
  744. rows, err := r.sql.QueryContext(ctx, query, args...)
  745. if err != nil {
  746. return nil, err
  747. }
  748. defer func() { _ = rows.Close() }()
  749. result := make([]int64, 0)
  750. for rows.Next() {
  751. var userID int64
  752. if scanErr := rows.Scan(&userID); scanErr != nil {
  753. return nil, scanErr
  754. }
  755. result = append(result, userID)
  756. }
  757. if err := rows.Err(); err != nil {
  758. return nil, err
  759. }
  760. return result, nil
  761. }
  762. func (r *userRepository) UpdateBalance(ctx context.Context, id int64, amount float64) error {
  763. client := clientFromContext(ctx, r.client)
  764. update := client.User.Update().Where(dbuser.IDEQ(id)).AddBalance(amount)
  765. // Track cumulative recharge amount for percentage-based notifications
  766. if amount > 0 {
  767. update = update.AddTotalRecharged(amount)
  768. }
  769. n, err := update.Save(ctx)
  770. if err != nil {
  771. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  772. }
  773. if n == 0 {
  774. return service.ErrUserNotFound
  775. }
  776. return nil
  777. }
  778. func (r *userRepository) ApplyRedeemBalanceAdjustment(ctx context.Context, id int64, delta float64) error {
  779. const updateSQL = `
  780. UPDATE users
  781. SET balance = GREATEST(balance + $1, 0), updated_at = NOW()
  782. WHERE id = $2 AND deleted_at IS NULL
  783. `
  784. client := clientFromContext(ctx, r.client)
  785. result, err := client.ExecContext(ctx, updateSQL, delta, id)
  786. if err != nil {
  787. return err
  788. }
  789. affected, err := result.RowsAffected()
  790. if err != nil {
  791. return err
  792. }
  793. if affected == 0 {
  794. return service.ErrUserNotFound
  795. }
  796. return nil
  797. }
  798. // DeductBalance 扣除用户余额
  799. // 透支策略:允许余额变为负数,确保当前请求能够完成
  800. // 中间件会阻止余额 <= 0 的用户发起后续请求
  801. func (r *userRepository) DeductBalance(ctx context.Context, id int64, amount float64) error {
  802. client := clientFromContext(ctx, r.client)
  803. n, err := client.User.Update().
  804. Where(dbuser.IDEQ(id), dbuser.BalanceGTE(amount)).
  805. AddBalance(-amount).
  806. Save(ctx)
  807. if err != nil {
  808. return err
  809. }
  810. if n > 0 {
  811. return nil
  812. }
  813. n, err = client.User.Update().
  814. Where(dbuser.IDEQ(id)).
  815. AddBalance(-amount).
  816. Save(ctx)
  817. if err != nil {
  818. return err
  819. }
  820. if n == 0 {
  821. return service.ErrUserNotFound
  822. }
  823. return nil
  824. }
  825. // DeductAvailableBalance atomically deducts min(amount, max(balance, 0)).
  826. // Unlike DeductBalance, this refund-specific operation never increases an
  827. // existing deficit or permits a concurrent deduction to cause an overdraft.
  828. func (r *userRepository) DeductAvailableBalance(ctx context.Context, id int64, amount float64) (deducted float64, err error) {
  829. if amount < 0 {
  830. return 0, fmt.Errorf("deduction amount must be nonnegative")
  831. }
  832. const updateSQL = `
  833. WITH target AS (
  834. SELECT id, balance
  835. FROM users
  836. WHERE id = $2 AND deleted_at IS NULL
  837. FOR UPDATE
  838. ), updated AS (
  839. UPDATE users AS u
  840. SET balance = target.balance - LEAST($1, GREATEST(target.balance, 0)), updated_at = NOW()
  841. FROM target
  842. WHERE u.id = target.id AND u.deleted_at IS NULL
  843. RETURNING target.balance - u.balance AS deducted
  844. )
  845. SELECT deducted FROM updated
  846. `
  847. rows, err := clientFromContext(ctx, r.client).QueryContext(ctx, updateSQL, amount, id)
  848. if err != nil {
  849. return 0, err
  850. }
  851. defer func() {
  852. if closeErr := rows.Close(); closeErr != nil && err == nil {
  853. err = closeErr
  854. }
  855. }()
  856. if !rows.Next() {
  857. if rowsErr := rows.Err(); rowsErr != nil {
  858. return 0, rowsErr
  859. }
  860. return 0, service.ErrUserNotFound
  861. }
  862. if err := rows.Scan(&deducted); err != nil {
  863. return 0, err
  864. }
  865. return deducted, rows.Err()
  866. }
  867. // AdjustBalance 原子地把 delta 累加到余额上,结果为负时整条语句不生效。
  868. // 相比"读余额 → 算新值 → 整行写回",这里把读与写压进同一条 UPDATE,
  869. // 并发的计费扣款不会被旧快照覆盖。
  870. func (r *userRepository) AdjustBalance(ctx context.Context, id int64, delta float64) (service.BalanceChange, error) {
  871. const updateSQL = `
  872. UPDATE users
  873. SET balance = balance + $1, updated_at = NOW()
  874. WHERE id = $2 AND deleted_at IS NULL AND balance + $1 >= 0
  875. RETURNING balance - $1, balance
  876. `
  877. change, ok, err := scanBalanceChange(ctx, clientFromContext(ctx, r.client), updateSQL, delta, id)
  878. if err != nil {
  879. return service.BalanceChange{}, err
  880. }
  881. if ok {
  882. return change, nil
  883. }
  884. // 0 行既可能是用户不存在,也可能是余额不足以承受这次扣减,需要区分。
  885. current, err := r.currentBalance(ctx, id)
  886. if err != nil {
  887. return service.BalanceChange{}, err
  888. }
  889. return service.BalanceChange{Old: current, New: current + delta}, service.ErrBalanceNegative
  890. }
  891. // SetBalance 原子地把余额置为 value,并返回变更前后的值。
  892. func (r *userRepository) SetBalance(ctx context.Context, id int64, value float64) (service.BalanceChange, error) {
  893. if value < 0 {
  894. // 连同当前余额一起返回,便于上层给出可读的错误信息。
  895. current, err := r.currentBalance(ctx, id)
  896. if err != nil {
  897. return service.BalanceChange{}, err
  898. }
  899. return service.BalanceChange{Old: current, New: value}, service.ErrBalanceNegative
  900. }
  901. const updateSQL = `
  902. UPDATE users AS u
  903. SET balance = $1, updated_at = NOW()
  904. FROM (SELECT id, balance FROM users WHERE id = $2 AND deleted_at IS NULL) AS prev
  905. WHERE u.id = prev.id AND u.deleted_at IS NULL
  906. RETURNING prev.balance, u.balance
  907. `
  908. change, ok, err := scanBalanceChange(ctx, clientFromContext(ctx, r.client), updateSQL, value, id)
  909. if err != nil {
  910. return service.BalanceChange{}, err
  911. }
  912. if !ok {
  913. return service.BalanceChange{}, service.ErrUserNotFound
  914. }
  915. return change, nil
  916. }
  917. // currentBalance 读取用户当前余额,用户不存在时返回 ErrUserNotFound。
  918. func (r *userRepository) currentBalance(ctx context.Context, id int64) (balance float64, err error) {
  919. rows, err := clientFromContext(ctx, r.client).QueryContext(ctx,
  920. `SELECT balance FROM users WHERE id = $1 AND deleted_at IS NULL`, id)
  921. if err != nil {
  922. return 0, err
  923. }
  924. defer func() {
  925. if closeErr := rows.Close(); closeErr != nil && err == nil {
  926. err = closeErr
  927. }
  928. }()
  929. if !rows.Next() {
  930. if rowsErr := rows.Err(); rowsErr != nil {
  931. return 0, rowsErr
  932. }
  933. return 0, service.ErrUserNotFound
  934. }
  935. if err := rows.Scan(&balance); err != nil {
  936. return 0, err
  937. }
  938. return balance, rows.Err()
  939. }
  940. // scanBalanceChange 执行一条 RETURNING 旧余额、新余额的语句。ok 为 false 表示语句未命中任何行。
  941. func scanBalanceChange(ctx context.Context, client *dbent.Client, query string, args ...any) (change service.BalanceChange, ok bool, err error) {
  942. rows, err := client.QueryContext(ctx, query, args...)
  943. if err != nil {
  944. return service.BalanceChange{}, false, err
  945. }
  946. defer func() {
  947. if closeErr := rows.Close(); closeErr != nil && err == nil {
  948. err = closeErr
  949. }
  950. }()
  951. if !rows.Next() {
  952. if rowsErr := rows.Err(); rowsErr != nil {
  953. return service.BalanceChange{}, false, rowsErr
  954. }
  955. return service.BalanceChange{}, false, nil
  956. }
  957. if err := rows.Scan(&change.Old, &change.New); err != nil {
  958. return service.BalanceChange{}, false, err
  959. }
  960. return change, true, rows.Err()
  961. }
  962. func (r *userRepository) UpdateConcurrency(ctx context.Context, id int64, amount int) error {
  963. client := clientFromContext(ctx, r.client)
  964. n, err := client.User.Update().Where(dbuser.IDEQ(id)).AddConcurrency(amount).Save(ctx)
  965. if err != nil {
  966. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  967. }
  968. if n == 0 {
  969. return service.ErrUserNotFound
  970. }
  971. return nil
  972. }
  973. func (r *userRepository) ApplyRedeemConcurrencyAdjustment(ctx context.Context, id int64, delta int) error {
  974. const updateSQL = `
  975. UPDATE users
  976. SET concurrency = GREATEST(concurrency + $1, 0), updated_at = NOW()
  977. WHERE id = $2 AND deleted_at IS NULL
  978. `
  979. client := clientFromContext(ctx, r.client)
  980. result, err := client.ExecContext(ctx, updateSQL, delta, id)
  981. if err != nil {
  982. return err
  983. }
  984. affected, err := result.RowsAffected()
  985. if err != nil {
  986. return err
  987. }
  988. if affected == 0 {
  989. return service.ErrUserNotFound
  990. }
  991. return nil
  992. }
  993. func (r *userRepository) BatchSetConcurrency(ctx context.Context, userIDs []int64, value int) (int, error) {
  994. if len(userIDs) == 0 {
  995. return 0, nil
  996. }
  997. if value < 0 {
  998. value = 0
  999. }
  1000. res, err := r.sql.ExecContext(ctx,
  1001. "UPDATE users SET concurrency = $1, updated_at = NOW() WHERE id = ANY($2) AND deleted_at IS NULL",
  1002. value, pq.Array(userIDs))
  1003. if err != nil {
  1004. return 0, fmt.Errorf("batch set concurrency: %w", err)
  1005. }
  1006. affected, _ := res.RowsAffected()
  1007. return int(affected), nil
  1008. }
  1009. func (r *userRepository) BatchAddConcurrency(ctx context.Context, userIDs []int64, delta int) (int, error) {
  1010. if len(userIDs) == 0 {
  1011. return 0, nil
  1012. }
  1013. res, err := r.sql.ExecContext(ctx,
  1014. "UPDATE users SET concurrency = GREATEST(concurrency + $1, 0), updated_at = NOW() WHERE id = ANY($2) AND deleted_at IS NULL",
  1015. delta, pq.Array(userIDs))
  1016. if err != nil {
  1017. return 0, fmt.Errorf("batch add concurrency: %w", err)
  1018. }
  1019. affected, _ := res.RowsAffected()
  1020. return int(affected), nil
  1021. }
  1022. func (r *userRepository) BatchUpdateLimits(ctx context.Context, userIDs []int64, concurrency, rpmLimit *int) (int, error) {
  1023. if len(userIDs) == 0 || (concurrency == nil && rpmLimit == nil) {
  1024. return 0, nil
  1025. }
  1026. setClauses := make([]string, 0, 3)
  1027. args := make([]any, 0, 3)
  1028. if concurrency != nil {
  1029. value := max(*concurrency, 0)
  1030. args = append(args, value)
  1031. setClauses = append(setClauses, fmt.Sprintf("concurrency = $%d", len(args)))
  1032. }
  1033. if rpmLimit != nil {
  1034. value := max(*rpmLimit, 0)
  1035. args = append(args, value)
  1036. setClauses = append(setClauses, fmt.Sprintf("rpm_limit = $%d", len(args)))
  1037. }
  1038. setClauses = append(setClauses, "updated_at = NOW()")
  1039. args = append(args, pq.Array(userIDs))
  1040. query := fmt.Sprintf(
  1041. "UPDATE users SET %s WHERE id = ANY($%d) AND deleted_at IS NULL",
  1042. strings.Join(setClauses, ", "),
  1043. len(args),
  1044. )
  1045. res, err := r.sql.ExecContext(ctx, query, args...)
  1046. if err != nil {
  1047. return 0, fmt.Errorf("batch update user limits: %w", err)
  1048. }
  1049. affected, _ := res.RowsAffected()
  1050. return int(affected), nil
  1051. }
  1052. func (r *userRepository) ExistsByEmail(ctx context.Context, email string) (bool, error) {
  1053. return r.client.User.Query().Where(userEmailLookupPredicate(email)).Exist(ctx)
  1054. }
  1055. // emailAliasCandidateLimit 限制一次别名查重最多取回的候选行数。探针都以去点后的
  1056. // 本地部分为前缀锚定(见 dotStrippedEmailExpr),正常收件箱的变体只有个位数;
  1057. // 上限只是兜底,避免公开未鉴权的注册/发码端点把大表整张读进内存。
  1058. const emailAliasCandidateLimit = 50
  1059. // ExistsByEmailAlias 见 service.UserRepository。软删除过滤沿用 ExistsByEmail 的默认行为。
  1060. func (r *userRepository) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
  1061. return existsByEmailAliasWithClient(ctx, clientFromContext(ctx, r.client), email)
  1062. }
  1063. func existsByEmailAliasWithClient(ctx context.Context, client *dbent.Client, email string) (bool, error) {
  1064. _, exists, err := emailAliasOwnerIDWithClient(ctx, client, email, 0)
  1065. return exists, err
  1066. }
  1067. func emailAliasOwnerIDWithClient(ctx context.Context, client *dbent.Client, email string, currentUserID int64) (int64, bool, error) {
  1068. if client == nil {
  1069. return 0, false, nil
  1070. }
  1071. probes := service.EmailAliasDedupProbes(email)
  1072. if len(probes) == 0 {
  1073. return 0, false, nil
  1074. }
  1075. preds := make([]predicate.User, 0, 2*len(probes))
  1076. for _, probe := range probes {
  1077. preds = append(preds,
  1078. dotStrippedEmailEQ(probe.Local+"@"+probe.Domain),
  1079. // "+后缀"的内容未知,只能按前缀匹配。
  1080. dotStrippedEmailLike(escapeLikeWildcards(probe.Local)+"+%@"+escapeLikeWildcards(probe.Domain)),
  1081. )
  1082. }
  1083. candidates, err := client.User.Query().
  1084. Where(dbuser.Or(preds...)).
  1085. Limit(emailAliasCandidateLimit).
  1086. Select(dbuser.FieldID, dbuser.FieldEmail).
  1087. All(ctx)
  1088. if err != nil {
  1089. return 0, false, err
  1090. }
  1091. // 探针会有过度匹配(点号只在 Gmail 家族无意义),最终判定必须回到完整归一化规则。
  1092. // 返回“其他用户”优先于当前用户,避免历史重复数据让调用方误判为仅当前用户占用。
  1093. identity := service.NormalizeEmailForAliasDedup(email)
  1094. var selfID int64
  1095. selfExists := false
  1096. for _, candidate := range candidates {
  1097. if service.NormalizeEmailForAliasDedup(candidate.Email) != identity {
  1098. continue
  1099. }
  1100. if candidate.ID != 0 && candidate.ID != currentUserID {
  1101. return candidate.ID, true, nil
  1102. }
  1103. if candidate.ID == currentUserID {
  1104. selfID = candidate.ID
  1105. selfExists = true
  1106. }
  1107. }
  1108. return selfID, selfExists, nil
  1109. }
  1110. // UpdateEmailWithAliasGuard 在调用方事务内更新主邮箱与密码哈希。
  1111. //
  1112. // 邮箱换绑不能只依赖服务层前置查重:两个并发请求可能同时看到同一收件箱未被占用。
  1113. // 这里先按“字面邮箱 + 收件箱身份”加锁,复查是否已被其他用户占用,再执行写入;
  1114. // PostgreSQL 使用事务级 advisory lock 跨实例互斥,测试内存库则由进程内锁兜底。
  1115. func (r *userRepository) UpdateEmailWithAliasGuard(
  1116. ctx context.Context,
  1117. userID int64,
  1118. email string,
  1119. passwordHash string,
  1120. ) error {
  1121. if userID <= 0 {
  1122. return service.ErrUserNotFound
  1123. }
  1124. if strings.TrimSpace(email) == "" || passwordHash == "" {
  1125. return fmt.Errorf("email identity update requires email and password hash")
  1126. }
  1127. tx := dbent.TxFromContext(ctx)
  1128. if tx == nil {
  1129. return fmt.Errorf("email identity update requires a transaction")
  1130. }
  1131. client := tx.Client()
  1132. releaseEmailLock, err := lockRepositoryScopedKeys(
  1133. ctx,
  1134. client,
  1135. txAwareSQLExecutor(ctx, r.sql, r.client),
  1136. normalizedEmailUniquenessLockKey(email),
  1137. emailAliasUniquenessLockKey(email),
  1138. )
  1139. if err != nil {
  1140. return err
  1141. }
  1142. defer releaseEmailLock()
  1143. ownerID, exists, err := emailAliasOwnerIDWithClient(ctx, client, email, userID)
  1144. if err != nil {
  1145. return err
  1146. }
  1147. if exists && ownerID != userID {
  1148. return service.ErrEmailExists
  1149. }
  1150. if _, err := client.User.UpdateOneID(userID).
  1151. SetEmail(email).
  1152. SetPasswordHash(passwordHash).
  1153. Save(ctx); err != nil {
  1154. return translatePersistenceError(err, service.ErrUserNotFound, service.ErrEmailExists)
  1155. }
  1156. return nil
  1157. }
  1158. // dotStrippedEmailExpr 渲染下面的表达式:去掉存量邮箱的大小写、首尾空白(与
  1159. // userEmailLookupPredicate 的精确匹配口径一致,历史数据存在带空白的行)以及全部点号。
  1160. //
  1161. // REPLACE(LOWER(TRIM(email)), '.', '')
  1162. //
  1163. // 两侧都去点,因此一个域名探针即可同时覆盖 Gmail 点号变体与 FQDN 根点(user@gmail.com.)。
  1164. // migrations/190 为同一表达式建了索引。
  1165. func dotStrippedEmailExpr(b *entsql.Builder, s *entsql.Selector) *entsql.Builder {
  1166. return b.WriteString("REPLACE(LOWER(TRIM(").
  1167. Ident(s.C(dbuser.FieldEmail)).
  1168. WriteString(")), '.', '')")
  1169. }
  1170. func dotStrippedEmailEQ(value string) predicate.User {
  1171. return predicate.User(func(s *entsql.Selector) {
  1172. s.Where(entsql.P(func(b *entsql.Builder) {
  1173. dotStrippedEmailExpr(b, s).WriteString(" = ").Arg(value)
  1174. }))
  1175. })
  1176. }
  1177. func dotStrippedEmailLike(pattern string) predicate.User {
  1178. return predicate.User(func(s *entsql.Selector) {
  1179. s.Where(entsql.P(func(b *entsql.Builder) {
  1180. dotStrippedEmailExpr(b, s).WriteString(" LIKE ").Arg(pattern).WriteString(` ESCAPE '\'`)
  1181. }))
  1182. })
  1183. }
  1184. // escapeLikeWildcards 转义 LIKE 元字符:本地部分合法可含 % 与 _,不转义会扩大匹配面。
  1185. var likeWildcardEscaper = strings.NewReplacer(`\`, `\\`, "%", `\%`, "_", `\_`)
  1186. func escapeLikeWildcards(value string) string {
  1187. return likeWildcardEscaper.Replace(value)
  1188. }
  1189. func ensureNormalizedEmailAvailableWithClient(ctx context.Context, client *dbent.Client, userID int64, email string) error {
  1190. client = clientFromContext(ctx, client)
  1191. if client == nil {
  1192. return nil
  1193. }
  1194. matches, err := client.User.Query().
  1195. Where(userEmailLookupPredicate(email)).
  1196. All(ctx)
  1197. if err != nil {
  1198. return err
  1199. }
  1200. for _, match := range matches {
  1201. if match.ID != userID {
  1202. return service.ErrEmailExists
  1203. }
  1204. }
  1205. return nil
  1206. }
  1207. func userEmailLookupPredicate(email string) predicate.User {
  1208. normalized := normalizeEmailLookupValue(email)
  1209. if normalized == "" {
  1210. return dbuser.EmailEQ(email)
  1211. }
  1212. return predicate.User(func(s *entsql.Selector) {
  1213. s.Where(entsql.P(func(b *entsql.Builder) {
  1214. b.WriteString("LOWER(TRIM(").
  1215. Ident(s.C(dbuser.FieldEmail)).
  1216. WriteString(")) = ").
  1217. Arg(normalized)
  1218. }))
  1219. })
  1220. }
  1221. func normalizeEmailLookupValue(email string) string {
  1222. return strings.ToLower(strings.TrimSpace(email))
  1223. }
  1224. func normalizedEmailUniquenessLockKey(email string) string {
  1225. normalized := normalizeEmailLookupValue(email)
  1226. if normalized == "" {
  1227. return ""
  1228. }
  1229. return "users:normalized-email:" + normalized
  1230. }
  1231. func registrationEmailDomainLockKey(domain string) string {
  1232. domain = normalizeEmailDomain(domain)
  1233. if domain == "" {
  1234. return ""
  1235. }
  1236. return "users:registration-email-domain:" + domain
  1237. }
  1238. func normalizeEmailDomain(domain string) string {
  1239. return service.NormalizeRegistrationEmailDomain(domain)
  1240. }
  1241. func countUsersByEmailDomainWithClient(ctx context.Context, client *dbent.Client, domain string) (int, error) {
  1242. client = clientFromContext(ctx, client)
  1243. domain = normalizeEmailDomain(domain)
  1244. if client == nil || domain == "" {
  1245. return 0, nil
  1246. }
  1247. return client.User.Query().Where(userEmailDomainPredicate(domain)).Count(ctx)
  1248. }
  1249. func userEmailDomainPredicate(domain string) predicate.User {
  1250. domain = normalizeEmailDomain(domain)
  1251. escapedDomain := escapeLikeWildcards(domain)
  1252. exactPattern := "%@" + escapedDomain
  1253. subdomainPattern := "%@%." + escapedDomain
  1254. return predicate.User(func(s *entsql.Selector) {
  1255. s.Where(entsql.P(func(b *entsql.Builder) {
  1256. b.WriteString("(RTRIM(LOWER(TRIM(").
  1257. Ident(s.C(dbuser.FieldEmail)).
  1258. WriteString(")), '.') LIKE ").
  1259. Arg(exactPattern).
  1260. WriteString(` ESCAPE '\' OR RTRIM(LOWER(TRIM(`).
  1261. Ident(s.C(dbuser.FieldEmail)).
  1262. WriteString(")), '.') LIKE ").
  1263. Arg(subdomainPattern).
  1264. WriteString(` ESCAPE '\'`).
  1265. WriteString(")")
  1266. }))
  1267. })
  1268. }
  1269. // emailAliasUniquenessLockKey 按收件箱身份(而非邮箱字面量)加锁,使同一收件箱的不同
  1270. // 别名变体在注册时互斥。
  1271. func emailAliasUniquenessLockKey(email string) string {
  1272. identity := service.NormalizeEmailForAliasDedup(email)
  1273. if identity == "" {
  1274. return ""
  1275. }
  1276. return "users:email-alias-identity:" + identity
  1277. }
  1278. func (r *userRepository) AddGroupToAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
  1279. client := clientFromContext(ctx, r.client)
  1280. err := client.UserAllowedGroup.Create().
  1281. SetUserID(userID).
  1282. SetGroupID(groupID).
  1283. OnConflictColumns(userallowedgroup.FieldUserID, userallowedgroup.FieldGroupID).
  1284. DoNothing().
  1285. Exec(ctx)
  1286. if isSQLNoRowsError(err) {
  1287. return nil
  1288. }
  1289. return err
  1290. }
  1291. func (r *userRepository) RemoveGroupFromAllowedGroups(ctx context.Context, groupID int64) (int64, error) {
  1292. // 仅操作 user_allowed_groups 联接表,legacy users.allowed_groups 列已弃用。
  1293. affected, err := r.client.UserAllowedGroup.Delete().
  1294. Where(userallowedgroup.GroupIDEQ(groupID)).
  1295. Exec(ctx)
  1296. if err != nil {
  1297. return 0, err
  1298. }
  1299. return int64(affected), nil
  1300. }
  1301. // RemoveGroupFromUserAllowedGroups 移除单个用户的指定分组权限
  1302. func (r *userRepository) RemoveGroupFromUserAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
  1303. client := clientFromContext(ctx, r.client)
  1304. _, err := client.UserAllowedGroup.Delete().
  1305. Where(userallowedgroup.UserIDEQ(userID), userallowedgroup.GroupIDEQ(groupID)).
  1306. Exec(ctx)
  1307. return err
  1308. }
  1309. func (r *userRepository) GetFirstAdmin(ctx context.Context) (*service.User, error) {
  1310. m, err := r.client.User.Query().
  1311. Where(
  1312. dbuser.RoleIn(service.RoleAdmin, service.RoleSuperAdmin),
  1313. dbuser.StatusEQ(service.StatusActive),
  1314. ).
  1315. Order(dbent.Desc(dbuser.FieldRole), dbent.Asc(dbuser.FieldID)).
  1316. First(ctx)
  1317. if err != nil {
  1318. return nil, translatePersistenceError(err, service.ErrUserNotFound, nil)
  1319. }
  1320. out := userEntityToService(m)
  1321. groups, err := r.loadAllowedGroups(ctx, []int64{m.ID})
  1322. if err != nil {
  1323. return nil, err
  1324. }
  1325. if v, ok := groups[m.ID]; ok {
  1326. out.AllowedGroups = v
  1327. }
  1328. return out, nil
  1329. }
  1330. func (r *userRepository) loadAllowedGroups(ctx context.Context, userIDs []int64) (map[int64][]int64, error) {
  1331. out := make(map[int64][]int64, len(userIDs))
  1332. if len(userIDs) == 0 {
  1333. return out, nil
  1334. }
  1335. rows, err := r.client.UserAllowedGroup.Query().
  1336. Where(userallowedgroup.UserIDIn(userIDs...)).
  1337. All(ctx)
  1338. if err != nil {
  1339. return nil, err
  1340. }
  1341. for i := range rows {
  1342. out[rows[i].UserID] = append(out[rows[i].UserID], rows[i].GroupID)
  1343. }
  1344. for userID := range out {
  1345. sort.Slice(out[userID], func(i, j int) bool { return out[userID][i] < out[userID][j] })
  1346. }
  1347. return out, nil
  1348. }
  1349. // syncUserAllowedGroupsWithClient 在 ent client/事务内同步用户允许分组:
  1350. // 仅操作 user_allowed_groups 联接表,legacy users.allowed_groups 列已弃用。
  1351. func (r *userRepository) syncUserAllowedGroupsWithClient(ctx context.Context, client *dbent.Client, userID int64, groupIDs []int64) error {
  1352. if client == nil {
  1353. return nil
  1354. }
  1355. existingRows, err := client.UserAllowedGroup.Query().
  1356. Where(userallowedgroup.UserIDEQ(userID)).
  1357. All(ctx)
  1358. if err != nil {
  1359. return err
  1360. }
  1361. desired := make(map[int64]struct{}, len(groupIDs))
  1362. for _, id := range groupIDs {
  1363. if id <= 0 {
  1364. continue
  1365. }
  1366. desired[id] = struct{}{}
  1367. }
  1368. existing := make(map[int64]struct{}, len(existingRows))
  1369. removed := make([]int64, 0)
  1370. for _, row := range existingRows {
  1371. existing[row.GroupID] = struct{}{}
  1372. if _, keep := desired[row.GroupID]; !keep {
  1373. removed = append(removed, row.GroupID)
  1374. }
  1375. }
  1376. if len(removed) > 0 {
  1377. if _, err := client.UserAllowedGroup.Delete().
  1378. Where(userallowedgroup.UserIDEQ(userID), userallowedgroup.GroupIDIn(removed...)).
  1379. Exec(ctx); err != nil {
  1380. return err
  1381. }
  1382. }
  1383. creates := make([]*dbent.UserAllowedGroupCreate, 0, len(desired))
  1384. for groupID := range desired {
  1385. if _, present := existing[groupID]; !present {
  1386. creates = append(creates, client.UserAllowedGroup.Create().SetUserID(userID).SetGroupID(groupID))
  1387. }
  1388. }
  1389. if len(creates) > 0 {
  1390. if err := client.UserAllowedGroup.
  1391. CreateBulk(creates...).
  1392. OnConflictColumns(userallowedgroup.FieldUserID, userallowedgroup.FieldGroupID).
  1393. DoNothing().
  1394. Exec(ctx); err != nil {
  1395. if isSQLNoRowsError(err) {
  1396. return nil
  1397. }
  1398. return err
  1399. }
  1400. }
  1401. return nil
  1402. }
  1403. func applyUserEntityToService(dst *service.User, src *dbent.User) {
  1404. if dst == nil || src == nil {
  1405. return
  1406. }
  1407. dst.ID = src.ID
  1408. dst.SignupSource = src.SignupSource
  1409. dst.LastLoginAt = src.LastLoginAt
  1410. dst.LastActiveAt = src.LastActiveAt
  1411. dst.CreatedAt = src.CreatedAt
  1412. dst.UpdatedAt = src.UpdatedAt
  1413. }
  1414. func userSignupSourceOrDefault(signupSource string) string {
  1415. switch strings.TrimSpace(strings.ToLower(signupSource)) {
  1416. case "", "email":
  1417. return "email"
  1418. case "linuxdo", "wechat", "oidc", "dingtalk":
  1419. return strings.TrimSpace(strings.ToLower(signupSource))
  1420. default:
  1421. return "email"
  1422. }
  1423. }
  1424. // marshalExtraEmails serializes notify email entries to JSON for storage.
  1425. func marshalExtraEmails(entries []service.NotifyEmailEntry) string {
  1426. return service.MarshalNotifyEmails(entries)
  1427. }
  1428. // UpdateTotpSecret 更新用户的 TOTP 加密密钥
  1429. func (r *userRepository) UpdateTotpSecret(ctx context.Context, userID int64, encryptedSecret *string) error {
  1430. client := clientFromContext(ctx, r.client)
  1431. update := client.User.UpdateOneID(userID)
  1432. if encryptedSecret == nil {
  1433. update = update.ClearTotpSecretEncrypted()
  1434. } else {
  1435. update = update.SetTotpSecretEncrypted(*encryptedSecret)
  1436. }
  1437. _, err := update.Save(ctx)
  1438. if err != nil {
  1439. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  1440. }
  1441. return nil
  1442. }
  1443. // EnableTotp 启用用户的 TOTP 双因素认证
  1444. func (r *userRepository) EnableTotp(ctx context.Context, userID int64) error {
  1445. client := clientFromContext(ctx, r.client)
  1446. _, err := client.User.UpdateOneID(userID).
  1447. SetTotpEnabled(true).
  1448. SetTotpEnabledAt(time.Now()).
  1449. Save(ctx)
  1450. if err != nil {
  1451. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  1452. }
  1453. return nil
  1454. }
  1455. // DisableTotp 禁用用户的 TOTP 双因素认证
  1456. func (r *userRepository) DisableTotp(ctx context.Context, userID int64) error {
  1457. client := clientFromContext(ctx, r.client)
  1458. _, err := client.User.UpdateOneID(userID).
  1459. SetTotpEnabled(false).
  1460. ClearTotpEnabledAt().
  1461. ClearTotpSecretEncrypted().
  1462. Save(ctx)
  1463. if err != nil {
  1464. return translatePersistenceError(err, service.ErrUserNotFound, nil)
  1465. }
  1466. return nil
  1467. }