dashboard_aggregation_repo.go 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689
  1. package repository
  2. import (
  3. "context"
  4. "database/sql"
  5. "fmt"
  6. "log"
  7. "sort"
  8. "strings"
  9. "time"
  10. "github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
  11. "github.com/Wei-Shaw/sub2api/internal/service"
  12. "github.com/lib/pq"
  13. )
  14. type dashboardAggregationRepository struct {
  15. sql sqlExecutor
  16. clock func() time.Time
  17. }
  18. const usageLogsCleanupBatchSize = 10000
  19. const usageBillingDedupCleanupBatchSize = 10000
  20. // NewDashboardAggregationRepository 创建仪表盘预聚合仓储。
  21. func NewDashboardAggregationRepository(sqlDB *sql.DB) service.DashboardAggregationRepository {
  22. if sqlDB == nil {
  23. return nil
  24. }
  25. if !isPostgresDriver(sqlDB) {
  26. log.Printf("[DashboardAggregation] 检测到非 PostgreSQL 驱动,已自动禁用预聚合")
  27. return nil
  28. }
  29. return newDashboardAggregationRepositoryWithSQL(sqlDB)
  30. }
  31. func newDashboardAggregationRepositoryWithSQL(sqlq sqlExecutor) *dashboardAggregationRepository {
  32. return &dashboardAggregationRepository{sql: sqlq, clock: time.Now}
  33. }
  34. func (r *dashboardAggregationRepository) now() time.Time {
  35. if r.clock != nil {
  36. return r.clock()
  37. }
  38. return time.Now()
  39. }
  40. func isPostgresDriver(db *sql.DB) bool {
  41. if db == nil {
  42. return false
  43. }
  44. _, ok := db.Driver().(*pq.Driver)
  45. return ok
  46. }
  47. func (r *dashboardAggregationRepository) AggregateRange(ctx context.Context, start, end time.Time) error {
  48. if r == nil || r.sql == nil {
  49. return nil
  50. }
  51. loc := timezone.Location()
  52. startLocal := start.In(loc)
  53. endLocal := end.In(loc)
  54. if !endLocal.After(startLocal) {
  55. return nil
  56. }
  57. hourStart := startLocal.Truncate(time.Hour)
  58. hourEnd := endLocal.Truncate(time.Hour)
  59. if endLocal.After(hourEnd) {
  60. hourEnd = hourEnd.Add(time.Hour)
  61. }
  62. dayStart := truncateToDay(startLocal)
  63. dayEnd := truncateToDay(endLocal)
  64. if endLocal.After(dayEnd) {
  65. dayEnd = dayEnd.Add(24 * time.Hour)
  66. }
  67. if db, ok := r.sql.(*sql.DB); ok {
  68. tx, err := db.BeginTx(ctx, nil)
  69. if err != nil {
  70. return err
  71. }
  72. txRepo := newDashboardAggregationRepositoryWithSQL(tx)
  73. if err := txRepo.aggregateRangeInTx(ctx, hourStart, hourEnd, dayStart, dayEnd); err != nil {
  74. _ = tx.Rollback()
  75. return err
  76. }
  77. return tx.Commit()
  78. }
  79. return r.aggregateRangeInTx(ctx, hourStart, hourEnd, dayStart, dayEnd)
  80. }
  81. func (r *dashboardAggregationRepository) aggregateRangeInTx(ctx context.Context, hourStart, hourEnd, dayStart, dayEnd time.Time) error {
  82. // 以桶边界聚合,允许覆盖 end 所在桶的剩余区间。
  83. if err := r.insertHourlyActiveUsers(ctx, hourStart, hourEnd); err != nil {
  84. return err
  85. }
  86. if err := r.insertDailyActiveUsers(ctx, hourStart, hourEnd); err != nil {
  87. return err
  88. }
  89. if err := r.upsertHourlyAggregates(ctx, hourStart, hourEnd); err != nil {
  90. return err
  91. }
  92. if err := r.upsertDailyAggregates(ctx, dayStart, dayEnd); err != nil {
  93. return err
  94. }
  95. return nil
  96. }
  97. func (r *dashboardAggregationRepository) RecomputeRange(ctx context.Context, start, end time.Time) error {
  98. if r == nil || r.sql == nil {
  99. return nil
  100. }
  101. loc := timezone.Location()
  102. startLocal := start.In(loc)
  103. endLocal := end.In(loc)
  104. if !endLocal.After(startLocal) {
  105. return nil
  106. }
  107. hourStart := startLocal.Truncate(time.Hour)
  108. hourEnd := endLocal.Truncate(time.Hour)
  109. if endLocal.After(hourEnd) {
  110. hourEnd = hourEnd.Add(time.Hour)
  111. }
  112. dayStart := truncateToDay(startLocal)
  113. dayEnd := truncateToDay(endLocal)
  114. if endLocal.After(dayEnd) {
  115. dayEnd = dayEnd.Add(24 * time.Hour)
  116. }
  117. // 尽量使用事务保证范围内的一致性(允许在非 *sql.DB 的情况下退化为非事务执行)。
  118. if db, ok := r.sql.(*sql.DB); ok {
  119. tx, err := db.BeginTx(ctx, nil)
  120. if err != nil {
  121. return err
  122. }
  123. if err := lockGroupUsageRollupState(ctx, tx); err != nil {
  124. _ = tx.Rollback()
  125. return err
  126. }
  127. if err := invalidateGroupUsageRollupsAt(ctx, tx, start); err != nil {
  128. _ = tx.Rollback()
  129. return err
  130. }
  131. txRepo := newDashboardAggregationRepositoryWithSQL(tx)
  132. if err := txRepo.recomputeRangeInTx(ctx, hourStart, hourEnd, dayStart, dayEnd); err != nil {
  133. _ = tx.Rollback()
  134. return err
  135. }
  136. if err := txRepo.syncGroupUsageRollupsInTx(ctx, service.GroupUsageTodayStart(r.now())); err != nil {
  137. _ = tx.Rollback()
  138. return err
  139. }
  140. return tx.Commit()
  141. }
  142. return r.recomputeRangeInTx(ctx, hourStart, hourEnd, dayStart, dayEnd)
  143. }
  144. func (r *dashboardAggregationRepository) recomputeRangeInTx(ctx context.Context, hourStart, hourEnd, dayStart, dayEnd time.Time) error {
  145. // 先清空范围内桶,再重建(避免仅增量插入导致活跃用户等指标无法回退)。
  146. if _, err := r.sql.ExecContext(ctx, "DELETE FROM usage_dashboard_hourly WHERE bucket_start >= $1 AND bucket_start < $2", hourStart, hourEnd); err != nil {
  147. return err
  148. }
  149. if _, err := r.sql.ExecContext(ctx, "DELETE FROM usage_dashboard_hourly_users WHERE bucket_start >= $1 AND bucket_start < $2", hourStart, hourEnd); err != nil {
  150. return err
  151. }
  152. if _, err := r.sql.ExecContext(ctx, "DELETE FROM usage_dashboard_daily WHERE bucket_date >= $1::date AND bucket_date < $2::date", dayStart, dayEnd); err != nil {
  153. return err
  154. }
  155. if _, err := r.sql.ExecContext(ctx, "DELETE FROM usage_dashboard_daily_users WHERE bucket_date >= $1::date AND bucket_date < $2::date", dayStart, dayEnd); err != nil {
  156. return err
  157. }
  158. if err := r.insertHourlyActiveUsers(ctx, hourStart, hourEnd); err != nil {
  159. return err
  160. }
  161. if err := r.insertDailyActiveUsers(ctx, hourStart, hourEnd); err != nil {
  162. return err
  163. }
  164. if err := r.upsertHourlyAggregates(ctx, hourStart, hourEnd); err != nil {
  165. return err
  166. }
  167. if err := r.upsertDailyAggregates(ctx, dayStart, dayEnd); err != nil {
  168. return err
  169. }
  170. return nil
  171. }
  172. func (r *dashboardAggregationRepository) GetAggregationWatermark(ctx context.Context) (time.Time, error) {
  173. var ts time.Time
  174. query := "SELECT last_aggregated_at FROM usage_dashboard_aggregation_watermark WHERE id = 1"
  175. if err := scanSingleRow(ctx, r.sql, query, nil, &ts); err != nil {
  176. if err == sql.ErrNoRows {
  177. return time.Unix(0, 0).UTC(), nil
  178. }
  179. return time.Time{}, err
  180. }
  181. return ts.UTC(), nil
  182. }
  183. func (r *dashboardAggregationRepository) UpdateAggregationWatermark(ctx context.Context, aggregatedAt time.Time) error {
  184. query := `
  185. INSERT INTO usage_dashboard_aggregation_watermark (id, last_aggregated_at, updated_at)
  186. VALUES (1, $1, NOW())
  187. ON CONFLICT (id)
  188. DO UPDATE SET last_aggregated_at = EXCLUDED.last_aggregated_at, updated_at = EXCLUDED.updated_at
  189. `
  190. _, err := r.sql.ExecContext(ctx, query, aggregatedAt.UTC())
  191. return err
  192. }
  193. func (r *dashboardAggregationRepository) CleanupAggregates(ctx context.Context, hourlyCutoff, dailyCutoff time.Time) error {
  194. hourlyCutoffUTC := hourlyCutoff.UTC()
  195. dailyCutoffUTC := dailyCutoff.UTC()
  196. if _, err := r.sql.ExecContext(ctx, "DELETE FROM usage_dashboard_hourly WHERE bucket_start < $1", hourlyCutoffUTC); err != nil {
  197. return err
  198. }
  199. if _, err := r.sql.ExecContext(ctx, "DELETE FROM usage_dashboard_hourly_users WHERE bucket_start < $1", hourlyCutoffUTC); err != nil {
  200. return err
  201. }
  202. if _, err := r.sql.ExecContext(ctx, "DELETE FROM usage_dashboard_daily WHERE bucket_date < $1::date", dailyCutoffUTC); err != nil {
  203. return err
  204. }
  205. if _, err := r.sql.ExecContext(ctx, "DELETE FROM usage_dashboard_daily_users WHERE bucket_date < $1::date", dailyCutoffUTC); err != nil {
  206. return err
  207. }
  208. return nil
  209. }
  210. func (r *dashboardAggregationRepository) CleanupUsageLogs(ctx context.Context, cutoff time.Time) error {
  211. isPartitioned, err := r.isUsageLogsPartitioned(ctx)
  212. if err != nil {
  213. return err
  214. }
  215. if isPartitioned {
  216. if err := r.dropUsageLogsPartitions(ctx, cutoff); err != nil {
  217. return err
  218. }
  219. } else if err := r.cleanupUsageLogsBatches(ctx, cutoff); err != nil {
  220. return err
  221. }
  222. return r.SyncGroupUsageRollups(ctx, service.GroupUsageTodayStart(r.now()))
  223. }
  224. func (r *dashboardAggregationRepository) cleanupUsageLogsBatches(ctx context.Context, cutoff time.Time) error {
  225. db, transactional := r.sql.(*sql.DB)
  226. for {
  227. if transactional {
  228. affected, err := cleanupUsageLogsBatchWithRollupInvalidation(ctx, db, cutoff)
  229. if err != nil {
  230. return err
  231. }
  232. if affected < usageLogsCleanupBatchSize {
  233. return nil
  234. }
  235. continue
  236. }
  237. res, err := r.sql.ExecContext(ctx, `
  238. WITH victims AS (
  239. SELECT ctid
  240. FROM usage_logs
  241. WHERE created_at < $1
  242. ORDER BY created_at ASC, id ASC
  243. LIMIT $2
  244. )
  245. DELETE FROM usage_logs
  246. WHERE ctid IN (SELECT ctid FROM victims)
  247. `, cutoff.UTC(), usageLogsCleanupBatchSize)
  248. if err != nil {
  249. return err
  250. }
  251. affected, err := res.RowsAffected()
  252. if err != nil {
  253. return err
  254. }
  255. if affected < usageLogsCleanupBatchSize {
  256. return nil
  257. }
  258. }
  259. }
  260. func cleanupUsageLogsBatchWithRollupInvalidation(ctx context.Context, db *sql.DB, cutoff time.Time) (int64, error) {
  261. tx, err := db.BeginTx(ctx, nil)
  262. if err != nil {
  263. return 0, err
  264. }
  265. rollback := func(err error) (int64, error) {
  266. _ = tx.Rollback()
  267. return 0, err
  268. }
  269. if err := lockGroupUsageRollupState(ctx, tx); err != nil {
  270. return rollback(err)
  271. }
  272. rows, err := tx.QueryContext(ctx, `
  273. WITH victims AS (
  274. SELECT ctid
  275. FROM usage_logs
  276. WHERE created_at < $1
  277. ORDER BY created_at ASC, id ASC
  278. LIMIT $2
  279. )
  280. DELETE FROM usage_logs
  281. WHERE ctid IN (SELECT ctid FROM victims)
  282. RETURNING created_at
  283. `, cutoff.UTC(), usageLogsCleanupBatchSize)
  284. if err != nil {
  285. return rollback(err)
  286. }
  287. var affected int64
  288. var earliestDeletedAt time.Time
  289. for rows.Next() {
  290. var deletedAt time.Time
  291. if err := rows.Scan(&deletedAt); err != nil {
  292. _ = rows.Close()
  293. return rollback(err)
  294. }
  295. affected++
  296. if earliestDeletedAt.IsZero() || deletedAt.Before(earliestDeletedAt) {
  297. earliestDeletedAt = deletedAt
  298. }
  299. }
  300. if err := rows.Err(); err != nil {
  301. _ = rows.Close()
  302. return rollback(err)
  303. }
  304. if err := rows.Close(); err != nil {
  305. return rollback(err)
  306. }
  307. if affected > 0 {
  308. if err := invalidateGroupUsageRollupsAt(ctx, tx, earliestDeletedAt); err != nil {
  309. return rollback(err)
  310. }
  311. }
  312. if err := tx.Commit(); err != nil {
  313. return 0, err
  314. }
  315. return affected, nil
  316. }
  317. func (r *dashboardAggregationRepository) CleanupUsageBillingDedup(ctx context.Context, cutoff time.Time) error {
  318. for {
  319. res, err := r.sql.ExecContext(ctx, `
  320. WITH victims AS (
  321. SELECT ctid, request_id, api_key_id, request_fingerprint, created_at
  322. FROM usage_billing_dedup
  323. WHERE created_at < $1
  324. LIMIT $2
  325. ), archived AS (
  326. INSERT INTO usage_billing_dedup_archive (request_id, api_key_id, request_fingerprint, created_at)
  327. SELECT request_id, api_key_id, request_fingerprint, created_at
  328. FROM victims
  329. ON CONFLICT (request_id, api_key_id) DO NOTHING
  330. )
  331. DELETE FROM usage_billing_dedup
  332. WHERE ctid IN (SELECT ctid FROM victims)
  333. `, cutoff.UTC(), usageBillingDedupCleanupBatchSize)
  334. if err != nil {
  335. return err
  336. }
  337. affected, err := res.RowsAffected()
  338. if err != nil {
  339. return err
  340. }
  341. if affected < usageBillingDedupCleanupBatchSize {
  342. return nil
  343. }
  344. }
  345. }
  346. func (r *dashboardAggregationRepository) EnsureUsageLogsPartitions(ctx context.Context, now time.Time) error {
  347. isPartitioned, err := r.isUsageLogsPartitioned(ctx)
  348. if err != nil || !isPartitioned {
  349. return err
  350. }
  351. monthStart := truncateToMonthUTC(now)
  352. prevMonth := monthStart.AddDate(0, -1, 0)
  353. nextMonth := monthStart.AddDate(0, 1, 0)
  354. for _, m := range []time.Time{prevMonth, monthStart, nextMonth} {
  355. if err := r.createUsageLogsPartition(ctx, m); err != nil {
  356. return err
  357. }
  358. }
  359. return nil
  360. }
  361. func (r *dashboardAggregationRepository) insertHourlyActiveUsers(ctx context.Context, start, end time.Time) error {
  362. tzName := timezone.Name()
  363. query := `
  364. INSERT INTO usage_dashboard_hourly_users (bucket_start, user_id)
  365. SELECT DISTINCT
  366. date_trunc('hour', created_at AT TIME ZONE $3) AT TIME ZONE $3 AS bucket_start,
  367. user_id
  368. FROM usage_logs
  369. WHERE created_at >= $1 AND created_at < $2
  370. ON CONFLICT DO NOTHING
  371. `
  372. _, err := r.sql.ExecContext(ctx, query, start, end, tzName)
  373. return err
  374. }
  375. func (r *dashboardAggregationRepository) insertDailyActiveUsers(ctx context.Context, start, end time.Time) error {
  376. tzName := timezone.Name()
  377. query := `
  378. INSERT INTO usage_dashboard_daily_users (bucket_date, user_id)
  379. SELECT DISTINCT
  380. (bucket_start AT TIME ZONE $3)::date AS bucket_date,
  381. user_id
  382. FROM usage_dashboard_hourly_users
  383. WHERE bucket_start >= $1 AND bucket_start < $2
  384. ON CONFLICT DO NOTHING
  385. `
  386. _, err := r.sql.ExecContext(ctx, query, start, end, tzName)
  387. return err
  388. }
  389. func (r *dashboardAggregationRepository) upsertHourlyAggregates(ctx context.Context, start, end time.Time) error {
  390. tzName := timezone.Name()
  391. query := `
  392. WITH hourly AS (
  393. SELECT
  394. date_trunc('hour', created_at AT TIME ZONE $3) AT TIME ZONE $3 AS bucket_start,
  395. COUNT(*) AS total_requests,
  396. COALESCE(SUM(input_tokens), 0) AS input_tokens,
  397. COALESCE(SUM(output_tokens), 0) AS output_tokens,
  398. COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens,
  399. COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
  400. COALESCE(SUM(total_cost), 0) AS total_cost,
  401. COALESCE(SUM(actual_cost), 0) AS actual_cost,
  402. COALESCE(SUM(COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1)), 0) AS account_cost,
  403. COALESCE(SUM(COALESCE(duration_ms, 0)), 0) AS total_duration_ms
  404. FROM usage_logs
  405. WHERE created_at >= $1 AND created_at < $2
  406. GROUP BY 1
  407. ),
  408. user_counts AS (
  409. SELECT bucket_start, COUNT(*) AS active_users
  410. FROM usage_dashboard_hourly_users
  411. WHERE bucket_start >= $1 AND bucket_start < $2
  412. GROUP BY bucket_start
  413. )
  414. INSERT INTO usage_dashboard_hourly (
  415. bucket_start,
  416. total_requests,
  417. input_tokens,
  418. output_tokens,
  419. cache_creation_tokens,
  420. cache_read_tokens,
  421. total_cost,
  422. actual_cost,
  423. account_cost,
  424. total_duration_ms,
  425. active_users,
  426. computed_at
  427. )
  428. SELECT
  429. hourly.bucket_start,
  430. hourly.total_requests,
  431. hourly.input_tokens,
  432. hourly.output_tokens,
  433. hourly.cache_creation_tokens,
  434. hourly.cache_read_tokens,
  435. hourly.total_cost,
  436. hourly.actual_cost,
  437. hourly.account_cost,
  438. hourly.total_duration_ms,
  439. COALESCE(user_counts.active_users, 0) AS active_users,
  440. NOW()
  441. FROM hourly
  442. LEFT JOIN user_counts ON user_counts.bucket_start = hourly.bucket_start
  443. ON CONFLICT (bucket_start)
  444. DO UPDATE SET
  445. total_requests = EXCLUDED.total_requests,
  446. input_tokens = EXCLUDED.input_tokens,
  447. output_tokens = EXCLUDED.output_tokens,
  448. cache_creation_tokens = EXCLUDED.cache_creation_tokens,
  449. cache_read_tokens = EXCLUDED.cache_read_tokens,
  450. total_cost = EXCLUDED.total_cost,
  451. actual_cost = EXCLUDED.actual_cost,
  452. account_cost = EXCLUDED.account_cost,
  453. total_duration_ms = EXCLUDED.total_duration_ms,
  454. active_users = EXCLUDED.active_users,
  455. computed_at = EXCLUDED.computed_at
  456. `
  457. _, err := r.sql.ExecContext(ctx, query, start, end, tzName)
  458. return err
  459. }
  460. func (r *dashboardAggregationRepository) upsertDailyAggregates(ctx context.Context, start, end time.Time) error {
  461. tzName := timezone.Name()
  462. query := `
  463. WITH daily AS (
  464. SELECT
  465. (bucket_start AT TIME ZONE $5)::date AS bucket_date,
  466. COALESCE(SUM(total_requests), 0) AS total_requests,
  467. COALESCE(SUM(input_tokens), 0) AS input_tokens,
  468. COALESCE(SUM(output_tokens), 0) AS output_tokens,
  469. COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens,
  470. COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
  471. COALESCE(SUM(total_cost), 0) AS total_cost,
  472. COALESCE(SUM(actual_cost), 0) AS actual_cost,
  473. COALESCE(SUM(account_cost), 0) AS account_cost,
  474. COALESCE(SUM(total_duration_ms), 0) AS total_duration_ms
  475. FROM usage_dashboard_hourly
  476. WHERE bucket_start >= $1 AND bucket_start < $2
  477. GROUP BY (bucket_start AT TIME ZONE $5)::date
  478. ),
  479. user_counts AS (
  480. SELECT bucket_date, COUNT(*) AS active_users
  481. FROM usage_dashboard_daily_users
  482. WHERE bucket_date >= $3::date AND bucket_date < $4::date
  483. GROUP BY bucket_date
  484. )
  485. INSERT INTO usage_dashboard_daily (
  486. bucket_date,
  487. total_requests,
  488. input_tokens,
  489. output_tokens,
  490. cache_creation_tokens,
  491. cache_read_tokens,
  492. total_cost,
  493. actual_cost,
  494. account_cost,
  495. total_duration_ms,
  496. active_users,
  497. computed_at
  498. )
  499. SELECT
  500. daily.bucket_date,
  501. daily.total_requests,
  502. daily.input_tokens,
  503. daily.output_tokens,
  504. daily.cache_creation_tokens,
  505. daily.cache_read_tokens,
  506. daily.total_cost,
  507. daily.actual_cost,
  508. daily.account_cost,
  509. daily.total_duration_ms,
  510. COALESCE(user_counts.active_users, 0) AS active_users,
  511. NOW()
  512. FROM daily
  513. LEFT JOIN user_counts ON user_counts.bucket_date = daily.bucket_date
  514. ON CONFLICT (bucket_date)
  515. DO UPDATE SET
  516. total_requests = EXCLUDED.total_requests,
  517. input_tokens = EXCLUDED.input_tokens,
  518. output_tokens = EXCLUDED.output_tokens,
  519. cache_creation_tokens = EXCLUDED.cache_creation_tokens,
  520. cache_read_tokens = EXCLUDED.cache_read_tokens,
  521. total_cost = EXCLUDED.total_cost,
  522. actual_cost = EXCLUDED.actual_cost,
  523. account_cost = EXCLUDED.account_cost,
  524. total_duration_ms = EXCLUDED.total_duration_ms,
  525. active_users = EXCLUDED.active_users,
  526. computed_at = EXCLUDED.computed_at
  527. `
  528. _, err := r.sql.ExecContext(ctx, query, start, end, start, end, tzName)
  529. return err
  530. }
  531. func (r *dashboardAggregationRepository) isUsageLogsPartitioned(ctx context.Context) (bool, error) {
  532. query := `
  533. SELECT EXISTS(
  534. SELECT 1
  535. FROM pg_partitioned_table pt
  536. JOIN pg_class c ON c.oid = pt.partrelid
  537. WHERE c.relname = 'usage_logs'
  538. )
  539. `
  540. var partitioned bool
  541. if err := scanSingleRow(ctx, r.sql, query, nil, &partitioned); err != nil {
  542. return false, err
  543. }
  544. return partitioned, nil
  545. }
  546. func (r *dashboardAggregationRepository) dropUsageLogsPartitions(ctx context.Context, cutoff time.Time) error {
  547. rows, err := r.sql.QueryContext(ctx, `
  548. SELECT c.relname
  549. FROM pg_inherits
  550. JOIN pg_class c ON c.oid = pg_inherits.inhrelid
  551. JOIN pg_class p ON p.oid = pg_inherits.inhparent
  552. WHERE p.relname = 'usage_logs'
  553. `)
  554. if err != nil {
  555. return err
  556. }
  557. cutoffMonth := truncateToMonthUTC(cutoff)
  558. type usageLogsPartition struct {
  559. name string
  560. month time.Time
  561. }
  562. partitions := make([]usageLogsPartition, 0)
  563. for rows.Next() {
  564. var name string
  565. if err := rows.Scan(&name); err != nil {
  566. _ = rows.Close()
  567. return err
  568. }
  569. if !strings.HasPrefix(name, "usage_logs_") {
  570. continue
  571. }
  572. suffix := strings.TrimPrefix(name, "usage_logs_")
  573. month, err := time.Parse("200601", suffix)
  574. if err != nil {
  575. continue
  576. }
  577. month = month.UTC()
  578. if month.Before(cutoffMonth) {
  579. partitions = append(partitions, usageLogsPartition{name: name, month: month})
  580. }
  581. }
  582. if err := rows.Err(); err != nil {
  583. _ = rows.Close()
  584. return err
  585. }
  586. if err := rows.Close(); err != nil {
  587. return err
  588. }
  589. sort.Slice(partitions, func(i, j int) bool {
  590. return partitions[i].month.Before(partitions[j].month)
  591. })
  592. if db, ok := r.sql.(*sql.DB); ok {
  593. for _, partition := range partitions {
  594. if err := dropUsageLogsPartitionWithRollupInvalidation(ctx, db, partition.name, partition.month); err != nil {
  595. return err
  596. }
  597. }
  598. return nil
  599. }
  600. for _, partition := range partitions {
  601. if _, err := r.sql.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", pq.QuoteIdentifier(partition.name))); err != nil {
  602. return err
  603. }
  604. }
  605. return nil
  606. }
  607. func dropUsageLogsPartitionWithRollupInvalidation(ctx context.Context, db *sql.DB, name string, monthStart time.Time) error {
  608. tx, err := db.BeginTx(ctx, nil)
  609. if err != nil {
  610. return err
  611. }
  612. rollback := func(err error) error {
  613. _ = tx.Rollback()
  614. return err
  615. }
  616. if err := lockGroupUsageRollupState(ctx, tx); err != nil {
  617. return rollback(err)
  618. }
  619. if err := invalidateGroupUsageRollupsAt(ctx, tx, monthStart); err != nil {
  620. return rollback(err)
  621. }
  622. if _, err := tx.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", pq.QuoteIdentifier(name))); err != nil {
  623. return rollback(err)
  624. }
  625. return tx.Commit()
  626. }
  627. func (r *dashboardAggregationRepository) createUsageLogsPartition(ctx context.Context, month time.Time) error {
  628. monthStart := truncateToMonthUTC(month)
  629. nextMonth := monthStart.AddDate(0, 1, 0)
  630. name := fmt.Sprintf("usage_logs_%s", monthStart.Format("200601"))
  631. query := fmt.Sprintf(
  632. "CREATE TABLE IF NOT EXISTS %s PARTITION OF usage_logs FOR VALUES FROM (%s) TO (%s)",
  633. pq.QuoteIdentifier(name),
  634. pq.QuoteLiteral(monthStart.Format("2006-01-02")),
  635. pq.QuoteLiteral(nextMonth.Format("2006-01-02")),
  636. )
  637. _, err := r.sql.ExecContext(ctx, query)
  638. return err
  639. }
  640. func truncateToDay(t time.Time) time.Time {
  641. return timezone.StartOfDay(t)
  642. }
  643. func truncateToMonthUTC(t time.Time) time.Time {
  644. t = t.UTC()
  645. return time.Date(t.Year(), t.Month(), 1, 0, 0, 0, 0, time.UTC)
  646. }