user_platform_quota_repo.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508
  1. package repository
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "strings"
  7. "time"
  8. dbent "github.com/Wei-Shaw/sub2api/ent"
  9. "github.com/Wei-Shaw/sub2api/ent/userplatformquota"
  10. "github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
  11. "github.com/lib/pq"
  12. )
  13. // UserPlatformQuotaRecord 是 repository 层的传输结构体,
  14. // 与 ent.UserPlatformQuota 实体解耦,供业务层使用。
  15. type UserPlatformQuotaRecord struct {
  16. UserID int64
  17. Platform string
  18. DailyLimitUSD *float64
  19. WeeklyLimitUSD *float64
  20. MonthlyLimitUSD *float64
  21. DailyUsageUSD float64
  22. WeeklyUsageUSD float64
  23. MonthlyUsageUSD float64
  24. DailyWindowStart *time.Time
  25. WeeklyWindowStart *time.Time
  26. MonthlyWindowStart *time.Time
  27. }
  28. // ErrUserPlatformQuotaNotFound 用于 ResetExpiredWindow 等需要"必须命中已有记录"的方法。
  29. var ErrUserPlatformQuotaNotFound = fmt.Errorf("user platform quota record not found")
  30. // ErrUserPlatformQuotaFKViolation 当批量 UPSERT 中存在 user_id 不在 users 表的记录时返回。
  31. var ErrUserPlatformQuotaFKViolation = errors.New("user platform quota snapshot FK violation")
  32. // UserPlatformQuotaSnapshot 是 BatchSnapshotUsage 的输入结构体,
  33. // 表示 Redis 当前窗口快照(用于绝对值覆盖写入 DB)。
  34. type UserPlatformQuotaSnapshot struct {
  35. UserID int64
  36. Platform string
  37. DailyUsageUSD float64
  38. WeeklyUsageUSD float64
  39. MonthlyUsageUSD float64
  40. DailyWindowStart time.Time
  41. WeeklyWindowStart time.Time
  42. MonthlyWindowStart time.Time
  43. }
  44. // UserPlatformQuotaRepository 定义用户平台配额的数据访问接口。
  45. type UserPlatformQuotaRepository interface {
  46. // BulkInsertInitial 幂等批量插入初始配额记录(ON CONFLICT DO NOTHING)。
  47. BulkInsertInitial(ctx context.Context, records []UserPlatformQuotaRecord) error
  48. // GetByUserPlatform 查询单条配额记录,未找到时返回 (nil, nil)。
  49. GetByUserPlatform(ctx context.Context, userID int64, platform string) (*UserPlatformQuotaRecord, error)
  50. // ListByUser 查询用户的所有平台配额记录(排除软删除)。
  51. ListByUser(ctx context.Context, userID int64) ([]UserPlatformQuotaRecord, error)
  52. // IncrementUsageWithReset 原子地累加用量,若窗口已过期则先重置再累加。
  53. IncrementUsageWithReset(ctx context.Context, userID int64, platform string, cost float64, now time.Time) error
  54. // ResetExpiredWindow 重置指定窗口(daily/weekly/monthly)的用量与起始时间。
  55. ResetExpiredWindow(ctx context.Context, userID int64, platform string, window string, newStart time.Time) error
  56. // UpsertForUser 全量替换该用户所有平台限额配置(详见 service.UserPlatformQuotaRepository.UpsertForUser)。
  57. UpsertForUser(ctx context.Context, userID int64, records []UserPlatformQuotaRecord) error
  58. // BatchSnapshotUsage 用一条多行 UPSERT 把整批 usage 以绝对值覆盖写入(非累加)。
  59. // usage/window_start 直接取 EXCLUDED(Redis 当前窗口快照),无 CASE。整批共用 now 作 created/updated_at。
  60. // 要求 snapshots 内 (user,platform) 不重复。FK 违反返回 ErrUserPlatformQuotaFKViolation。
  61. BatchSnapshotUsage(ctx context.Context, snapshots []UserPlatformQuotaSnapshot, now time.Time) error
  62. }
  63. type userPlatformQuotaRepository struct {
  64. client *dbent.Client
  65. }
  66. // NewUserPlatformQuotaRepository 创建 UserPlatformQuotaRepository 实现。
  67. func NewUserPlatformQuotaRepository(client *dbent.Client) UserPlatformQuotaRepository {
  68. return &userPlatformQuotaRepository{client: client}
  69. }
  70. // BulkInsertInitial 用原生 SQL ON CONFLICT 实现幂等批量插入(带条件 limit 覆盖)。
  71. // 仅插入 limit_usd 与元数据,usage_usd 用 DB 默认 0,window_start 留 NULL。
  72. // FK 约束要求 user_id 在 users 表中存在,调用方负责保证。
  73. //
  74. // 冲突策略:CASE WHEN existing.*_limit_usd IS NULL THEN EXCLUDED.*_limit_usd ELSE existing ...
  75. // - 若 IncrementUsageWithReset 因时序问题已先建行(limit 全 NULL),
  76. // 此处会把注册时的默认 limit 写入,避免该用户在该平台永久无限额。
  77. // - 若管理员已通过 UpsertForUser 设置了非 NULL 个性化 limit,**保留不动**
  78. // —— 旧实现无条件 EXCLUDED 覆盖会丢失个性化配置。
  79. // - 不会改 usage_usd / window_start,保留累计的用量。
  80. // - 仅命中 deleted_at IS NULL 的活跃记录(partial unique index 作用域)。
  81. func (r *userPlatformQuotaRepository) BulkInsertInitial(ctx context.Context, records []UserPlatformQuotaRecord) error {
  82. if len(records) == 0 {
  83. return nil
  84. }
  85. client := clientFromContext(ctx, r.client)
  86. var sb strings.Builder
  87. _, _ = sb.WriteString("INSERT INTO user_platform_quotas (user_id, platform, daily_limit_usd, weekly_limit_usd, monthly_limit_usd, daily_usage_usd, weekly_usage_usd, monthly_usage_usd, created_at, updated_at) VALUES ")
  88. args := make([]any, 0, len(records)*6)
  89. // 统一时间戳:避免循环内多次 time.Now() 让同一批记录的 created_at/updated_at
  90. // 出现亚毫秒级偏差(与 UpsertForUser 的 now := time.Now() 风格一致)。
  91. now := time.Now()
  92. for i, rec := range records {
  93. base := i * 6
  94. if i > 0 {
  95. _, _ = sb.WriteString(",")
  96. }
  97. fmt.Fprintf(&sb, "($%d,$%d,$%d,$%d,$%d,0,0,0,$%d,$%d)",
  98. base+1, base+2, base+3, base+4, base+5, base+6, base+6)
  99. args = append(args,
  100. rec.UserID, rec.Platform,
  101. rec.DailyLimitUSD, rec.WeeklyLimitUSD, rec.MonthlyLimitUSD,
  102. now,
  103. )
  104. }
  105. // 精确命中 partial unique index(deleted_at IS NULL),避免对软删记录的歧义冲突。
  106. // 条件覆盖:仅在现有 limit 为 NULL 时才写入 EXCLUDED,否则保留现有非 NULL 值。
  107. // - 修复 IncrementUsageWithReset 已用 NULL limit 建行的场景(NULL → 注册默认)
  108. // - 保护管理员通过 UpsertForUser 设置的个性化 limit 不被静默覆盖
  109. _, _ = sb.WriteString(` ON CONFLICT (user_id, platform) WHERE deleted_at IS NULL
  110. DO UPDATE SET
  111. daily_limit_usd = COALESCE(user_platform_quotas.daily_limit_usd, EXCLUDED.daily_limit_usd),
  112. weekly_limit_usd = COALESCE(user_platform_quotas.weekly_limit_usd, EXCLUDED.weekly_limit_usd),
  113. monthly_limit_usd = COALESCE(user_platform_quotas.monthly_limit_usd, EXCLUDED.monthly_limit_usd),
  114. updated_at = EXCLUDED.updated_at`)
  115. _, err := client.ExecContext(ctx, sb.String(), args...)
  116. return err
  117. }
  118. // GetByUserPlatform 通过 ent 查询单条配额(排除软删除)。未找到返回 (nil, nil)。
  119. func (r *userPlatformQuotaRepository) GetByUserPlatform(ctx context.Context, userID int64, platform string) (*UserPlatformQuotaRecord, error) {
  120. client := clientFromContext(ctx, r.client)
  121. entity, err := client.UserPlatformQuota.Query().
  122. Where(
  123. userplatformquota.UserIDEQ(userID),
  124. userplatformquota.PlatformEQ(platform),
  125. userplatformquota.DeletedAtIsNil(),
  126. ).
  127. Only(ctx)
  128. if dbent.IsNotFound(err) {
  129. return nil, nil
  130. }
  131. if err != nil {
  132. return nil, err
  133. }
  134. return entQuotaToRecord(entity), nil
  135. }
  136. // ListByUser 查询用户的所有平台配额记录(排除软删除)。
  137. func (r *userPlatformQuotaRepository) ListByUser(ctx context.Context, userID int64) ([]UserPlatformQuotaRecord, error) {
  138. client := clientFromContext(ctx, r.client)
  139. rows, err := client.UserPlatformQuota.Query().
  140. Where(
  141. userplatformquota.UserIDEQ(userID),
  142. userplatformquota.DeletedAtIsNil(),
  143. ).
  144. All(ctx)
  145. if err != nil {
  146. return nil, err
  147. }
  148. out := make([]UserPlatformQuotaRecord, 0, len(rows))
  149. for _, e := range rows {
  150. out = append(out, *entQuotaToRecord(e))
  151. }
  152. return out, nil
  153. }
  154. // IncrementUsageWithReset 原子累加 cost 到 (user, platform) 三个窗口的 *_usage_usd。
  155. // 行为:
  156. // - 若记录存在:在事务内 SELECT FOR UPDATE,按 (prev_window_start vs current_window_start)
  157. // 判断是否需要重置(不同 = 重置为 cost;相同 = 累加 cost)
  158. // - 若记录不存在(fail-open create 分支):插入新记录,**limit 字段保留 nil(无限制)**
  159. // —— 这是预期行为:billing 链路不能因 quota 表缺失而阻断请求,未注册路径
  160. // 的用户 quota 默认放行,由调度层指标观测 + 后台对账补建 limit
  161. //
  162. // 上层正常路径(注册时 BulkInsertInitial)保证 limit 在记录创建时就被写入。
  163. func (r *userPlatformQuotaRepository) IncrementUsageWithReset(ctx context.Context, userID int64, platform string, cost float64, now time.Time) error {
  164. return r.withTx(ctx, func(txCtx context.Context, txClient *dbent.Client) error {
  165. existing, err := txClient.UserPlatformQuota.Query().
  166. Where(
  167. userplatformquota.UserIDEQ(userID),
  168. userplatformquota.PlatformEQ(platform),
  169. userplatformquota.DeletedAtIsNil(),
  170. ).
  171. ForUpdate().
  172. Only(txCtx)
  173. if dbent.IsNotFound(err) {
  174. // fail-open 建行:limit_* 保留 NULL(无限额)。
  175. // 用 ON CONFLICT DO UPDATE 累加,而非裸 INSERT:并发下另一请求可能在本事务
  176. // SELECT FOR UPDATE 之后、INSERT 之前刚建行,裸 INSERT 会撞 partial unique index
  177. // 致事务回滚、本次 cost 丢失;DO UPDATE 把 cost 累加到既有 usage 上。
  178. // 写法与本文件 insertLimitsRow / BulkInsertInitial 的 ON CONFLICT 一致。
  179. const insertSQL = `INSERT INTO user_platform_quotas
  180. (user_id, platform, daily_usage_usd, weekly_usage_usd, monthly_usage_usd,
  181. daily_window_start, weekly_window_start, monthly_window_start, created_at, updated_at)
  182. VALUES ($1, $2, $3, $3, $3, $4, $5, $6, $7, $7)
  183. ON CONFLICT (user_id, platform) WHERE deleted_at IS NULL DO UPDATE SET
  184. daily_usage_usd = user_platform_quotas.daily_usage_usd + EXCLUDED.daily_usage_usd,
  185. weekly_usage_usd = user_platform_quotas.weekly_usage_usd + EXCLUDED.weekly_usage_usd,
  186. monthly_usage_usd = user_platform_quotas.monthly_usage_usd + EXCLUDED.monthly_usage_usd,
  187. updated_at = EXCLUDED.updated_at`
  188. // $6 = now:30 天滚动月度窗口以当前时刻为起始
  189. _, e := txClient.ExecContext(txCtx, insertSQL,
  190. userID, platform, cost,
  191. timezone.StartOfDay(now), timezone.StartOfWeek(now), now, now)
  192. return e
  193. }
  194. if err != nil {
  195. return err
  196. }
  197. newDaily := maybeReset(existing.DailyUsageUsd, existing.DailyWindowStart, timezone.StartOfDay(now), cost)
  198. newWeekly := maybeReset(existing.WeeklyUsageUsd, existing.WeeklyWindowStart, timezone.StartOfWeek(now), cost)
  199. // 30 天滚动月度窗口:过期时重置为 cost 并以 now 为新起始,否则累加保留原起始
  200. newMonthly, newMonthlyStart := monthlyMaybeReset(existing.MonthlyUsageUsd, existing.MonthlyWindowStart, cost, now)
  201. _, e := existing.Update().
  202. SetDailyUsageUsd(newDaily).
  203. SetWeeklyUsageUsd(newWeekly).
  204. SetMonthlyUsageUsd(newMonthly).
  205. SetDailyWindowStart(timezone.StartOfDay(now)).
  206. SetWeeklyWindowStart(timezone.StartOfWeek(now)).
  207. SetMonthlyWindowStart(newMonthlyStart). // 30 天滚动:仅过期时更新起始
  208. Save(txCtx)
  209. return e
  210. })
  211. }
  212. // ResetExpiredWindow 无条件重置指定窗口(daily/weekly/monthly)的用量与起始时间。
  213. //
  214. // ⚠️ 命名警告(NOT a "check-then-reset" helper):
  215. //
  216. // 名字里的 "Expired" 是历史遗留,**实现并不校验窗口是否真的过期**。
  217. // 任何调用都会无条件把对应窗口的 *_usage_usd 清零并重写 *_window_start。
  218. // 目前唯一合法 caller 是 admin POST /reset 接口(管理员强制归零)。
  219. //
  220. // 如果你想要"仅在窗口过期才重置"的语义,请直接使用 IncrementUsageWithReset
  221. // 的内部判断(maybeReset / monthlyMaybeReset),或新增独立函数;
  222. // 不要复用这里的函数,否则会出现"明明窗口未过期,用量却被清零"的隐蔽 bug。
  223. //
  224. // 未命中活跃记录时返回 ErrUserPlatformQuotaNotFound。
  225. func (r *userPlatformQuotaRepository) ResetExpiredWindow(ctx context.Context, userID int64, platform string, window string, newStart time.Time) error {
  226. client := clientFromContext(ctx, r.client)
  227. upd := client.UserPlatformQuota.Update().
  228. Where(
  229. userplatformquota.UserIDEQ(userID),
  230. userplatformquota.PlatformEQ(platform),
  231. userplatformquota.DeletedAtIsNil(),
  232. )
  233. switch window {
  234. case "daily":
  235. upd = upd.SetDailyUsageUsd(0).SetDailyWindowStart(newStart)
  236. case "weekly":
  237. upd = upd.SetWeeklyUsageUsd(0).SetWeeklyWindowStart(newStart)
  238. case "monthly":
  239. upd = upd.SetMonthlyUsageUsd(0).SetMonthlyWindowStart(newStart)
  240. default:
  241. return fmt.Errorf("unknown window %q", window)
  242. }
  243. n, err := upd.Save(ctx)
  244. if err != nil {
  245. return err
  246. }
  247. if n == 0 {
  248. return ErrUserPlatformQuotaNotFound
  249. }
  250. return nil
  251. }
  252. // withTx 在事务中执行 fn,若 ctx 中已有事务则复用。
  253. func (r *userPlatformQuotaRepository) withTx(ctx context.Context, fn func(txCtx context.Context, txClient *dbent.Client) error) error {
  254. if tx := dbent.TxFromContext(ctx); tx != nil {
  255. return fn(ctx, tx.Client())
  256. }
  257. tx, err := r.client.Tx(ctx)
  258. if err != nil {
  259. return fmt.Errorf("begin user_platform_quota transaction: %w", err)
  260. }
  261. defer func() { _ = tx.Rollback() }()
  262. txCtx := dbent.NewTxContext(ctx, tx)
  263. if err := fn(txCtx, tx.Client()); err != nil {
  264. return err
  265. }
  266. if err := tx.Commit(); err != nil {
  267. return fmt.Errorf("commit user_platform_quota transaction: %w", err)
  268. }
  269. return nil
  270. }
  271. // entQuotaToRecord 将 ent entity 映射为 repository record。
  272. // 注意 ent 生成字段名为 DailyLimitUsd(非 DailyLimitUSD)。
  273. func entQuotaToRecord(e *dbent.UserPlatformQuota) *UserPlatformQuotaRecord {
  274. return &UserPlatformQuotaRecord{
  275. UserID: e.UserID,
  276. Platform: e.Platform,
  277. DailyLimitUSD: e.DailyLimitUsd,
  278. WeeklyLimitUSD: e.WeeklyLimitUsd,
  279. MonthlyLimitUSD: e.MonthlyLimitUsd,
  280. DailyUsageUSD: e.DailyUsageUsd,
  281. WeeklyUsageUSD: e.WeeklyUsageUsd,
  282. MonthlyUsageUSD: e.MonthlyUsageUsd,
  283. DailyWindowStart: e.DailyWindowStart,
  284. WeeklyWindowStart: e.WeeklyWindowStart,
  285. MonthlyWindowStart: e.MonthlyWindowStart,
  286. }
  287. }
  288. // maybeReset 判断是否需要重置窗口用量:
  289. // - 若 prevStart 为 nil 或与 currStart 不同,表示窗口已过期,返回 cost(重置)
  290. // - 否则返回 prevUsage + cost(累加)
  291. func maybeReset(prevUsage float64, prevStart *time.Time, currStart time.Time, cost float64) float64 {
  292. if prevStart == nil || !prevStart.Equal(currStart) {
  293. return cost
  294. }
  295. return prevUsage + cost
  296. }
  297. // monthlyMaybeReset 判断 30 天滚动月度窗口是否需要重置。
  298. // 过期条件:prevStart 为 nil 或 now - prevStart >= 30×24h(与订阅模式 NeedsMonthlyReset 语义一致)。
  299. // 过期时重置为 cost,否则累加。返回 (newUsage, newWindowStart)。
  300. func monthlyMaybeReset(prevUsage float64, prevStart *time.Time, cost float64, now time.Time) (float64, time.Time) {
  301. if prevStart == nil || now.Sub(*prevStart) >= 30*24*time.Hour {
  302. return cost, now
  303. }
  304. return prevUsage + cost, *prevStart
  305. }
  306. // UpsertForUser 全量替换该用户的所有平台限额(事务内):
  307. // 1. 软删除未在 records 中出现的所有 active 行
  308. // 2. 对每条 record 尝试 UPDATE(含 deleted_at = NULL 兼容重激活);
  309. // UPDATE 行数为 0 时 INSERT 新行
  310. //
  311. // 仅改 *_limit_usd + deleted_at + updated_at,保留 *_usage_usd / *_window_start。
  312. func (r *userPlatformQuotaRepository) UpsertForUser(ctx context.Context, userID int64, records []UserPlatformQuotaRecord) error {
  313. return r.withTx(ctx, func(txCtx context.Context, txClient *dbent.Client) error {
  314. platforms := make([]string, 0, len(records))
  315. for _, rec := range records {
  316. platforms = append(platforms, rec.Platform)
  317. }
  318. now := time.Now()
  319. if err := softDeleteMissingPlatforms(txCtx, txClient, userID, platforms, now); err != nil {
  320. return err
  321. }
  322. for _, rec := range records {
  323. affected, err := updateLimitsRow(txCtx, txClient, userID, rec, now)
  324. if err != nil {
  325. return err
  326. }
  327. if affected == 0 {
  328. if err := insertLimitsRow(txCtx, txClient, userID, rec, now); err != nil {
  329. return err
  330. }
  331. }
  332. }
  333. return nil
  334. })
  335. }
  336. // softDeleteMissingPlatforms 软删除该用户所有不在 keepPlatforms 中的 active 行。
  337. // keepPlatforms 为空时 → 软删用户所有 active 行。
  338. // now 由调用方传入,与 updateLimitsRow / insertLimitsRow 共享同一个 Go time.Now(),
  339. // 保证事务内所有时间戳一致(避免 Postgres NOW() 与 Go time.Now() 的微小偏差)。
  340. func softDeleteMissingPlatforms(ctx context.Context, client *dbent.Client, userID int64, keepPlatforms []string, now time.Time) error {
  341. var (
  342. query string
  343. args []any
  344. )
  345. if len(keepPlatforms) == 0 {
  346. query = `UPDATE user_platform_quotas SET deleted_at = $2, updated_at = $2
  347. WHERE user_id = $1 AND deleted_at IS NULL`
  348. args = []any{userID, now}
  349. } else {
  350. placeholders := make([]string, len(keepPlatforms))
  351. args = make([]any, 0, len(keepPlatforms)+2)
  352. args = append(args, userID, now)
  353. for i, p := range keepPlatforms {
  354. placeholders[i] = fmt.Sprintf("$%d", i+3)
  355. args = append(args, p)
  356. }
  357. query = fmt.Sprintf(`UPDATE user_platform_quotas SET deleted_at = $2, updated_at = $2
  358. WHERE user_id = $1 AND deleted_at IS NULL AND platform NOT IN (%s)`,
  359. strings.Join(placeholders, ","))
  360. }
  361. _, err := client.ExecContext(ctx, query, args...)
  362. return err
  363. }
  364. // updateLimitsRow 尝试 UPDATE active 行(deleted_at IS NULL),返回受影响行数。
  365. // 仅更新 active 行:若存在多条历史软删记录,加 deleted_at IS NULL 守卫可避免
  366. // 批量重激活导致的 partial unique index(userplatformquota_user_id_platform_uq)冲突。
  367. // affected=0 时由调用方 UpsertForUser 走 insertLimitsRow 路径创建新行。
  368. func updateLimitsRow(ctx context.Context, client *dbent.Client, userID int64, rec UserPlatformQuotaRecord, now time.Time) (int64, error) {
  369. const query = `UPDATE user_platform_quotas
  370. SET daily_limit_usd = $1, weekly_limit_usd = $2, monthly_limit_usd = $3,
  371. deleted_at = NULL, updated_at = $4
  372. WHERE user_id = $5 AND platform = $6 AND deleted_at IS NULL`
  373. res, err := client.ExecContext(ctx, query,
  374. rec.DailyLimitUSD, rec.WeeklyLimitUSD, rec.MonthlyLimitUSD, now,
  375. userID, rec.Platform)
  376. if err != nil {
  377. return 0, err
  378. }
  379. return res.RowsAffected()
  380. }
  381. // insertLimitsRow 插入新限额行(usage 默认 0,window_start 默认 NULL)。
  382. // 带 ON CONFLICT ... DO NOTHING 守卫:防止两个并发请求同时为同一 user/platform 新建行时
  383. // 触发 unique constraint 违反(userplatformquota_user_id_platform_uq 部分唯一索引)。
  384. // affected=0 时说明另一个并发请求刚完成 INSERT,fallback 到 updateLimitsRow 覆写 limits 值。
  385. func insertLimitsRow(ctx context.Context, client *dbent.Client, userID int64, rec UserPlatformQuotaRecord, now time.Time) error {
  386. const query = `INSERT INTO user_platform_quotas
  387. (user_id, platform, daily_limit_usd, weekly_limit_usd, monthly_limit_usd,
  388. daily_usage_usd, weekly_usage_usd, monthly_usage_usd, created_at, updated_at)
  389. VALUES ($1, $2, $3, $4, $5, 0, 0, 0, $6, $6)
  390. ON CONFLICT (user_id, platform) WHERE deleted_at IS NULL DO NOTHING`
  391. res, err := client.ExecContext(ctx, query,
  392. userID, rec.Platform,
  393. rec.DailyLimitUSD, rec.WeeklyLimitUSD, rec.MonthlyLimitUSD,
  394. now)
  395. if err != nil {
  396. return err
  397. }
  398. affected, err := res.RowsAffected()
  399. if err != nil {
  400. return err
  401. }
  402. if affected == 0 {
  403. // 并发情形:另一请求已插入该行,fallback 到 UPDATE 覆写 limits 值(last-writer-wins)。
  404. _, err = updateLimitsRow(ctx, client, userID, rec, now)
  405. return err
  406. }
  407. return nil
  408. }
  409. // batchRows 是 BatchSnapshotUsage 每批最大行数(9 参/行 × 6000 ≈ 54000 参,低于 Postgres 65535 上限)。
  410. const batchRows = 6000
  411. // BatchSnapshotUsage 用一条多行 UPSERT 把整批 usage 以绝对值覆盖写入(非累加)。
  412. // 每批最多 batchRows 行;$1=now 共用;每行 8 个 per-row 参(user_id, platform, 3×usage, 3×window_start)。
  413. // FK 违反(user_id 不存在)返回 ErrUserPlatformQuotaFKViolation。
  414. //
  415. // 注意:snapshots 超过 batchRows 会分多条 SQL 执行且【非单事务】——若某子批 FK 失败,
  416. // 先前子批已写入无法回滚。调用方(flusher)应保证单次 batchSize ≤ batchRows
  417. // (默认 flush_batch_size=1000 < 6000,安全)。
  418. // 另注:启用 flusher 后,本绝对值覆盖与 admin 直写 DB(ResetExpiredWindow/UpsertForUser)存在覆盖竞态,
  419. // 详见 service/user_platform_quota_flusher.go 中 flushOneBatch 的"已知竞态"注释。
  420. func (r *userPlatformQuotaRepository) BatchSnapshotUsage(ctx context.Context, snapshots []UserPlatformQuotaSnapshot, now time.Time) error {
  421. if len(snapshots) == 0 {
  422. return nil
  423. }
  424. client := clientFromContext(ctx, r.client)
  425. for start := 0; start < len(snapshots); start += batchRows {
  426. end := start + batchRows
  427. if end > len(snapshots) {
  428. end = len(snapshots)
  429. }
  430. batch := snapshots[start:end]
  431. var sb strings.Builder
  432. _, _ = sb.WriteString(
  433. "INSERT INTO user_platform_quotas" +
  434. " (user_id, platform, daily_usage_usd, weekly_usage_usd, monthly_usage_usd," +
  435. " daily_window_start, weekly_window_start, monthly_window_start, created_at, updated_at)" +
  436. " VALUES ")
  437. // $1 = now(共用);每行 8 个 per-row 参,从 $2 起连续编号。
  438. args := []any{now}
  439. for i, s := range batch {
  440. if i > 0 {
  441. _, _ = sb.WriteString(",")
  442. }
  443. b := len(args) // 当前 per-row 第一个参数的 0-based 索引,实际占位符 = b+1
  444. fmt.Fprintf(&sb, "($%d,$%d,$%d,$%d,$%d,$%d,$%d,$%d,$1,$1)",
  445. b+1, b+2, b+3, b+4, b+5, b+6, b+7, b+8)
  446. args = append(args,
  447. s.UserID, s.Platform,
  448. s.DailyUsageUSD, s.WeeklyUsageUSD, s.MonthlyUsageUSD,
  449. s.DailyWindowStart, s.WeeklyWindowStart, s.MonthlyWindowStart,
  450. )
  451. }
  452. _, _ = sb.WriteString(
  453. " ON CONFLICT (user_id, platform) WHERE deleted_at IS NULL DO UPDATE SET" +
  454. " daily_usage_usd = EXCLUDED.daily_usage_usd," +
  455. " weekly_usage_usd = EXCLUDED.weekly_usage_usd," +
  456. " monthly_usage_usd = EXCLUDED.monthly_usage_usd," +
  457. " daily_window_start = EXCLUDED.daily_window_start," +
  458. " weekly_window_start = EXCLUDED.weekly_window_start," +
  459. " monthly_window_start = EXCLUDED.monthly_window_start," +
  460. " updated_at = EXCLUDED.updated_at")
  461. if _, err := client.ExecContext(ctx, sb.String(), args...); err != nil {
  462. var pqErr *pq.Error
  463. if errors.As(err, &pqErr) && pqErr.Code == "23503" {
  464. return ErrUserPlatformQuotaFKViolation
  465. }
  466. return err
  467. }
  468. }
  469. return nil
  470. }