mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-07 12:57:57 +08:00
Merge pull request #5649 from lyen1688/feat/group-usage-daily-rollups
feat: 优化分组用量统计
This commit is contained in:
Wesley Liddick
@@ -32,7 +32,7 @@ jobs:
|
||||
cache-dependency-path: backend/go.sum
|
||||
- name: Verify Go version
|
||||
run: |
|
||||
go version | grep -q 'go1.26.5'
|
||||
go version | grep -q 'go1.26.6'
|
||||
- name: Unit tests
|
||||
working-directory: backend
|
||||
run: make test-unit
|
||||
@@ -72,7 +72,7 @@ jobs:
|
||||
cache-dependency-path: backend/go.sum
|
||||
- name: Verify Go version
|
||||
run: |
|
||||
go version | grep -q 'go1.26.5'
|
||||
go version | grep -q 'go1.26.6'
|
||||
- name: golangci-lint
|
||||
uses: golangci/golangci-lint-action@v9
|
||||
with:
|
||||
|
||||
@@ -115,7 +115,7 @@ jobs:
|
||||
|
||||
- name: Verify Go version
|
||||
run: |
|
||||
go version | grep -q 'go1.26.5'
|
||||
go version | grep -q 'go1.26.6'
|
||||
|
||||
# Docker setup for GoReleaser
|
||||
- name: Set up QEMU
|
||||
|
||||
@@ -23,7 +23,7 @@ jobs:
|
||||
cache-dependency-path: backend/go.sum
|
||||
- name: Verify Go version
|
||||
run: |
|
||||
go version | grep -q 'go1.26.5'
|
||||
go version | grep -q 'go1.26.6'
|
||||
- name: Run govulncheck
|
||||
working-directory: backend
|
||||
run: |
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
module github.com/Wei-Shaw/sub2api
|
||||
|
||||
go 1.26.5
|
||||
go 1.26.6
|
||||
|
||||
require (
|
||||
entgo.io/ent v0.14.5
|
||||
|
||||
@@ -1703,6 +1703,10 @@ func load(allowMissingJWTSecret bool) (*Config, error) {
|
||||
// 环境变量支持
|
||||
viper.AutomaticEnv()
|
||||
viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
|
||||
if tz, ok := os.LookupEnv("TZ"); ok && strings.TrimSpace(tz) != "" {
|
||||
// AutomaticEnv 会先把 timezone 映射到 TIMEZONE；显式 Set 保证标准 TZ 变量优先。
|
||||
viper.Set("timezone", strings.TrimSpace(tz))
|
||||
}
|
||||
if err := viper.BindEnv("server.enable_server_timing", "ENABLE_SERVER_TIMING"); err != nil {
|
||||
return nil, fmt.Errorf("bind ENABLE_SERVER_TIMING: %w", err)
|
||||
}
|
||||
|
||||
@@ -23,6 +23,38 @@ func resetViperWithJWTSecret(t *testing.T) {
|
||||
t.Setenv("JWT_SECRET", strings.Repeat("x", 32))
|
||||
}
|
||||
|
||||
func TestLoadTimezonePrecedence(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
fileTimezone string
|
||||
timezoneEnv string
|
||||
tzEnv string
|
||||
want string
|
||||
}{
|
||||
{name: "default", want: "Asia/Shanghai"},
|
||||
{name: "config_file", fileTimezone: "Europe/London", want: "Europe/London"},
|
||||
{name: "timezone_env", fileTimezone: "Europe/London", timezoneEnv: "UTC", want: "UTC"},
|
||||
{name: "tz_env", fileTimezone: "Europe/London", timezoneEnv: "UTC", tzEnv: "America/New_York", want: "America/New_York"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("TIMEZONE", tt.timezoneEnv)
|
||||
t.Setenv("TZ", tt.tzEnv)
|
||||
if tt.fileTimezone != "" {
|
||||
configFile := filepath.Join(t.TempDir(), "config.yaml")
|
||||
require.NoError(t, os.WriteFile(configFile, []byte("timezone: "+tt.fileTimezone+"\n"), 0o600))
|
||||
t.Setenv("CONFIG_FILE", configFile)
|
||||
}
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tt.want, cfg.Timezone)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadServerTimingConfig(t *testing.T) {
|
||||
t.Run("disabled by default", func(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
|
||||
@@ -8,11 +8,11 @@ import (
|
||||
"log/slog"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/Wei-Shaw/sub2api/internal/platform/liveattestation"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
@@ -734,12 +734,10 @@ func (h *GroupHandler) GetStats(c *gin.Context) {
|
||||
_ = groupID // TODO: implement actual stats
|
||||
}
|
||||
|
||||
// GetUsageSummary returns today's and cumulative cost for all groups.
|
||||
// GET /api/v1/admin/groups/usage-summary?timezone=Asia/Shanghai
|
||||
// GetUsageSummary returns today's, yesterday's, and cumulative cost for all groups.
|
||||
// GET /api/v1/admin/groups/usage-summary
|
||||
func (h *GroupHandler) GetUsageSummary(c *gin.Context) {
|
||||
userTZ := c.Query("timezone")
|
||||
now := timezone.NowInUserLocation(userTZ)
|
||||
todayStart := timezone.StartOfDayInUserLocation(now, userTZ)
|
||||
todayStart := service.GroupUsageTodayStart(time.Now())
|
||||
|
||||
results, err := h.dashboardService.GetGroupUsageSummary(c.Request.Context(), todayStart)
|
||||
if err != nil {
|
||||
|
||||
@@ -115,11 +115,12 @@ type EndpointStat struct {
|
||||
ActualCost float64 `json:"actual_cost"` // 实际扣除
|
||||
}
|
||||
|
||||
// GroupUsageSummary represents today's and cumulative cost for a single group.
|
||||
// GroupUsageSummary represents today's, yesterday's, and cumulative cost for a single group.
|
||||
type GroupUsageSummary struct {
|
||||
GroupID int64 `json:"group_id"`
|
||||
TodayCost float64 `json:"today_cost"`
|
||||
TotalCost float64 `json:"total_cost"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
TodayCost float64 `json:"today_cost"`
|
||||
YesterdayCost float64 `json:"yesterday_cost"`
|
||||
TotalCost float64 `json:"total_cost"`
|
||||
}
|
||||
|
||||
// GroupStat represents usage statistics for a single group
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
func (r *usageLogRepository) getAllGroupUsageSummaryFromRollups(ctx context.Context, todayStart time.Time) (results []usagestats.GroupUsageSummary, err error) {
|
||||
todayStart = service.GroupUsageTodayStart(todayStart)
|
||||
yesterdayStart := service.GroupUsageYesterdayStart(todayStart)
|
||||
timezoneName := service.GroupUsageTimezoneName()
|
||||
todayDate := service.GroupUsageDate(todayStart)
|
||||
yesterdayDate := service.GroupUsageDate(yesterdayStart)
|
||||
|
||||
const query = `
|
||||
WITH state_values AS (
|
||||
SELECT
|
||||
COUNT(*) = 1
|
||||
AND MAX(timezone_name) = $3
|
||||
AND MAX(closed_before) <= $4::date AS valid,
|
||||
MAX(closed_before) AS closed_before,
|
||||
MAX(retained_from) AS retained_from
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
),
|
||||
state AS (
|
||||
SELECT
|
||||
CASE WHEN valid THEN closed_before ELSE DATE '1970-01-01' END AS closed_before,
|
||||
CASE WHEN valid THEN retained_from ELSE TIMESTAMPTZ '1970-01-01 00:00:00+00' END AS retained_from,
|
||||
CASE
|
||||
WHEN valid THEN closed_before::timestamp AT TIME ZONE $3::text
|
||||
ELSE TIMESTAMPTZ '1970-01-01 00:00:00+00'
|
||||
END AS tail_start,
|
||||
valid
|
||||
FROM state_values
|
||||
),
|
||||
historical AS (
|
||||
SELECT
|
||||
rollup.group_id,
|
||||
COALESCE(SUM(rollup.actual_cost), 0) AS actual_cost,
|
||||
COALESCE(SUM(rollup.actual_cost) FILTER (
|
||||
WHERE rollup.bucket_date = $5::date
|
||||
), 0) AS yesterday_cost
|
||||
FROM usage_group_daily_rollups rollup
|
||||
CROSS JOIN state
|
||||
WHERE state.valid
|
||||
AND rollup.bucket_date >= (state.retained_from AT TIME ZONE $3::text)::date
|
||||
AND rollup.bucket_date < state.closed_before
|
||||
GROUP BY rollup.group_id
|
||||
),
|
||||
tail AS (
|
||||
SELECT
|
||||
ul.group_id,
|
||||
COALESCE(SUM(ul.actual_cost), 0) AS actual_cost,
|
||||
COALESCE(SUM(ul.actual_cost) FILTER (WHERE ul.created_at >= $1), 0) AS today_cost,
|
||||
COALESCE(SUM(ul.actual_cost) FILTER (
|
||||
WHERE ul.created_at >= $2
|
||||
AND ul.created_at < $1
|
||||
), 0) AS yesterday_cost
|
||||
FROM usage_logs ul
|
||||
CROSS JOIN state
|
||||
WHERE ul.created_at >= state.tail_start
|
||||
GROUP BY ul.group_id
|
||||
)
|
||||
SELECT
|
||||
g.id AS group_id,
|
||||
COALESCE(historical.actual_cost, 0) + COALESCE(tail.actual_cost, 0) AS total_cost,
|
||||
COALESCE(tail.today_cost, 0) AS today_cost,
|
||||
COALESCE(historical.yesterday_cost, 0) + COALESCE(tail.yesterday_cost, 0) AS yesterday_cost
|
||||
FROM groups g
|
||||
LEFT JOIN historical ON historical.group_id = g.id
|
||||
LEFT JOIN tail ON tail.group_id = g.id
|
||||
ORDER BY g.id
|
||||
`
|
||||
|
||||
rows, err := r.sql.QueryContext(
|
||||
ctx,
|
||||
query,
|
||||
todayStart,
|
||||
yesterdayStart,
|
||||
timezoneName,
|
||||
todayDate,
|
||||
yesterdayDate,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() {
|
||||
if closeErr := rows.Close(); closeErr != nil && err == nil {
|
||||
err = closeErr
|
||||
results = nil
|
||||
}
|
||||
}()
|
||||
|
||||
results = make([]usagestats.GroupUsageSummary, 0)
|
||||
for rows.Next() {
|
||||
var row usagestats.GroupUsageSummary
|
||||
if err := rows.Scan(&row.GroupID, &row.TotalCost, &row.TodayCost, &row.YesterdayCost); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
results = append(results, row)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// SyncGroupUsageRollups 将服务端配置时区今日以前的用量发布为分组日桶。
|
||||
func (r *dashboardAggregationRepository) SyncGroupUsageRollups(ctx context.Context, todayStart time.Time) error {
|
||||
if r == nil || r.sql == nil {
|
||||
return nil
|
||||
}
|
||||
todayStart = service.GroupUsageTodayStart(todayStart)
|
||||
if db, ok := r.sql.(*sql.DB); ok {
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
txRepo := newDashboardAggregationRepositoryWithSQL(tx)
|
||||
if err := txRepo.syncGroupUsageRollupsInTx(ctx, todayStart); err != nil {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
return r.syncGroupUsageRollupsInTx(ctx, todayStart)
|
||||
}
|
||||
|
||||
func (r *dashboardAggregationRepository) syncGroupUsageRollupsInTx(ctx context.Context, todayStart time.Time) error {
|
||||
var closedBefore string
|
||||
var previousRetainedFrom time.Time
|
||||
var stateTimezoneName string
|
||||
if err := scanSingleRow(ctx, r.sql, `
|
||||
SELECT closed_before::text, retained_from, timezone_name
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
FOR UPDATE
|
||||
`, nil, &closedBefore, &previousRetainedFrom, &stateTimezoneName); err != nil {
|
||||
return fmt.Errorf("读取分组用量汇总水位: %w", err)
|
||||
}
|
||||
|
||||
todayDate := service.GroupUsageDate(todayStart)
|
||||
timezoneName := service.GroupUsageTimezoneName()
|
||||
timezoneChanged := stateTimezoneName != timezoneName
|
||||
var closedTime time.Time
|
||||
if !timezoneChanged {
|
||||
var err error
|
||||
closedTime, err = service.ParseGroupUsageDate(closedBefore)
|
||||
if err != nil {
|
||||
return fmt.Errorf("解析分组用量汇总水位 %q: %w", closedBefore, err)
|
||||
}
|
||||
todayDateTime, err := service.ParseGroupUsageDate(todayDate)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if closedTime.After(todayDateTime) {
|
||||
return fmt.Errorf("分组用量汇总水位位于未来: %s", closedBefore)
|
||||
}
|
||||
if closedBefore == todayDate {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
var earliest sql.NullTime
|
||||
if err := scanSingleRow(ctx, r.sql, "SELECT MIN(created_at) FROM usage_logs", nil, &earliest); err != nil {
|
||||
return fmt.Errorf("读取最早用量记录: %w", err)
|
||||
}
|
||||
retainedFrom := todayStart
|
||||
if earliest.Valid {
|
||||
retainedFrom = earliest.Time.UTC()
|
||||
}
|
||||
retainedDate := service.GroupUsageDate(retainedFrom)
|
||||
retainedDateTime, err := service.ParseGroupUsageDate(retainedDate)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rebuildStartDate := retainedDate
|
||||
if !timezoneChanged && closedTime.After(retainedDateTime) {
|
||||
rebuildStartDate = closedBefore
|
||||
}
|
||||
rebuildStart, err := service.ParseGroupUsageDate(rebuildStartDate)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if _, err := r.sql.ExecContext(ctx, `
|
||||
DELETE FROM usage_group_daily_rollups
|
||||
WHERE bucket_date < $1::date
|
||||
OR (bucket_date >= $2::date AND bucket_date < $3::date)
|
||||
OR bucket_date >= $3::date
|
||||
`, retainedDate, rebuildStartDate, todayDate); err != nil {
|
||||
return fmt.Errorf("清理分组用量日桶: %w", err)
|
||||
}
|
||||
|
||||
if _, err := r.sql.ExecContext(ctx, `
|
||||
INSERT INTO usage_group_daily_rollups (bucket_date, group_id, actual_cost, computed_at)
|
||||
SELECT
|
||||
(created_at AT TIME ZONE $3::text)::date AS bucket_date,
|
||||
group_id,
|
||||
COALESCE(SUM(actual_cost), 0) AS actual_cost,
|
||||
NOW()
|
||||
FROM usage_logs
|
||||
WHERE group_id IS NOT NULL
|
||||
AND created_at >= $1
|
||||
AND created_at < $2
|
||||
GROUP BY 1, 2
|
||||
ON CONFLICT (bucket_date, group_id)
|
||||
DO UPDATE SET
|
||||
actual_cost = EXCLUDED.actual_cost,
|
||||
computed_at = EXCLUDED.computed_at
|
||||
`, rebuildStart.UTC(), todayStart.UTC(), timezoneName); err != nil {
|
||||
return fmt.Errorf("重建分组用量日桶: %w", err)
|
||||
}
|
||||
|
||||
if _, err := r.sql.ExecContext(ctx, `
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = $1::date,
|
||||
retained_from = $2,
|
||||
timezone_name = $3,
|
||||
updated_at = NOW()
|
||||
WHERE id = 1
|
||||
`, todayDate, retainedFrom, timezoneName); err != nil {
|
||||
return fmt.Errorf("更新分组用量汇总水位: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func lockGroupUsageRollupState(ctx context.Context, tx *sql.Tx) error {
|
||||
var id int16
|
||||
if err := tx.QueryRowContext(ctx, `
|
||||
SELECT id
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
FOR UPDATE
|
||||
`).Scan(&id); err != nil {
|
||||
return fmt.Errorf("锁定分组用量汇总水位: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func invalidateGroupUsageRollupsAt(ctx context.Context, tx *sql.Tx, affectedAt time.Time) error {
|
||||
timezoneName := service.GroupUsageTimezoneName()
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = LEAST(
|
||||
closed_before,
|
||||
($1::timestamptz AT TIME ZONE $2::text)::date
|
||||
),
|
||||
updated_at = NOW()
|
||||
WHERE id = 1
|
||||
`, affectedAt.UTC(), timezoneName)
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
appTimezone "github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func useGroupUsageRepositoryTestTimezone(t *testing.T, name string) {
|
||||
t.Helper()
|
||||
|
||||
previousName := appTimezone.Name()
|
||||
require.NoError(t, appTimezone.Init(name))
|
||||
t.Cleanup(func() { require.NoError(t, appTimezone.Init(previousName)) })
|
||||
}
|
||||
@@ -0,0 +1,319 @@
|
||||
//go:build unit
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestDashboardAggregationRepositorySyncGroupUsageRollupsNoopsAtCurrentDate(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow("2026-08-14", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectCommit()
|
||||
|
||||
require.NoError(t, repo.SyncGroupUsageRollups(context.Background(), todayStart))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositorySyncGroupUsageRollupsRebuildsWhenTimezoneChanges(t *testing.T) {
|
||||
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
|
||||
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
|
||||
retainedFrom := time.Date(2026, 3, 1, 5, 0, 0, 0, time.UTC)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow("2026-03-09", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectQuery(`SELECT MIN\(created_at\) FROM usage_logs`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(retainedFrom))
|
||||
mock.ExpectExec(`DELETE FROM usage_group_daily_rollups`).
|
||||
WithArgs("2026-03-01", "2026-03-01", "2026-03-09").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`INSERT INTO usage_group_daily_rollups`).
|
||||
WithArgs(retainedFrom, todayStart, "America/New_York").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs("2026-03-09", retainedFrom, "America/New_York").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
require.NoError(t, repo.SyncGroupUsageRollups(context.Background(), todayStart))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositorySyncGroupUsageRollupsPublishesWatermarkLast(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
|
||||
retainedFrom := time.Date(2026, 5, 1, 3, 0, 0, 0, time.UTC)
|
||||
rebuildStart := time.Date(2026, 8, 12, 16, 0, 0, 0, time.UTC)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow("2026-08-13", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectQuery(`SELECT MIN\(created_at\) FROM usage_logs`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(retainedFrom))
|
||||
mock.ExpectExec(`DELETE FROM usage_group_daily_rollups`).
|
||||
WithArgs("2026-05-01", "2026-08-13", "2026-08-14").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`INSERT INTO usage_group_daily_rollups`).
|
||||
WithArgs(rebuildStart, todayStart, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 2))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs("2026-08-14", retainedFrom, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
require.NoError(t, repo.SyncGroupUsageRollups(context.Background(), todayStart))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositorySyncGroupUsageRollupsRejectsFutureWatermark(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow("2026-08-15", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectRollback()
|
||||
|
||||
err := repo.SyncGroupUsageRollups(context.Background(), todayStart)
|
||||
require.ErrorContains(t, err, "未来")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositoryRecomputeRangeInvalidatesGroupRollupsBeforeDashboardRebuild(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
start := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
|
||||
end := start.Add(time.Hour)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(start, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`DELETE FROM usage_dashboard_hourly`).
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
mock.ExpectRollback()
|
||||
|
||||
err := repo.RecomputeRange(context.Background(), start, end)
|
||||
require.ErrorIs(t, err, sql.ErrConnDone)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositoryRecomputeRangeRebuildsGroupRollupsBeforeCommit(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
start := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
|
||||
end := start.Add(time.Hour)
|
||||
fixedNow := time.Date(2026, 8, 14, 8, 0, 0, 0, time.UTC)
|
||||
repo.clock = func() time.Time { return fixedNow }
|
||||
todayStart := service.GroupUsageTodayStart(fixedNow)
|
||||
startDate := service.GroupUsageDate(start)
|
||||
rebuildStart, err := service.ParseGroupUsageDate(startDate)
|
||||
require.NoError(t, err)
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(start, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
for _, query := range []string{
|
||||
`DELETE FROM usage_dashboard_hourly WHERE`,
|
||||
`DELETE FROM usage_dashboard_hourly_users WHERE`,
|
||||
`DELETE FROM usage_dashboard_daily WHERE`,
|
||||
`DELETE FROM usage_dashboard_daily_users WHERE`,
|
||||
`INSERT INTO usage_dashboard_hourly_users`,
|
||||
`INSERT INTO usage_dashboard_daily_users`,
|
||||
`INSERT INTO usage_dashboard_hourly`,
|
||||
`INSERT INTO usage_dashboard_daily`,
|
||||
} {
|
||||
mock.ExpectExec(query).WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
}
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow(startDate, time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectQuery(`SELECT MIN\(created_at\) FROM usage_logs`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(start))
|
||||
mock.ExpectExec(`DELETE FROM usage_group_daily_rollups`).
|
||||
WithArgs(startDate, startDate, service.GroupUsageDate(todayStart)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`INSERT INTO usage_group_daily_rollups`).
|
||||
WithArgs(rebuildStart.UTC(), todayStart.UTC(), "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(service.GroupUsageDate(todayStart), start, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
require.NoError(t, repo.RecomputeRange(context.Background(), start, end))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositoryCleanupUsageLogsNonPartitionedInvalidatesEachBatchAndSyncs(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
cutoff := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)
|
||||
earliestDeletedAt := time.Date(2026, 5, 3, 2, 0, 0, 0, time.UTC)
|
||||
fixedNow := time.Date(2026, 8, 14, 8, 0, 0, 0, time.UTC)
|
||||
repo.clock = func() time.Time { return fixedNow }
|
||||
todayStart := service.GroupUsageTodayStart(fixedNow)
|
||||
|
||||
mock.ExpectQuery(`SELECT EXISTS`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false))
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
|
||||
WithArgs(cutoff, usageLogsCleanupBatchSize).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).
|
||||
AddRow(earliestDeletedAt.Add(time.Hour)).
|
||||
AddRow(earliestDeletedAt))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(earliestDeletedAt, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow(service.GroupUsageDate(todayStart), time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectCommit()
|
||||
|
||||
require.NoError(t, repo.CleanupUsageLogs(context.Background(), cutoff))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositoryCleanupUsageLogsPartitionedSortsAndInvalidatesEachDropBeforeSync(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
cutoff := time.Date(2026, 7, 18, 0, 0, 0, 0, time.UTC)
|
||||
aprilStart := time.Date(2026, 4, 1, 0, 0, 0, 0, time.UTC)
|
||||
juneStart := time.Date(2026, 6, 1, 0, 0, 0, 0, time.UTC)
|
||||
fixedNow := time.Date(2026, 8, 14, 8, 0, 0, 0, time.UTC)
|
||||
repo.clock = func() time.Time { return fixedNow }
|
||||
todayStart := service.GroupUsageTodayStart(fixedNow)
|
||||
|
||||
mock.ExpectQuery(`SELECT EXISTS`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
|
||||
mock.ExpectQuery(`SELECT c.relname`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"relname"}).
|
||||
AddRow("usage_logs_202606").
|
||||
AddRow("usage_logs_invalid").
|
||||
AddRow("usage_logs_202604").
|
||||
AddRow("usage_logs_202607"))
|
||||
|
||||
for _, partition := range []struct {
|
||||
name string
|
||||
start time.Time
|
||||
}{
|
||||
{name: "usage_logs_202604", start: aprilStart},
|
||||
{name: "usage_logs_202606", start: juneStart},
|
||||
} {
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(partition.start, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`DROP TABLE IF EXISTS "` + partition.name + `"`).
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
mock.ExpectCommit()
|
||||
}
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
|
||||
AddRow(service.GroupUsageDate(todayStart), time.Unix(0, 0).UTC(), "Asia/Shanghai"))
|
||||
mock.ExpectCommit()
|
||||
|
||||
require.NoError(t, repo.CleanupUsageLogs(context.Background(), cutoff))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositoryCleanupUsageLogsNonPartitionedFailureRollsBackWithoutSync(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
cutoff := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)
|
||||
deletedAt := time.Date(2026, 5, 3, 2, 0, 0, 0, time.UTC)
|
||||
|
||||
mock.ExpectQuery(`SELECT EXISTS`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false))
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery(`(?s)SELECT ctid.*ORDER BY created_at ASC, id ASC.*DELETE FROM usage_logs.*RETURNING created_at`).
|
||||
WithArgs(cutoff, usageLogsCleanupBatchSize).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(deletedAt))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(deletedAt, "Asia/Shanghai").
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
mock.ExpectRollback()
|
||||
|
||||
err := repo.CleanupUsageLogs(context.Background(), cutoff)
|
||||
require.ErrorIs(t, err, sql.ErrConnDone)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDashboardAggregationRepositoryCleanupUsageLogsPartitionFailureRollsBackAndStops(t *testing.T) {
|
||||
setGroupUsageRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newDashboardAggregationRepositoryWithSQL(db)
|
||||
cutoff := time.Date(2026, 7, 18, 0, 0, 0, 0, time.UTC)
|
||||
aprilStart := time.Date(2026, 4, 1, 0, 0, 0, 0, time.UTC)
|
||||
dropErr := errors.New("drop partition failed")
|
||||
|
||||
mock.ExpectQuery(`SELECT EXISTS`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
|
||||
mock.ExpectQuery(`SELECT c.relname`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"relname"}).
|
||||
AddRow("usage_logs_202606").
|
||||
AddRow("usage_logs_202604"))
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(aprilStart, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`DROP TABLE IF EXISTS "usage_logs_202604"`).
|
||||
WillReturnError(dropErr)
|
||||
mock.ExpectRollback()
|
||||
|
||||
err := repo.CleanupUsageLogs(context.Background(), cutoff)
|
||||
require.ErrorIs(t, err, dropErr)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func setGroupUsageRollupTestTimezone(t *testing.T) {
|
||||
t.Helper()
|
||||
useGroupUsageRepositoryTestTimezone(t, "Asia/Shanghai")
|
||||
}
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"log"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -14,7 +15,8 @@ import (
|
||||
)
|
||||
|
||||
type dashboardAggregationRepository struct {
|
||||
sql sqlExecutor
|
||||
sql sqlExecutor
|
||||
clock func() time.Time
|
||||
}
|
||||
|
||||
const usageLogsCleanupBatchSize = 10000
|
||||
@@ -33,7 +35,14 @@ func NewDashboardAggregationRepository(sqlDB *sql.DB) service.DashboardAggregati
|
||||
}
|
||||
|
||||
func newDashboardAggregationRepositoryWithSQL(sqlq sqlExecutor) *dashboardAggregationRepository {
|
||||
return &dashboardAggregationRepository{sql: sqlq}
|
||||
return &dashboardAggregationRepository{sql: sqlq, clock: time.Now}
|
||||
}
|
||||
|
||||
func (r *dashboardAggregationRepository) now() time.Time {
|
||||
if r.clock != nil {
|
||||
return r.clock()
|
||||
}
|
||||
return time.Now()
|
||||
}
|
||||
|
||||
func isPostgresDriver(db *sql.DB) bool {
|
||||
@@ -128,11 +137,23 @@ func (r *dashboardAggregationRepository) RecomputeRange(ctx context.Context, sta
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := lockGroupUsageRollupState(ctx, tx); err != nil {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
if err := invalidateGroupUsageRollupsAt(ctx, tx, start); err != nil {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
txRepo := newDashboardAggregationRepositoryWithSQL(tx)
|
||||
if err := txRepo.recomputeRangeInTx(ctx, hourStart, hourEnd, dayStart, dayEnd); err != nil {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
if err := txRepo.syncGroupUsageRollupsInTx(ctx, service.GroupUsageTodayStart(r.now())); err != nil {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
return r.recomputeRangeInTx(ctx, hourStart, hourEnd, dayStart, dayEnd)
|
||||
@@ -215,14 +236,35 @@ func (r *dashboardAggregationRepository) CleanupUsageLogs(ctx context.Context, c
|
||||
return err
|
||||
}
|
||||
if isPartitioned {
|
||||
return r.dropUsageLogsPartitions(ctx, cutoff)
|
||||
if err := r.dropUsageLogsPartitions(ctx, cutoff); err != nil {
|
||||
return err
|
||||
}
|
||||
} else if err := r.cleanupUsageLogsBatches(ctx, cutoff); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.SyncGroupUsageRollups(ctx, service.GroupUsageTodayStart(r.now()))
|
||||
}
|
||||
|
||||
func (r *dashboardAggregationRepository) cleanupUsageLogsBatches(ctx context.Context, cutoff time.Time) error {
|
||||
db, transactional := r.sql.(*sql.DB)
|
||||
for {
|
||||
if transactional {
|
||||
affected, err := cleanupUsageLogsBatchWithRollupInvalidation(ctx, db, cutoff)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if affected < usageLogsCleanupBatchSize {
|
||||
return nil
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
res, err := r.sql.ExecContext(ctx, `
|
||||
WITH victims AS (
|
||||
SELECT ctid
|
||||
FROM usage_logs
|
||||
WHERE created_at < $1
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $2
|
||||
)
|
||||
DELETE FROM usage_logs
|
||||
@@ -241,6 +283,66 @@ func (r *dashboardAggregationRepository) CleanupUsageLogs(ctx context.Context, c
|
||||
}
|
||||
}
|
||||
|
||||
func cleanupUsageLogsBatchWithRollupInvalidation(ctx context.Context, db *sql.DB, cutoff time.Time) (int64, error) {
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
rollback := func(err error) (int64, error) {
|
||||
_ = tx.Rollback()
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if err := lockGroupUsageRollupState(ctx, tx); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
rows, err := tx.QueryContext(ctx, `
|
||||
WITH victims AS (
|
||||
SELECT ctid
|
||||
FROM usage_logs
|
||||
WHERE created_at < $1
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $2
|
||||
)
|
||||
DELETE FROM usage_logs
|
||||
WHERE ctid IN (SELECT ctid FROM victims)
|
||||
RETURNING created_at
|
||||
`, cutoff.UTC(), usageLogsCleanupBatchSize)
|
||||
if err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
|
||||
var affected int64
|
||||
var earliestDeletedAt time.Time
|
||||
for rows.Next() {
|
||||
var deletedAt time.Time
|
||||
if err := rows.Scan(&deletedAt); err != nil {
|
||||
_ = rows.Close()
|
||||
return rollback(err)
|
||||
}
|
||||
affected++
|
||||
if earliestDeletedAt.IsZero() || deletedAt.Before(earliestDeletedAt) {
|
||||
earliestDeletedAt = deletedAt
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
_ = rows.Close()
|
||||
return rollback(err)
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
if affected > 0 {
|
||||
if err := invalidateGroupUsageRollupsAt(ctx, tx, earliestDeletedAt); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return affected, nil
|
||||
}
|
||||
|
||||
func (r *dashboardAggregationRepository) CleanupUsageBillingDedup(ctx context.Context, cutoff time.Time) error {
|
||||
for {
|
||||
res, err := r.sql.ExecContext(ctx, `
|
||||
@@ -489,14 +591,16 @@ func (r *dashboardAggregationRepository) dropUsageLogsPartitions(ctx context.Con
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
_ = rows.Close()
|
||||
}()
|
||||
|
||||
cutoffMonth := truncateToMonthUTC(cutoff)
|
||||
type usageLogsPartition struct {
|
||||
name string
|
||||
month time.Time
|
||||
}
|
||||
partitions := make([]usageLogsPartition, 0)
|
||||
for rows.Next() {
|
||||
var name string
|
||||
if err := rows.Scan(&name); err != nil {
|
||||
_ = rows.Close()
|
||||
return err
|
||||
}
|
||||
if !strings.HasPrefix(name, "usage_logs_") {
|
||||
@@ -509,12 +613,56 @@ func (r *dashboardAggregationRepository) dropUsageLogsPartitions(ctx context.Con
|
||||
}
|
||||
month = month.UTC()
|
||||
if month.Before(cutoffMonth) {
|
||||
if _, err := r.sql.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", pq.QuoteIdentifier(name))); err != nil {
|
||||
partitions = append(partitions, usageLogsPartition{name: name, month: month})
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
_ = rows.Close()
|
||||
return err
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
sort.Slice(partitions, func(i, j int) bool {
|
||||
return partitions[i].month.Before(partitions[j].month)
|
||||
})
|
||||
if db, ok := r.sql.(*sql.DB); ok {
|
||||
for _, partition := range partitions {
|
||||
if err := dropUsageLogsPartitionWithRollupInvalidation(ctx, db, partition.name, partition.month); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return rows.Err()
|
||||
for _, partition := range partitions {
|
||||
if _, err := r.sql.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", pq.QuoteIdentifier(partition.name))); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func dropUsageLogsPartitionWithRollupInvalidation(ctx context.Context, db *sql.DB, name string, monthStart time.Time) error {
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rollback := func(err error) error {
|
||||
_ = tx.Rollback()
|
||||
return err
|
||||
}
|
||||
|
||||
if err := lockGroupUsageRollupState(ctx, tx); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
if err := invalidateGroupUsageRollupsAt(ctx, tx, monthStart); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", pq.QuoteIdentifier(name))); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func (r *dashboardAggregationRepository) createUsageLogsPartition(ctx context.Context, month time.Time) error {
|
||||
|
||||
@@ -0,0 +1,511 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/migrations"
|
||||
"github.com/lib/pq"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGroupUsageRollupTriggerInvalidatesCascadedHistoricalDelete(t *testing.T) {
|
||||
for _, partitioned := range []bool{false, true} {
|
||||
name := "ordinary"
|
||||
if partitioned {
|
||||
name = "partitioned"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, partitioned)
|
||||
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
|
||||
VALUES (1, 1, 10, 1.25, TIMESTAMPTZ '2020-01-02 08:00:00+08');
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
|
||||
WHERE id = 1;
|
||||
DELETE FROM users WHERE id = 1;
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
var closedBefore string
|
||||
err = tx.QueryRowContext(ctx, `
|
||||
SELECT closed_before::text
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
`).Scan(&closedBefore)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "2020-01-02", closedBefore)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupUsageRollupTriggerSerializesLateHistoricalInsertWithPublish(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
seedTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
_, err := seedTx.ExecContext(ctx, `
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = DATE '2020-01-02'
|
||||
WHERE id = 1;
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, seedTx.Commit())
|
||||
|
||||
syncTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = syncTx.Rollback() }()
|
||||
var stateID int16
|
||||
require.NoError(t, syncTx.QueryRowContext(ctx, `
|
||||
SELECT id
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
FOR UPDATE
|
||||
`).Scan(&stateID))
|
||||
|
||||
lateTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = lateTx.Rollback() }()
|
||||
var lateBackendPID int
|
||||
require.NoError(t, lateTx.QueryRowContext(ctx, "SELECT pg_backend_pid()").Scan(&lateBackendPID))
|
||||
|
||||
insertResult := make(chan error, 1)
|
||||
go func() {
|
||||
_, insertErr := lateTx.ExecContext(ctx, `
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
|
||||
VALUES (1, 1, 10, 1.25, TIMESTAMPTZ '2020-01-02 09:00:00+08')
|
||||
`)
|
||||
insertResult <- insertErr
|
||||
}()
|
||||
|
||||
blocked, err := waitForGroupUsageRollupStateLock(ctx, lateBackendPID, insertResult)
|
||||
if err != nil || !blocked {
|
||||
_ = syncTx.Rollback()
|
||||
_ = lateTx.Rollback()
|
||||
require.NoError(t, err)
|
||||
require.True(t, blocked, "迟到写入必须等待正在发布水位的事务")
|
||||
}
|
||||
|
||||
_, err = syncTx.ExecContext(ctx, `
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
|
||||
WHERE id = 1
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, syncTx.Commit())
|
||||
|
||||
select {
|
||||
case err = <-insertResult:
|
||||
require.NoError(t, err)
|
||||
case <-ctx.Done():
|
||||
t.Fatal("等待迟到写入完成超时")
|
||||
}
|
||||
require.NoError(t, lateTx.Commit())
|
||||
|
||||
var closedBefore string
|
||||
err = integrationDB.QueryRowContext(ctx, fmt.Sprintf(
|
||||
"SELECT closed_before::text FROM %s.usage_group_rollup_state WHERE id = 1",
|
||||
pq.QuoteIdentifier(schema),
|
||||
)).Scan(&closedBefore)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "2020-01-02", closedBefore)
|
||||
}
|
||||
|
||||
func TestGroupUsageRollupTriggerSerializesInsertTransactionAcrossMidnight(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
seedTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
_, err := seedTx.ExecContext(ctx, `
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
|
||||
WHERE id = 1;
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, seedTx.Commit())
|
||||
|
||||
syncTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = syncTx.Rollback() }()
|
||||
var stateID int16
|
||||
require.NoError(t, syncTx.QueryRowContext(ctx, `
|
||||
SELECT id
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
FOR UPDATE
|
||||
`).Scan(&stateID))
|
||||
|
||||
insertTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = insertTx.Rollback() }()
|
||||
var insertBackendPID int
|
||||
require.NoError(t, insertTx.QueryRowContext(ctx, "SELECT pg_backend_pid()").Scan(&insertBackendPID))
|
||||
|
||||
insertResult := make(chan error, 1)
|
||||
go func() {
|
||||
_, insertErr := insertTx.ExecContext(ctx, `
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
|
||||
VALUES (1, 1, 10, 1.25, CURRENT_TIMESTAMP)
|
||||
`)
|
||||
insertResult <- insertErr
|
||||
}()
|
||||
|
||||
blocked, err := waitForGroupUsageRollupStateLock(ctx, insertBackendPID, insertResult)
|
||||
if err != nil || !blocked {
|
||||
_ = syncTx.Rollback()
|
||||
_ = insertTx.Rollback()
|
||||
require.NoError(t, err)
|
||||
require.True(t, blocked, "跨越零点的在途写入必须与水位发布串行化")
|
||||
}
|
||||
|
||||
_, err = syncTx.ExecContext(ctx, `
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date + 1
|
||||
WHERE id = 1
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, syncTx.Commit())
|
||||
|
||||
select {
|
||||
case err = <-insertResult:
|
||||
require.NoError(t, err)
|
||||
case <-ctx.Done():
|
||||
t.Fatal("等待跨零点写入完成超时")
|
||||
}
|
||||
require.NoError(t, insertTx.Commit())
|
||||
|
||||
var currentDate string
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx, `
|
||||
SELECT (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date::text
|
||||
`).Scan(¤tDate))
|
||||
var closedBefore string
|
||||
err = integrationDB.QueryRowContext(ctx, fmt.Sprintf(
|
||||
"SELECT closed_before::text FROM %s.usage_group_rollup_state WHERE id = 1",
|
||||
pq.QuoteIdentifier(schema),
|
||||
)).Scan(&closedBefore)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, currentDate, closedBefore)
|
||||
}
|
||||
|
||||
func TestGroupUsageRollupTriggerKeepsWatermarkForTodayInsert(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
|
||||
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
|
||||
WHERE id = 1;
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
|
||||
VALUES (1, 1, 10, 1.25, CURRENT_TIMESTAMP);
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
var unchanged bool
|
||||
err = tx.QueryRowContext(ctx, `
|
||||
SELECT closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
`).Scan(&unchanged)
|
||||
require.NoError(t, err)
|
||||
require.True(t, unchanged)
|
||||
}
|
||||
|
||||
func TestGroupUsageRollupTriggerUsesSessionTimezoneAcrossDST(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
|
||||
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
SET LOCAL TIME ZONE 'America/New_York';
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = DATE '2026-03-09',
|
||||
timezone_name = 'America/New_York'
|
||||
WHERE id = 1;
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
|
||||
VALUES (1, 1, 10, 1.25, TIMESTAMPTZ '2026-03-08 04:30:00+00');
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
var closedBefore string
|
||||
err = tx.QueryRowContext(ctx, `
|
||||
SELECT closed_before::text
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
`).Scan(&closedBefore)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "2026-03-07", closedBefore)
|
||||
}
|
||||
|
||||
func TestGroupUsageSummaryIncludesYesterdayAcrossWatermark(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
useGroupUsageRepositoryTestTimezone(t, "Asia/Shanghai")
|
||||
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
closedBefore string
|
||||
includeYesterday bool
|
||||
}{
|
||||
{name: "closed_rollup", closedBefore: "2026-08-14", includeYesterday: true},
|
||||
{name: "raw_tail", closedBefore: "2026-08-13", includeYesterday: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at) VALUES
|
||||
(1, 1, 10, 2, TIMESTAMPTZ '2026-08-12 12:00:00+08'),
|
||||
(2, 1, 10, 3, TIMESTAMPTZ '2026-08-13 12:00:00+08'),
|
||||
(3, 1, 10, 4, TIMESTAMPTZ '2026-08-14 12:00:00+08');
|
||||
INSERT INTO usage_group_daily_rollups (bucket_date, group_id, actual_cost, computed_at)
|
||||
VALUES (DATE '2026-08-12', 10, 2, NOW());
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
if tt.includeYesterday {
|
||||
_, err = tx.ExecContext(ctx, `
|
||||
INSERT INTO usage_group_daily_rollups (bucket_date, group_id, actual_cost, computed_at)
|
||||
VALUES (DATE '2026-08-13', 10, 3, NOW())
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
_, err = tx.ExecContext(ctx, `
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = $1::date,
|
||||
retained_from = TIMESTAMPTZ '2026-08-12 00:00:00+08'
|
||||
WHERE id = 1
|
||||
`, tt.closedBefore)
|
||||
require.NoError(t, err)
|
||||
|
||||
repo := newUsageLogRepositoryWithSQL(nil, tx)
|
||||
result, err := repo.GetAllGroupUsageSummary(ctx, todayStart)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, result, 1)
|
||||
require.InDelta(t, 9, result[0].TotalCost, 0.0000001)
|
||||
require.InDelta(t, 4, result[0].TodayCost, 0.0000001)
|
||||
require.InDelta(t, 3, result[0].YesterdayCost, 0.0000001)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupUsageRollupSyncRebuildsAfterTimezoneChange(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
|
||||
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
SET LOCAL TIME ZONE 'UTC';
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at) VALUES
|
||||
(1, 1, 10, 3, TIMESTAMPTZ '2026-03-08 05:30:00+00'),
|
||||
(2, 1, 10, 5, TIMESTAMPTZ '2026-03-09 04:30:00+00');
|
||||
INSERT INTO usage_group_daily_rollups (bucket_date, group_id, actual_cost, computed_at)
|
||||
VALUES (DATE '2026-03-08', 10, 99, NOW());
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = DATE '2026-03-09',
|
||||
retained_from = TIMESTAMPTZ '2026-03-08 05:30:00+00',
|
||||
timezone_name = 'Asia/Shanghai'
|
||||
WHERE id = 1;
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
repo := newDashboardAggregationRepositoryWithSQL(tx)
|
||||
require.NoError(t, repo.SyncGroupUsageRollups(ctx, todayStart))
|
||||
|
||||
var stateTimezone string
|
||||
var closedBefore string
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
SELECT timezone_name, closed_before::text
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
`).Scan(&stateTimezone, &closedBefore))
|
||||
require.Equal(t, "America/New_York", stateTimezone)
|
||||
require.Equal(t, "2026-03-09", closedBefore)
|
||||
|
||||
var rollupCost float64
|
||||
require.NoError(t, tx.QueryRowContext(ctx, `
|
||||
SELECT actual_cost
|
||||
FROM usage_group_daily_rollups
|
||||
WHERE bucket_date = DATE '2026-03-08' AND group_id = 10
|
||||
`).Scan(&rollupCost))
|
||||
require.InDelta(t, 3, rollupCost, 0.0000001)
|
||||
|
||||
usageRepo := newUsageLogRepositoryWithSQL(nil, tx)
|
||||
result, err := usageRepo.GetAllGroupUsageSummary(ctx, todayStart)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, result, 1)
|
||||
require.InDelta(t, 8, result[0].TotalCost, 0.0000001)
|
||||
require.InDelta(t, 5, result[0].TodayCost, 0.0000001)
|
||||
require.InDelta(t, 3, result[0].YesterdayCost, 0.0000001)
|
||||
}
|
||||
|
||||
func TestGroupUsageSummaryUsesConfiguredDSTBoundaries(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
|
||||
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
|
||||
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
|
||||
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
_, err := tx.ExecContext(ctx, `
|
||||
SET LOCAL TIME ZONE 'UTC';
|
||||
INSERT INTO groups (id) VALUES (10);
|
||||
INSERT INTO users (id) VALUES (1);
|
||||
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at) VALUES
|
||||
(1, 1, 10, 100, TIMESTAMPTZ '2026-03-08 04:30:00+00'),
|
||||
(2, 1, 10, 3, TIMESTAMPTZ '2026-03-08 05:30:00+00'),
|
||||
(3, 1, 10, 4, TIMESTAMPTZ '2026-03-09 03:30:00+00'),
|
||||
(4, 1, 10, 5, TIMESTAMPTZ '2026-03-09 04:30:00+00');
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = DATE '1970-01-01',
|
||||
retained_from = TIMESTAMPTZ '1970-01-01 00:00:00+00',
|
||||
timezone_name = 'America/New_York'
|
||||
WHERE id = 1;
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
repo := newUsageLogRepositoryWithSQL(nil, tx)
|
||||
result, err := repo.GetAllGroupUsageSummary(ctx, todayStart)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, result, 1)
|
||||
require.InDelta(t, 112, result[0].TotalCost, 0.0000001)
|
||||
require.InDelta(t, 5, result[0].TodayCost, 0.0000001)
|
||||
require.InDelta(t, 7, result[0].YesterdayCost, 0.0000001)
|
||||
}
|
||||
|
||||
func createGroupUsageRollupTriggerTestSchema(t *testing.T, ctx context.Context, partitioned bool) string {
|
||||
t.Helper()
|
||||
|
||||
schema := fmt.Sprintf("group_usage_rollup_trigger_%d", time.Now().UnixNano())
|
||||
quotedSchema := pq.QuoteIdentifier(schema)
|
||||
_, err := integrationDB.ExecContext(ctx, "CREATE SCHEMA "+quotedSchema)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() {
|
||||
_, _ = integrationDB.ExecContext(context.Background(), "DROP SCHEMA IF EXISTS "+quotedSchema+" CASCADE")
|
||||
})
|
||||
|
||||
tx, err := integrationDB.BeginTx(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
require.NoError(t, setGroupUsageRollupTriggerSearchPath(ctx, tx, quotedSchema))
|
||||
|
||||
usageLogsDDL := `
|
||||
CREATE TABLE usage_logs (
|
||||
id BIGINT PRIMARY KEY,
|
||||
user_id BIGINT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
group_id BIGINT REFERENCES groups(id) ON DELETE SET NULL,
|
||||
actual_cost NUMERIC(20, 10) NOT NULL,
|
||||
created_at TIMESTAMPTZ NOT NULL
|
||||
);
|
||||
`
|
||||
if partitioned {
|
||||
usageLogsDDL = `
|
||||
CREATE TABLE usage_logs (
|
||||
id BIGINT NOT NULL,
|
||||
user_id BIGINT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
group_id BIGINT REFERENCES groups(id) ON DELETE SET NULL,
|
||||
actual_cost NUMERIC(20, 10) NOT NULL,
|
||||
created_at TIMESTAMPTZ NOT NULL
|
||||
) PARTITION BY RANGE (created_at);
|
||||
CREATE TABLE usage_logs_default PARTITION OF usage_logs DEFAULT;
|
||||
`
|
||||
}
|
||||
|
||||
_, err = tx.ExecContext(ctx, `
|
||||
CREATE TABLE users (id BIGINT PRIMARY KEY);
|
||||
CREATE TABLE groups (id BIGINT PRIMARY KEY);
|
||||
`+usageLogsDDL)
|
||||
require.NoError(t, err)
|
||||
|
||||
for _, migrationName := range []string{
|
||||
"222_group_usage_daily_rollups.sql",
|
||||
"223_group_usage_rollup_timezone.sql",
|
||||
} {
|
||||
migrationSQL, readErr := migrations.FS.ReadFile(migrationName)
|
||||
require.NoError(t, readErr)
|
||||
for range 2 {
|
||||
_, err = tx.ExecContext(ctx, string(migrationSQL))
|
||||
require.NoError(t, err)
|
||||
}
|
||||
}
|
||||
require.NoError(t, tx.Commit())
|
||||
|
||||
return schema
|
||||
}
|
||||
|
||||
func beginGroupUsageRollupTriggerTestTx(t *testing.T, ctx context.Context, schema string) *sql.Tx {
|
||||
t.Helper()
|
||||
|
||||
tx, err := integrationDB.BeginTx(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, setGroupUsageRollupTriggerSearchPath(ctx, tx, pq.QuoteIdentifier(schema)))
|
||||
return tx
|
||||
}
|
||||
|
||||
func setGroupUsageRollupTriggerSearchPath(ctx context.Context, tx *sql.Tx, quotedSchema string) error {
|
||||
_, err := tx.ExecContext(ctx, "SET LOCAL search_path TO "+quotedSchema)
|
||||
return err
|
||||
}
|
||||
|
||||
func waitForGroupUsageRollupStateLock(
|
||||
ctx context.Context,
|
||||
backendPID int,
|
||||
insertResult <-chan error,
|
||||
) (bool, error) {
|
||||
ticker := time.NewTicker(10 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case err := <-insertResult:
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return false, nil
|
||||
case <-ticker.C:
|
||||
var waitEventType sql.NullString
|
||||
err := integrationDB.QueryRowContext(ctx, `
|
||||
SELECT wait_event_type
|
||||
FROM pg_stat_activity
|
||||
WHERE pid = $1
|
||||
`, backendPID).Scan(&waitEventType)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if waitEventType.Valid && waitEventType.String == "Lock" {
|
||||
return true, nil
|
||||
}
|
||||
case <-ctx.Done():
|
||||
return false, ctx.Err()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -291,6 +291,9 @@ func (r *usageCleanupRepository) DeleteUsageLogsBatch(ctx context.Context, filte
|
||||
return 0, fmt.Errorf("cleanup filters missing time range")
|
||||
}
|
||||
args = append(args, limit)
|
||||
if db, ok := r.sql.(*sql.DB); ok {
|
||||
return r.deleteUsageLogsBatchWithRollupInvalidation(ctx, db, whereClause, args)
|
||||
}
|
||||
query := fmt.Sprintf(`
|
||||
WITH target AS (
|
||||
SELECT id
|
||||
@@ -301,7 +304,7 @@ func (r *usageCleanupRepository) DeleteUsageLogsBatch(ctx context.Context, filte
|
||||
)
|
||||
DELETE FROM usage_logs
|
||||
WHERE id IN (SELECT id FROM target)
|
||||
RETURNING id
|
||||
RETURNING created_at
|
||||
`, whereClause, len(args))
|
||||
|
||||
rows, err := r.sql.QueryContext(ctx, query, args...)
|
||||
@@ -320,6 +323,68 @@ func (r *usageCleanupRepository) DeleteUsageLogsBatch(ctx context.Context, filte
|
||||
return deleted, nil
|
||||
}
|
||||
|
||||
func (r *usageCleanupRepository) deleteUsageLogsBatchWithRollupInvalidation(ctx context.Context, db *sql.DB, whereClause string, args []any) (int64, error) {
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
rollback := func(err error) (int64, error) {
|
||||
_ = tx.Rollback()
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if err := lockGroupUsageRollupState(ctx, tx); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
query := fmt.Sprintf(`
|
||||
WITH target AS (
|
||||
SELECT id
|
||||
FROM usage_logs
|
||||
WHERE %s
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $%d
|
||||
)
|
||||
DELETE FROM usage_logs
|
||||
WHERE id IN (SELECT id FROM target)
|
||||
RETURNING created_at
|
||||
`, whereClause, len(args))
|
||||
rows, err := tx.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
|
||||
var deleted int64
|
||||
var earliestDeletedAt time.Time
|
||||
for rows.Next() {
|
||||
var deletedAt time.Time
|
||||
if err := rows.Scan(&deletedAt); err != nil {
|
||||
_ = rows.Close()
|
||||
return rollback(err)
|
||||
}
|
||||
deleted++
|
||||
if earliestDeletedAt.IsZero() || deletedAt.Before(earliestDeletedAt) {
|
||||
earliestDeletedAt = deletedAt
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
_ = rows.Close()
|
||||
return rollback(err)
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
|
||||
if deleted > 0 {
|
||||
if err := invalidateGroupUsageRollupsAt(ctx, tx, earliestDeletedAt); err != nil {
|
||||
return rollback(err)
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return deleted, nil
|
||||
}
|
||||
|
||||
func buildUsageCleanupWhere(filters service.UsageCleanupFilters) (string, []any) {
|
||||
conditions := make([]string, 0, 8)
|
||||
args := make([]any, 0, 8)
|
||||
|
||||
@@ -21,6 +21,11 @@ func newSQLMock(t *testing.T) (*sql.DB, sqlmock.Sqlmock) {
|
||||
return db, mock
|
||||
}
|
||||
|
||||
func setUsageCleanupRollupTestTimezone(t *testing.T) {
|
||||
t.Helper()
|
||||
useGroupUsageRepositoryTestTimezone(t, "Asia/Shanghai")
|
||||
}
|
||||
|
||||
func TestNewUsageCleanupRepository(t *testing.T) {
|
||||
db, _ := newSQLMock(t)
|
||||
repo := NewUsageCleanupRepository(nil, db)
|
||||
@@ -398,6 +403,7 @@ func TestUsageCleanupRepositoryDeleteUsageLogsBatchMissingRange(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestUsageCleanupRepositoryDeleteUsageLogsBatch(t *testing.T) {
|
||||
setUsageCleanupRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := &usageCleanupRepository{sql: db}
|
||||
|
||||
@@ -412,9 +418,16 @@ func TestUsageCleanupRepositoryDeleteUsageLogsBatch(t *testing.T) {
|
||||
Model: &model,
|
||||
}
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery("DELETE FROM usage_logs").
|
||||
WithArgs(start, end, userID, "gpt-4", 2).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(int64(1)).AddRow(int64(2)))
|
||||
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(start.Add(time.Hour)).AddRow(start.Add(2 * time.Hour)))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(start.Add(time.Hour), "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
deleted, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 2)
|
||||
require.NoError(t, err)
|
||||
@@ -422,6 +435,62 @@ func TestUsageCleanupRepositoryDeleteUsageLogsBatch(t *testing.T) {
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestUsageCleanupRepositoryDeleteUsageLogsBatchAtomicallyInvalidatesGroupRollups(t *testing.T) {
|
||||
setUsageCleanupRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := &usageCleanupRepository{sql: db}
|
||||
|
||||
start := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
|
||||
end := start.Add(24 * time.Hour)
|
||||
firstDeletedAt := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
|
||||
secondDeletedAt := firstDeletedAt.Add(time.Hour)
|
||||
filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
|
||||
WithArgs(start, end, 2).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).
|
||||
AddRow(firstDeletedAt).
|
||||
AddRow(secondDeletedAt))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(firstDeletedAt, "Asia/Shanghai").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
deleted, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 2)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(2), deleted)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestUsageCleanupRepositoryDeleteUsageLogsBatchRollsBackWhenInvalidationFails(t *testing.T) {
|
||||
setUsageCleanupRollupTestTimezone(t)
|
||||
db, mock := newSQLMock(t)
|
||||
repo := &usageCleanupRepository{sql: db}
|
||||
|
||||
start := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
|
||||
end := start.Add(24 * time.Hour)
|
||||
deletedAt := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
|
||||
filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
|
||||
WithArgs(start, end, 1).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(deletedAt))
|
||||
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
|
||||
WithArgs(deletedAt, "Asia/Shanghai").
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
mock.ExpectRollback()
|
||||
|
||||
_, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 1)
|
||||
require.ErrorIs(t, err, sql.ErrConnDone)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestUsageCleanupRepositoryDeleteUsageLogsBatchQueryError(t *testing.T) {
|
||||
db, mock := newSQLMock(t)
|
||||
repo := &usageCleanupRepository{sql: db}
|
||||
@@ -430,9 +499,13 @@ func TestUsageCleanupRepositoryDeleteUsageLogsBatchQueryError(t *testing.T) {
|
||||
end := start.Add(24 * time.Hour)
|
||||
filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery("DELETE FROM usage_logs").
|
||||
WithArgs(start, end, 5).
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
mock.ExpectRollback()
|
||||
|
||||
_, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 5)
|
||||
require.Error(t, err)
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
//go:build unit
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestUsageLogRepositoryGetAllGroupUsageSummaryUsesRollupTail(t *testing.T) {
|
||||
db, mock := newSQLMock(t)
|
||||
repo := newUsageLogRepositoryWithSQL(nil, db)
|
||||
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
|
||||
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
|
||||
yesterdayStart := time.Date(2026, 3, 8, 5, 0, 0, 0, time.UTC)
|
||||
|
||||
mock.ExpectQuery(`(?s)usage_group_rollup_state.*usage_group_daily_rollups.*created_at >= state\.tail_start`).
|
||||
WithArgs(todayStart, yesterdayStart, "America/New_York", "2026-03-09", "2026-03-08").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"group_id", "total_cost", "today_cost", "yesterday_cost"}).
|
||||
AddRow(int64(7), 12.5, 1.25, 2.5))
|
||||
|
||||
result, err := repo.GetAllGroupUsageSummary(context.Background(), todayStart)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(7), result[0].GroupID)
|
||||
require.InDelta(t, 12.5, result[0].TotalCost, 0.0000001)
|
||||
require.InDelta(t, 1.25, result[0].TodayCost, 0.0000001)
|
||||
require.InDelta(t, 2.5, result[0].YesterdayCost, 0.0000001)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
@@ -715,39 +715,9 @@ func (r *usageLogRepository) GetUserBreakdownStats(ctx context.Context, startTim
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// GetAllGroupUsageSummary returns today's and cumulative actual_cost for every group.
|
||||
// todayStart is the start-of-day in the caller's timezone (UTC-based).
|
||||
// TODO(perf): This query scans ALL usage_logs rows for total_cost aggregation.
|
||||
// When usage_logs exceeds ~1M rows, consider adding a short-lived cache (30s)
|
||||
// or a materialized view / pre-aggregation table for cumulative costs.
|
||||
// GetAllGroupUsageSummary 返回所有分组在服务端配置时区内的今日、昨日与当前保留记录累计金额。
|
||||
func (r *usageLogRepository) GetAllGroupUsageSummary(ctx context.Context, todayStart time.Time) ([]usagestats.GroupUsageSummary, error) {
|
||||
query := `
|
||||
SELECT
|
||||
g.id AS group_id,
|
||||
COALESCE(SUM(ul.actual_cost), 0) AS total_cost,
|
||||
COALESCE(SUM(CASE WHEN ul.created_at >= $1 THEN ul.actual_cost ELSE 0 END), 0) AS today_cost
|
||||
FROM groups g
|
||||
LEFT JOIN usage_logs ul ON ul.group_id = g.id
|
||||
GROUP BY g.id
|
||||
`
|
||||
|
||||
rows, err := r.sql.QueryContext(ctx, query, todayStart)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
var results []usagestats.GroupUsageSummary
|
||||
for rows.Next() {
|
||||
var row usagestats.GroupUsageSummary
|
||||
if err := rows.Scan(&row.GroupID, &row.TotalCost, &row.TodayCost); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
results = append(results, row)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return results, nil
|
||||
return r.getAllGroupUsageSummaryFromRollups(ctx, todayStart)
|
||||
}
|
||||
|
||||
// resolveModelDimensionExpression maps model source type to a safe SQL expression.
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
appTimezone "github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
)
|
||||
|
||||
const groupUsageDateFormat = "2006-01-02"
|
||||
|
||||
// GroupUsageRollupRepository 是分组日汇总的可选持久化能力。
|
||||
type GroupUsageRollupRepository interface {
|
||||
SyncGroupUsageRollups(ctx context.Context, todayStart time.Time) error
|
||||
}
|
||||
|
||||
// GroupUsageTimezoneName 返回服务端配置的时区名称。
|
||||
func GroupUsageTimezoneName() string {
|
||||
return appTimezone.Location().String()
|
||||
}
|
||||
|
||||
// GroupUsageTodayStart 返回指定时刻在服务端配置时区内的自然日 UTC 起点。
|
||||
func GroupUsageTodayStart(at time.Time) time.Time {
|
||||
return appTimezone.StartOfDay(at).UTC()
|
||||
}
|
||||
|
||||
// GroupUsageYesterdayStart 返回指定时刻前一个本地日历日的 UTC 起点。
|
||||
func GroupUsageYesterdayStart(at time.Time) time.Time {
|
||||
return appTimezone.StartOfDay(at).AddDate(0, 0, -1).UTC()
|
||||
}
|
||||
|
||||
// GroupUsageDate 返回指定时刻在服务端配置时区内的日期。
|
||||
func GroupUsageDate(at time.Time) string {
|
||||
return at.In(appTimezone.Location()).Format(groupUsageDateFormat)
|
||||
}
|
||||
|
||||
// ParseGroupUsageDate 解析服务端配置时区内的日期并返回其零点。
|
||||
func ParseGroupUsageDate(value string) (time.Time, error) {
|
||||
return appTimezone.ParseInLocation(groupUsageDateFormat, value)
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
appTimezone "github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func useGroupUsageTestTimezone(t *testing.T, name string) {
|
||||
t.Helper()
|
||||
|
||||
previousName := appTimezone.Name()
|
||||
require.NoError(t, appTimezone.Init(name))
|
||||
t.Cleanup(func() { require.NoError(t, appTimezone.Init(previousName)) })
|
||||
}
|
||||
|
||||
func TestGroupUsageDateUsesConfiguredTimezoneBoundary(t *testing.T) {
|
||||
useGroupUsageTestTimezone(t, "America/New_York")
|
||||
|
||||
beforeMidnight := time.Date(2026, 3, 9, 3, 59, 59, 0, time.UTC)
|
||||
atMidnight := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
|
||||
|
||||
require.Equal(t, "2026-03-08", GroupUsageDate(beforeMidnight))
|
||||
require.Equal(t, "2026-03-09", GroupUsageDate(atMidnight))
|
||||
require.Equal(t, atMidnight, GroupUsageTodayStart(atMidnight))
|
||||
}
|
||||
|
||||
func TestGroupUsageParseDateUsesConfiguredTimezone(t *testing.T) {
|
||||
useGroupUsageTestTimezone(t, "America/New_York")
|
||||
|
||||
parsed, err := ParseGroupUsageDate("2026-03-08")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, time.Date(2026, 3, 8, 5, 0, 0, 0, time.UTC), parsed.UTC())
|
||||
require.Equal(t, "America/New_York", parsed.Location().String())
|
||||
}
|
||||
|
||||
func TestGroupUsageYesterdayStartHandlesDST(t *testing.T) {
|
||||
useGroupUsageTestTimezone(t, "America/New_York")
|
||||
|
||||
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
|
||||
yesterdayStart := GroupUsageYesterdayStart(todayStart)
|
||||
|
||||
require.Equal(t, time.Date(2026, 3, 8, 5, 0, 0, 0, time.UTC), yesterdayStart)
|
||||
require.Equal(t, 23*time.Hour, todayStart.Sub(yesterdayStart))
|
||||
require.Equal(t, "America/New_York", GroupUsageTimezoneName())
|
||||
}
|
||||
@@ -18,12 +18,14 @@ const (
|
||||
defaultDashboardAggregationBackfillTimeout = 30 * time.Minute
|
||||
dashboardAggregationRetentionInterval = 6 * time.Hour
|
||||
|
||||
// dashboardAggregationLeaderLockKey gates the periodic scheduled aggregation so
|
||||
// that only one instance runs it per cycle in a multi-replica deployment.
|
||||
// dashboardAggregationLeaderLockKey 保证多副本部署中每个周期只有一个实例执行聚合。
|
||||
dashboardAggregationLeaderLockKey = "dashboard:aggregation:leader"
|
||||
// dashboardAggregationLeaderLockTTL must exceed the job's worst-case runtime
|
||||
// (defaultDashboardAggregationTimeout) so the lock never expires mid-run.
|
||||
// TTL 必须覆盖 dashboard 聚合与分组日汇总两个有界阶段，避免任务中途失锁。
|
||||
dashboardAggregationLeaderLockTTL = 5 * time.Minute
|
||||
|
||||
// 启动回填耗时可能远长于周期聚合，因此使用独立锁并让 TTL 严格大于回填超时。
|
||||
dashboardAggregationGroupUsageBackfillLeaderLockKey = "dashboard:aggregation:group-usage-backfill:leader"
|
||||
dashboardAggregationGroupUsageBackfillLeaderLockTTL = defaultDashboardAggregationBackfillTimeout + time.Minute
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -95,6 +97,7 @@ func (s *DashboardAggregationService) Start() {
|
||||
logger.LegacyPrintf("service.dashboard_aggregation", "[DashboardAggregation] 聚合作业已禁用")
|
||||
return
|
||||
}
|
||||
go s.runStartupGroupUsageSync()
|
||||
|
||||
interval := time.Duration(s.cfg.IntervalSeconds) * time.Second
|
||||
if interval <= 0 {
|
||||
@@ -229,6 +232,7 @@ func (s *DashboardAggregationService) runScheduledAggregation() {
|
||||
return
|
||||
}
|
||||
defer release()
|
||||
defer s.runScheduledGroupUsageSync()
|
||||
|
||||
now := time.Now().UTC()
|
||||
last, err := s.repo.GetAggregationWatermark(ctx)
|
||||
@@ -269,6 +273,35 @@ func (s *DashboardAggregationService) runScheduledAggregation() {
|
||||
s.maybeCleanupRetention(ctx, now)
|
||||
}
|
||||
|
||||
func (s *DashboardAggregationService) runScheduledGroupUsageSync() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), defaultDashboardAggregationTimeout)
|
||||
defer cancel()
|
||||
if err := s.syncGroupUsageRollups(ctx, time.Now().UTC()); err != nil {
|
||||
logger.LegacyPrintf("service.dashboard_aggregation", "[DashboardAggregation] 分组用量日汇总失败: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *DashboardAggregationService) runStartupGroupUsageSync() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), defaultDashboardAggregationBackfillTimeout)
|
||||
defer cancel()
|
||||
release, ok := tryAcquireSingletonLeaderLock(ctx, s.lockCache, s.db, dashboardAggregationGroupUsageBackfillLeaderLockKey, s.instanceID, dashboardAggregationGroupUsageBackfillLeaderLockTTL)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
defer release()
|
||||
if err := s.syncGroupUsageRollups(ctx, time.Now().UTC()); err != nil {
|
||||
logger.LegacyPrintf("service.dashboard_aggregation", "[DashboardAggregation] 启动分组用量回填失败: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *DashboardAggregationService) syncGroupUsageRollups(ctx context.Context, now time.Time) error {
|
||||
repo, ok := s.repo.(GroupUsageRollupRepository)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return repo.SyncGroupUsageRollups(ctx, GroupUsageTodayStart(now))
|
||||
}
|
||||
|
||||
func (s *DashboardAggregationService) backfillRange(ctx context.Context, start, end time.Time) error {
|
||||
if !atomic.CompareAndSwapInt32(&s.running, 0, 1) {
|
||||
return errDashboardAggregationRunning
|
||||
|
||||
@@ -24,12 +24,36 @@ type dashboardAggregationRepoTestStub struct {
|
||||
cleanupUsageErr error
|
||||
cleanupDedupErr error
|
||||
ensurePartitionErr error
|
||||
aggregateCtx context.Context
|
||||
events *[]string
|
||||
}
|
||||
|
||||
type dashboardAggregationRollupRepoTestStub struct {
|
||||
*dashboardAggregationRepoTestStub
|
||||
groupRollupCalls int
|
||||
groupRollupAt time.Time
|
||||
groupRollupErr error
|
||||
groupRollupCtx context.Context
|
||||
}
|
||||
|
||||
func (s *dashboardAggregationRollupRepoTestStub) SyncGroupUsageRollups(ctx context.Context, todayStart time.Time) error {
|
||||
s.groupRollupCalls++
|
||||
s.groupRollupAt = todayStart
|
||||
s.groupRollupCtx = ctx
|
||||
if s.events != nil {
|
||||
*s.events = append(*s.events, "group_rollup")
|
||||
}
|
||||
return s.groupRollupErr
|
||||
}
|
||||
|
||||
func (s *dashboardAggregationRepoTestStub) AggregateRange(ctx context.Context, start, end time.Time) error {
|
||||
s.aggregateCalls++
|
||||
s.aggregateCtx = ctx
|
||||
s.lastStart = start
|
||||
s.lastEnd = end
|
||||
if s.events != nil {
|
||||
*s.events = append(*s.events, "dashboard_aggregation")
|
||||
}
|
||||
return s.aggregateErr
|
||||
}
|
||||
|
||||
@@ -88,6 +112,103 @@ func TestDashboardAggregationService_RunScheduledAggregation_EpochUsesRetentionS
|
||||
require.Equal(t, truncateToDayUTC(repo.lastEnd.AddDate(0, 0, -1)), repo.lastStart)
|
||||
}
|
||||
|
||||
func TestDashboardAggregationService_RunScheduledAggregationSyncsGroupUsageRollups(t *testing.T) {
|
||||
baseRepo := &dashboardAggregationRepoTestStub{watermark: time.Now().UTC()}
|
||||
repo := &dashboardAggregationRollupRepoTestStub{dashboardAggregationRepoTestStub: baseRepo}
|
||||
svc := &DashboardAggregationService{
|
||||
repo: repo,
|
||||
cfg: config.DashboardAggregationConfig{
|
||||
Enabled: true,
|
||||
IntervalSeconds: 60,
|
||||
LookbackSeconds: 120,
|
||||
Retention: config.DashboardAggregationRetentionConfig{
|
||||
UsageLogsDays: 1,
|
||||
UsageBillingDedupDays: 2,
|
||||
HourlyDays: 1,
|
||||
DailyDays: 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
before := GroupUsageTodayStart(time.Now())
|
||||
svc.runScheduledAggregation()
|
||||
after := GroupUsageTodayStart(time.Now())
|
||||
|
||||
require.Equal(t, 1, repo.groupRollupCalls)
|
||||
require.Contains(t, []time.Time{before, after}, repo.groupRollupAt)
|
||||
}
|
||||
|
||||
func TestDashboardAggregationService_RunScheduledAggregationSyncsGroupAfterDashboardEarlyReturn(t *testing.T) {
|
||||
events := make([]string, 0, 2)
|
||||
baseRepo := &dashboardAggregationRepoTestStub{
|
||||
watermark: time.Now().UTC(),
|
||||
aggregateErr: errors.New("dashboard aggregation failed"),
|
||||
events: &events,
|
||||
}
|
||||
repo := &dashboardAggregationRollupRepoTestStub{
|
||||
dashboardAggregationRepoTestStub: baseRepo,
|
||||
groupRollupErr: errors.New("group rollup failed"),
|
||||
}
|
||||
svc := &DashboardAggregationService{
|
||||
repo: repo,
|
||||
cfg: config.DashboardAggregationConfig{
|
||||
LookbackSeconds: 120,
|
||||
Retention: config.DashboardAggregationRetentionConfig{
|
||||
UsageLogsDays: 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc.runScheduledAggregation()
|
||||
|
||||
require.Equal(t, []string{"dashboard_aggregation", "group_rollup"}, events)
|
||||
require.NotNil(t, repo.aggregateCtx)
|
||||
require.NotNil(t, repo.groupRollupCtx)
|
||||
if repo.aggregateCtx == repo.groupRollupCtx {
|
||||
t.Fatal("分组日汇总必须使用独立于 dashboard 聚合的 context")
|
||||
}
|
||||
groupDeadline, ok := repo.groupRollupCtx.Deadline()
|
||||
require.True(t, ok, "group rollup context must be bounded")
|
||||
require.LessOrEqual(t, time.Until(groupDeadline), defaultDashboardAggregationTimeout)
|
||||
}
|
||||
|
||||
type dashboardAggregationLeaderLockRecordingCache struct {
|
||||
delegate *fakeLeaderLockCache
|
||||
acquireKeys []string
|
||||
acquireTTLs []time.Duration
|
||||
}
|
||||
|
||||
func (c *dashboardAggregationLeaderLockRecordingCache) TryAcquireLeaderLock(ctx context.Context, key, owner string, ttl time.Duration) (bool, error) {
|
||||
c.acquireKeys = append(c.acquireKeys, key)
|
||||
c.acquireTTLs = append(c.acquireTTLs, ttl)
|
||||
return c.delegate.TryAcquireLeaderLock(ctx, key, owner, ttl)
|
||||
}
|
||||
|
||||
func (c *dashboardAggregationLeaderLockRecordingCache) ReleaseLeaderLock(ctx context.Context, key, owner string) error {
|
||||
return c.delegate.ReleaseLeaderLock(ctx, key, owner)
|
||||
}
|
||||
|
||||
func TestDashboardAggregationService_StartupGroupSyncUsesIndependentLongLivedLeaderLock(t *testing.T) {
|
||||
delegate := &fakeLeaderLockCache{}
|
||||
_, err := delegate.TryAcquireLeaderLock(context.Background(), dashboardAggregationLeaderLockKey, "periodic-peer", time.Hour)
|
||||
require.NoError(t, err)
|
||||
cache := &dashboardAggregationLeaderLockRecordingCache{delegate: delegate}
|
||||
repo := &dashboardAggregationRollupRepoTestStub{dashboardAggregationRepoTestStub: &dashboardAggregationRepoTestStub{}}
|
||||
svc := &DashboardAggregationService{
|
||||
repo: repo,
|
||||
lockCache: cache,
|
||||
instanceID: "startup-instance",
|
||||
}
|
||||
|
||||
svc.runStartupGroupUsageSync()
|
||||
|
||||
require.Len(t, cache.acquireKeys, 1)
|
||||
require.NotEqual(t, dashboardAggregationLeaderLockKey, cache.acquireKeys[0])
|
||||
require.Len(t, cache.acquireTTLs, 1)
|
||||
require.Greater(t, cache.acquireTTLs[0], defaultDashboardAggregationBackfillTimeout)
|
||||
require.Equal(t, 1, repo.groupRollupCalls)
|
||||
}
|
||||
|
||||
func TestDashboardAggregationService_CleanupRetentionFailure_DoesNotRecord(t *testing.T) {
|
||||
repo := &dashboardAggregationRepoTestStub{cleanupAggregatesErr: errors.New("清理失败")}
|
||||
svc := &DashboardAggregationService{
|
||||
|
||||
@@ -212,7 +212,7 @@ func (s *DashboardService) GetGroupStatsWithUsageFilters(ctx context.Context, st
|
||||
return s.GetGroupStatsWithFilters(ctx, startTime, endTime, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.RequestType, filters.Stream, filters.BillingType)
|
||||
}
|
||||
|
||||
// GetGroupUsageSummary returns today's and cumulative cost for all groups.
|
||||
// GetGroupUsageSummary returns today's, yesterday's, and cumulative cost for all groups.
|
||||
func (s *DashboardService) GetGroupUsageSummary(ctx context.Context, todayStart time.Time) ([]usagestats.GroupUsageSummary, error) {
|
||||
results, err := s.usageRepo.GetAllGroupUsageSummary(ctx, todayStart)
|
||||
if err != nil {
|
||||
|
||||
@@ -0,0 +1,138 @@
|
||||
-- /admin/groups 分组用量日汇总。
|
||||
-- 迁移创建结构与源表失效触发器，历史数据由后台聚合作业按持久水位回填。
|
||||
|
||||
CREATE TABLE IF NOT EXISTS usage_group_daily_rollups (
|
||||
bucket_date DATE NOT NULL,
|
||||
group_id BIGINT NOT NULL,
|
||||
actual_cost DECIMAL(20, 10) NOT NULL DEFAULT 0,
|
||||
computed_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
PRIMARY KEY (bucket_date, group_id)
|
||||
);
|
||||
|
||||
COMMENT ON TABLE usage_group_daily_rollups IS '按北京时间自然日聚合的分组实际费用。';
|
||||
COMMENT ON COLUMN usage_group_daily_rollups.bucket_date IS '北京时间自然日。';
|
||||
|
||||
CREATE TABLE IF NOT EXISTS usage_group_rollup_state (
|
||||
id SMALLINT PRIMARY KEY CHECK (id = 1),
|
||||
closed_before DATE NOT NULL DEFAULT DATE '1970-01-01',
|
||||
retained_from TIMESTAMPTZ NOT NULL DEFAULT TIMESTAMPTZ '1970-01-01 00:00:00+00',
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||
);
|
||||
|
||||
COMMENT ON TABLE usage_group_rollup_state IS '分组日汇总的单行发布水位。';
|
||||
COMMENT ON COLUMN usage_group_rollup_state.closed_before IS '已完整发布日桶的北京时间日期排他上界。';
|
||||
|
||||
INSERT INTO usage_group_rollup_state (id, closed_before, retained_from)
|
||||
VALUES (1, DATE '1970-01-01', TIMESTAMPTZ '1970-01-01 00:00:00+00')
|
||||
ON CONFLICT (id) DO NOTHING;
|
||||
|
||||
-- 已发布范围的源记录发生变化时，必须在同一事务内后退发布水位。
|
||||
-- DELETE/UPDATE 使用行级触发器，以覆盖外键级联、分区表和直接分区写入。
|
||||
CREATE OR REPLACE FUNCTION invalidate_group_usage_rollup_state()
|
||||
RETURNS TRIGGER
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
DECLARE
|
||||
affected_date DATE;
|
||||
published_before DATE;
|
||||
BEGIN
|
||||
IF TG_OP = 'DELETE' THEN
|
||||
affected_date := (OLD.created_at AT TIME ZONE 'Asia/Shanghai')::date;
|
||||
ELSE
|
||||
IF OLD.group_id IS NULL THEN
|
||||
affected_date := (NEW.created_at AT TIME ZONE 'Asia/Shanghai')::date;
|
||||
ELSIF NEW.group_id IS NULL THEN
|
||||
affected_date := (OLD.created_at AT TIME ZONE 'Asia/Shanghai')::date;
|
||||
ELSE
|
||||
affected_date := LEAST(
|
||||
(OLD.created_at AT TIME ZONE 'Asia/Shanghai')::date,
|
||||
(NEW.created_at AT TIME ZONE 'Asia/Shanghai')::date
|
||||
);
|
||||
END IF;
|
||||
END IF;
|
||||
|
||||
-- 即使当前已发布水位尚未越过受影响日期，也必须先锁行。
|
||||
-- 否则并发关闭作业可能在本事务之后把水位推进，覆盖本次失效。
|
||||
SELECT closed_before
|
||||
INTO published_before
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
FOR UPDATE;
|
||||
|
||||
IF published_before > affected_date THEN
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = LEAST(closed_before, affected_date),
|
||||
updated_at = NOW()
|
||||
WHERE id = 1;
|
||||
END IF;
|
||||
|
||||
IF TG_OP = 'DELETE' THEN
|
||||
RETURN OLD;
|
||||
END IF;
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$;
|
||||
|
||||
-- INSERT 是网关高频路径。transition table 让每个批量 INSERT 只锁一次状态行；
|
||||
-- KEY SHARE 在普通写入之间兼容，但会与关闭作业的 FOR UPDATE 串行化。
|
||||
CREATE OR REPLACE FUNCTION invalidate_group_usage_rollup_state_after_insert()
|
||||
RETURNS TRIGGER
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
DECLARE
|
||||
affected_date DATE;
|
||||
published_before DATE;
|
||||
BEGIN
|
||||
SELECT MIN((created_at AT TIME ZONE 'Asia/Shanghai')::date)
|
||||
INTO affected_date
|
||||
FROM inserted_usage_logs
|
||||
WHERE group_id IS NOT NULL;
|
||||
|
||||
IF affected_date IS NULL THEN
|
||||
RETURN NULL;
|
||||
END IF;
|
||||
|
||||
SELECT closed_before
|
||||
INTO published_before
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
FOR KEY SHARE;
|
||||
|
||||
IF published_before > affected_date THEN
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = LEAST(closed_before, affected_date),
|
||||
updated_at = NOW()
|
||||
WHERE id = 1;
|
||||
END IF;
|
||||
|
||||
RETURN NULL;
|
||||
END;
|
||||
$$;
|
||||
|
||||
DROP TRIGGER IF EXISTS usage_logs_group_rollup_invalidate_insert ON usage_logs;
|
||||
CREATE TRIGGER usage_logs_group_rollup_invalidate_insert
|
||||
AFTER INSERT ON usage_logs
|
||||
REFERENCING NEW TABLE AS inserted_usage_logs
|
||||
FOR EACH STATEMENT
|
||||
EXECUTE FUNCTION invalidate_group_usage_rollup_state_after_insert();
|
||||
|
||||
DROP TRIGGER IF EXISTS usage_logs_group_rollup_invalidate_delete ON usage_logs;
|
||||
CREATE TRIGGER usage_logs_group_rollup_invalidate_delete
|
||||
AFTER DELETE ON usage_logs
|
||||
FOR EACH ROW
|
||||
WHEN (OLD.group_id IS NOT NULL)
|
||||
EXECUTE FUNCTION invalidate_group_usage_rollup_state();
|
||||
|
||||
DROP TRIGGER IF EXISTS usage_logs_group_rollup_invalidate_update ON usage_logs;
|
||||
CREATE TRIGGER usage_logs_group_rollup_invalidate_update
|
||||
AFTER UPDATE OF created_at, group_id, actual_cost ON usage_logs
|
||||
FOR EACH ROW
|
||||
WHEN (
|
||||
(
|
||||
OLD.created_at IS DISTINCT FROM NEW.created_at
|
||||
OR OLD.group_id IS DISTINCT FROM NEW.group_id
|
||||
OR OLD.actual_cost IS DISTINCT FROM NEW.actual_cost
|
||||
)
|
||||
AND (OLD.group_id IS NOT NULL OR NEW.group_id IS NOT NULL)
|
||||
)
|
||||
EXECUTE FUNCTION invalidate_group_usage_rollup_state();
|
||||
@@ -0,0 +1,89 @@
|
||||
-- 让 /admin/groups 分组日汇总跟随服务端配置时区。
|
||||
-- 222 迁移生成的存量日桶均为北京时间，因此新增状态默认标记为 Asia/Shanghai；
|
||||
-- 服务启动后若当前 TZ 不同，后台同步会检测到不一致并重建日桶。
|
||||
|
||||
ALTER TABLE usage_group_rollup_state
|
||||
ADD COLUMN IF NOT EXISTS timezone_name TEXT NOT NULL DEFAULT 'Asia/Shanghai';
|
||||
|
||||
COMMENT ON COLUMN usage_group_rollup_state.timezone_name IS '当前分组日桶采用的 IANA 时区名称。';
|
||||
COMMENT ON COLUMN usage_group_rollup_state.closed_before IS '已完整发布日桶的配置时区日期排他上界。';
|
||||
COMMENT ON COLUMN usage_group_daily_rollups.bucket_date IS 'timezone_name 对应时区的自然日。';
|
||||
|
||||
CREATE OR REPLACE FUNCTION invalidate_group_usage_rollup_state()
|
||||
RETURNS TRIGGER
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
DECLARE
|
||||
affected_date DATE;
|
||||
published_before DATE;
|
||||
configured_timezone TEXT := current_setting('TimeZone');
|
||||
BEGIN
|
||||
IF TG_OP = 'DELETE' THEN
|
||||
affected_date := (OLD.created_at AT TIME ZONE configured_timezone)::date;
|
||||
ELSE
|
||||
IF OLD.group_id IS NULL THEN
|
||||
affected_date := (NEW.created_at AT TIME ZONE configured_timezone)::date;
|
||||
ELSIF NEW.group_id IS NULL THEN
|
||||
affected_date := (OLD.created_at AT TIME ZONE configured_timezone)::date;
|
||||
ELSE
|
||||
affected_date := LEAST(
|
||||
(OLD.created_at AT TIME ZONE configured_timezone)::date,
|
||||
(NEW.created_at AT TIME ZONE configured_timezone)::date
|
||||
);
|
||||
END IF;
|
||||
END IF;
|
||||
|
||||
SELECT closed_before
|
||||
INTO published_before
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
FOR UPDATE;
|
||||
|
||||
IF published_before > affected_date THEN
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = LEAST(closed_before, affected_date),
|
||||
updated_at = NOW()
|
||||
WHERE id = 1;
|
||||
END IF;
|
||||
|
||||
IF TG_OP = 'DELETE' THEN
|
||||
RETURN OLD;
|
||||
END IF;
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE OR REPLACE FUNCTION invalidate_group_usage_rollup_state_after_insert()
|
||||
RETURNS TRIGGER
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
DECLARE
|
||||
affected_date DATE;
|
||||
published_before DATE;
|
||||
configured_timezone TEXT := current_setting('TimeZone');
|
||||
BEGIN
|
||||
SELECT MIN((created_at AT TIME ZONE configured_timezone)::date)
|
||||
INTO affected_date
|
||||
FROM inserted_usage_logs
|
||||
WHERE group_id IS NOT NULL;
|
||||
|
||||
IF affected_date IS NULL THEN
|
||||
RETURN NULL;
|
||||
END IF;
|
||||
|
||||
SELECT closed_before
|
||||
INTO published_before
|
||||
FROM usage_group_rollup_state
|
||||
WHERE id = 1
|
||||
FOR KEY SHARE;
|
||||
|
||||
IF published_before > affected_date THEN
|
||||
UPDATE usage_group_rollup_state
|
||||
SET closed_before = LEAST(closed_before, affected_date),
|
||||
updated_at = NOW()
|
||||
WHERE id = 1;
|
||||
END IF;
|
||||
|
||||
RETURN NULL;
|
||||
END;
|
||||
$$;
|
||||
@@ -0,0 +1,52 @@
|
||||
//go:build unit
|
||||
|
||||
package migrations
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestMigration222CreatesGroupUsageRollups(t *testing.T) {
|
||||
content, err := FS.ReadFile("222_group_usage_daily_rollups.sql")
|
||||
require.NoError(t, err)
|
||||
|
||||
sql := string(content)
|
||||
require.Contains(t, sql, "CREATE TABLE IF NOT EXISTS usage_group_daily_rollups")
|
||||
require.Contains(t, sql, "actual_cost DECIMAL(20, 10)")
|
||||
require.Contains(t, sql, "PRIMARY KEY (bucket_date, group_id)")
|
||||
require.Contains(t, sql, "CREATE TABLE IF NOT EXISTS usage_group_rollup_state")
|
||||
require.Contains(t, sql, "CHECK (id = 1)")
|
||||
require.Contains(t, sql, "TIMESTAMPTZ '1970-01-01 00:00:00+00'")
|
||||
require.Contains(t, sql, "ON CONFLICT (id) DO NOTHING")
|
||||
}
|
||||
|
||||
func TestMigration222InvalidatesClosedBucketsWhenUsageLogsChange(t *testing.T) {
|
||||
content, err := FS.ReadFile("222_group_usage_daily_rollups.sql")
|
||||
require.NoError(t, err)
|
||||
|
||||
sql := string(content)
|
||||
require.Contains(t, sql, "CREATE OR REPLACE FUNCTION invalidate_group_usage_rollup_state")
|
||||
require.Contains(t, sql, "SELECT closed_before")
|
||||
require.Contains(t, sql, "FOR UPDATE")
|
||||
require.Contains(t, sql, "FOR KEY SHARE")
|
||||
require.Contains(t, sql, "REFERENCING NEW TABLE AS inserted_usage_logs")
|
||||
require.Contains(t, sql, "closed_before = LEAST(closed_before, affected_date)")
|
||||
require.Contains(t, sql, "CREATE TRIGGER usage_logs_group_rollup_invalidate_insert")
|
||||
require.Contains(t, sql, "CREATE TRIGGER usage_logs_group_rollup_invalidate_delete")
|
||||
require.Contains(t, sql, "CREATE TRIGGER usage_logs_group_rollup_invalidate_update")
|
||||
require.Contains(t, sql, "AFTER UPDATE OF created_at, group_id, actual_cost")
|
||||
}
|
||||
|
||||
func TestMigration223TracksConfiguredTimezone(t *testing.T) {
|
||||
content, err := FS.ReadFile("223_group_usage_rollup_timezone.sql")
|
||||
require.NoError(t, err)
|
||||
|
||||
sql := string(content)
|
||||
require.Contains(t, sql, "ADD COLUMN IF NOT EXISTS timezone_name TEXT")
|
||||
require.Contains(t, sql, "DEFAULT 'Asia/Shanghai'")
|
||||
require.Contains(t, sql, "current_setting('TimeZone')")
|
||||
require.Contains(t, sql, "CREATE OR REPLACE FUNCTION invalidate_group_usage_rollup_state")
|
||||
require.Contains(t, sql, "CREATE OR REPLACE FUNCTION invalidate_group_usage_rollup_state_after_insert")
|
||||
}
|
||||
Generated
+4
-4
@@ -3322,8 +3322,8 @@ packages:
|
||||
mz@2.7.0:
|
||||
resolution: {integrity: sha512-z81GNO7nnYMEhrGh9LeymoE4+Yr0Wn5McHIZMK5cfQCl+NDX08sCZgUc9/6MHni9IWuFLm1Z3HTCXu2z9fN62Q==}
|
||||
|
||||
nanoid@3.3.17:
|
||||
resolution: {integrity: sha512-xQLf0A3HOMlgHq0n247/LRuAOYmB7dXJ/DvAxGvsSBij45XtBSmQycu+F8ODbHwns/XyFZagyL1+J0Offw1E0g==}
|
||||
nanoid@3.3.18:
|
||||
resolution: {integrity: sha512-DTg4MJbGMWkfi6VZFdNt2/caMbQy4Ou+Op/hJQvGEWcnVfoA1QA+xzRKAzw9jD6+GVOOeYr/mIcuDSdug6F6+w==}
|
||||
engines: {node: ^10 || ^12 || ^13.7 || ^14 || >=15.0.1}
|
||||
hasBin: true
|
||||
|
||||
@@ -8443,7 +8443,7 @@ snapshots:
|
||||
object-assign: 4.1.1
|
||||
thenify-all: 1.6.0
|
||||
|
||||
nanoid@3.3.17: {}
|
||||
nanoid@3.3.18: {}
|
||||
|
||||
natural-compare@1.4.0: {}
|
||||
|
||||
@@ -8641,7 +8641,7 @@ snapshots:
|
||||
|
||||
postcss@8.5.23:
|
||||
dependencies:
|
||||
nanoid: 3.3.17
|
||||
nanoid: 3.3.18
|
||||
picocolors: 1.1.1
|
||||
source-map-js: 1.2.1
|
||||
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
const { get } = vi.hoisted(() => ({
|
||||
get: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('@/api/client', () => ({
|
||||
apiClient: { get },
|
||||
}))
|
||||
|
||||
import { getUsageSummary } from '@/api/admin/groups'
|
||||
|
||||
describe('admin group usage summary API', () => {
|
||||
beforeEach(() => {
|
||||
get.mockReset()
|
||||
get.mockResolvedValue({ data: [] })
|
||||
})
|
||||
|
||||
it('does not send browser timezone parameters', async () => {
|
||||
const summary = [
|
||||
{ group_id: 1, today_cost: 1.25, yesterday_cost: 2.5, total_cost: 9.75 },
|
||||
]
|
||||
get.mockResolvedValue({ data: summary })
|
||||
|
||||
await expect(getUsageSummary()).resolves.toEqual(summary)
|
||||
|
||||
expect(get).toHaveBeenCalledWith('/admin/groups/usage-summary')
|
||||
})
|
||||
})
|
||||
@@ -446,18 +446,15 @@ export async function clearGroupRPMOverrides(id: number): Promise<{ message: str
|
||||
}
|
||||
|
||||
/**
|
||||
* Get usage summary (today + cumulative cost) for all groups
|
||||
* @param timezone - IANA timezone string (e.g. "Asia/Shanghai")
|
||||
* Get usage summary (today + yesterday + cumulative cost) for all groups
|
||||
* @returns Array of group usage summaries
|
||||
*/
|
||||
export async function getUsageSummary(
|
||||
timezone?: string
|
||||
): Promise<{ group_id: number; today_cost: number; total_cost: number }[]> {
|
||||
export async function getUsageSummary(): Promise<
|
||||
{ group_id: number; today_cost: number; yesterday_cost: number; total_cost: number }[]
|
||||
> {
|
||||
const { data } = await apiClient.get<
|
||||
{ group_id: number; today_cost: number; total_cost: number }[]
|
||||
>('/admin/groups/usage-summary', {
|
||||
params: timezone ? { timezone } : undefined
|
||||
})
|
||||
{ group_id: number; today_cost: number; yesterday_cost: number; total_cost: number }[]
|
||||
>('/admin/groups/usage-summary')
|
||||
return data
|
||||
}
|
||||
|
||||
|
||||
@@ -821,6 +821,7 @@ export default {
|
||||
userStatus: 'Status'
|
||||
},
|
||||
usageToday: 'Today',
|
||||
usageYesterday: 'Yesterday',
|
||||
usageTotal: 'Total',
|
||||
accountsAvailable: 'Avail:',
|
||||
accountsRateLimited: 'Limited:',
|
||||
|
||||
@@ -814,6 +814,7 @@ export default {
|
||||
userStatus: '状态'
|
||||
},
|
||||
usageToday: '今日',
|
||||
usageYesterday: '昨日',
|
||||
usageTotal: '累计',
|
||||
accountsAvailable: '可用:',
|
||||
accountsRateLimited: '限流:',
|
||||
|
||||
@@ -333,6 +333,16 @@
|
||||
}}</span
|
||||
>
|
||||
</div>
|
||||
<div class="text-gray-500 dark:text-gray-400">
|
||||
<span class="text-gray-400 dark:text-gray-500">{{
|
||||
t("admin.groups.usageYesterday")
|
||||
}}</span>
|
||||
<span class="ml-1 font-medium text-gray-700 dark:text-gray-300"
|
||||
>${{
|
||||
formatCost(usageMap.get(row.id)?.yesterday_cost ?? 0)
|
||||
}}</span
|
||||
>
|
||||
</div>
|
||||
<div class="text-gray-500 dark:text-gray-400">
|
||||
<span class="text-gray-400 dark:text-gray-500">{{
|
||||
t("admin.groups.usageTotal")
|
||||
@@ -4894,6 +4904,7 @@ const groups = ref<AdminGroup[]>([]);
|
||||
const loading = ref(false);
|
||||
type GroupUsageSummary = {
|
||||
today_cost: number;
|
||||
yesterday_cost: number;
|
||||
total_cost: number;
|
||||
};
|
||||
|
||||
@@ -5739,12 +5750,12 @@ const loadUsageSummary = async () => {
|
||||
}
|
||||
usageLoading.value = true;
|
||||
try {
|
||||
const tz = Intl.DateTimeFormat().resolvedOptions().timeZone;
|
||||
const data = await adminAPI.groups.getUsageSummary(tz);
|
||||
const data = await adminAPI.groups.getUsageSummary();
|
||||
const map = new Map<number, GroupUsageSummary>();
|
||||
for (const item of data) {
|
||||
map.set(item.group_id, {
|
||||
today_cost: item.today_cost,
|
||||
yesterday_cost: item.yesterday_cost,
|
||||
total_cost: item.total_cost,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -43,6 +43,9 @@ const messages: Record<string, string> = {
|
||||
'admin.groups.columns.usage': 'Usage',
|
||||
'admin.groups.columns.status': 'Status',
|
||||
'admin.groups.columns.actions': 'Actions',
|
||||
'admin.groups.usageToday': 'Today',
|
||||
'admin.groups.usageYesterday': 'Yesterday',
|
||||
'admin.groups.usageTotal': 'Total',
|
||||
}
|
||||
|
||||
vi.mock('@/api/admin', () => ({
|
||||
@@ -151,6 +154,9 @@ const DataTableStub = {
|
||||
<div>
|
||||
<div data-test="columns">{{ columns.map((col) => col.key).join(',') }}</div>
|
||||
<div data-test="rows">{{ data.map((row) => row.name).join(',') }}</div>
|
||||
<div v-if="data.length" data-test="usage-cell">
|
||||
<slot name="cell-usage" :row="data[0]" />
|
||||
</div>
|
||||
</div>
|
||||
`,
|
||||
}
|
||||
@@ -378,10 +384,26 @@ describe('admin GroupsView column settings', () => {
|
||||
await openColumnSettings(wrapper)
|
||||
await clickColumnToggle(wrapper, 'Usage')
|
||||
expect(getUsageSummary).toHaveBeenCalledTimes(1)
|
||||
expect(getUsageSummary).toHaveBeenCalledWith()
|
||||
expect(getCapacitySummary).not.toHaveBeenCalled()
|
||||
|
||||
await clickColumnToggle(wrapper, 'Capacity')
|
||||
expect(getUsageSummary).toHaveBeenCalledTimes(1)
|
||||
expect(getCapacitySummary).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
it('renders yesterday usage between today and total', async () => {
|
||||
getUsageSummary.mockResolvedValue([
|
||||
{ group_id: 1, today_cost: 1.25, yesterday_cost: 2.5, total_cost: 9.75 },
|
||||
])
|
||||
|
||||
const wrapper = await mountView()
|
||||
const text = wrapper.get('[data-test="usage-cell"]').text()
|
||||
|
||||
expect(text).toContain('Today$1.25')
|
||||
expect(text).toContain('Yesterday$2.50')
|
||||
expect(text).toContain('Total$9.75')
|
||||
expect(text.indexOf('Today')).toBeLessThan(text.indexOf('Yesterday'))
|
||||
expect(text.indexOf('Yesterday')).toBeLessThan(text.indexOf('Total'))
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user