diff --git a/.github/workflows/backend-ci.yml b/.github/workflows/backend-ci.yml index 21a4fea79c..6c4a028c96 100644 --- a/.github/workflows/backend-ci.yml +++ b/.github/workflows/backend-ci.yml @@ -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: diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 2ba01833c1..4ecc8a26ba 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -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 diff --git a/.github/workflows/security-scan.yml b/.github/workflows/security-scan.yml index 96a7ae4edd..0e68d950ac 100644 --- a/.github/workflows/security-scan.yml +++ b/.github/workflows/security-scan.yml @@ -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: | diff --git a/backend/go.mod b/backend/go.mod index d6f1c10ed6..d5188b3ca7 100644 --- a/backend/go.mod +++ b/backend/go.mod @@ -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 diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index 0ed2b17213..2c5d5cf545 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -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) } diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index 718913a044..d8cfe820b1 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -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) diff --git a/backend/internal/handler/admin/group_handler.go b/backend/internal/handler/admin/group_handler.go index b8b9587f72..dfbd4eba6f 100644 --- a/backend/internal/handler/admin/group_handler.go +++ b/backend/internal/handler/admin/group_handler.go @@ -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 { diff --git a/backend/internal/pkg/usagestats/usage_log_types.go b/backend/internal/pkg/usagestats/usage_log_types.go index a100b2ccaa..ad97026e2b 100644 --- a/backend/internal/pkg/usagestats/usage_log_types.go +++ b/backend/internal/pkg/usagestats/usage_log_types.go @@ -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 diff --git a/backend/internal/repository/custom_group_usage_rollup_repo.go b/backend/internal/repository/custom_group_usage_rollup_repo.go new file mode 100644 index 0000000000..ac30edfc19 --- /dev/null +++ b/backend/internal/repository/custom_group_usage_rollup_repo.go @@ -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 +} diff --git a/backend/internal/repository/custom_group_usage_timezone_test.go b/backend/internal/repository/custom_group_usage_timezone_test.go new file mode 100644 index 0000000000..c77c5b41bc --- /dev/null +++ b/backend/internal/repository/custom_group_usage_timezone_test.go @@ -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)) }) +} diff --git a/backend/internal/repository/dashboard_aggregation_group_usage_test.go b/backend/internal/repository/dashboard_aggregation_group_usage_test.go new file mode 100644 index 0000000000..cbb5d3c6b6 --- /dev/null +++ b/backend/internal/repository/dashboard_aggregation_group_usage_test.go @@ -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") +} diff --git a/backend/internal/repository/dashboard_aggregation_repo.go b/backend/internal/repository/dashboard_aggregation_repo.go index 5e09e75d40..6701a3c999 100644 --- a/backend/internal/repository/dashboard_aggregation_repo.go +++ b/backend/internal/repository/dashboard_aggregation_repo.go @@ -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 { diff --git a/backend/internal/repository/group_usage_rollup_trigger_integration_test.go b/backend/internal/repository/group_usage_rollup_trigger_integration_test.go new file mode 100644 index 0000000000..db182be495 --- /dev/null +++ b/backend/internal/repository/group_usage_rollup_trigger_integration_test.go @@ -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() + } + } +} diff --git a/backend/internal/repository/usage_cleanup_repo.go b/backend/internal/repository/usage_cleanup_repo.go index 1a25696e4c..f8aa67477f 100644 --- a/backend/internal/repository/usage_cleanup_repo.go +++ b/backend/internal/repository/usage_cleanup_repo.go @@ -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) diff --git a/backend/internal/repository/usage_cleanup_repo_test.go b/backend/internal/repository/usage_cleanup_repo_test.go index 1ac7cca569..b147d03a4f 100644 --- a/backend/internal/repository/usage_cleanup_repo_test.go +++ b/backend/internal/repository/usage_cleanup_repo_test.go @@ -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) diff --git a/backend/internal/repository/usage_log_repo_group_summary_test.go b/backend/internal/repository/usage_log_repo_group_summary_test.go new file mode 100644 index 0000000000..f67facb4b9 --- /dev/null +++ b/backend/internal/repository/usage_log_repo_group_summary_test.go @@ -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()) +} diff --git a/backend/internal/repository/usage_log_repo_trend.go b/backend/internal/repository/usage_log_repo_trend.go index 6d8bd79eaa..3ed3aa51b1 100644 --- a/backend/internal/repository/usage_log_repo_trend.go +++ b/backend/internal/repository/usage_log_repo_trend.go @@ -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. diff --git a/backend/internal/service/custom_group_usage_rollup.go b/backend/internal/service/custom_group_usage_rollup.go new file mode 100644 index 0000000000..d14dc9bfd9 --- /dev/null +++ b/backend/internal/service/custom_group_usage_rollup.go @@ -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) +} diff --git a/backend/internal/service/custom_group_usage_rollup_test.go b/backend/internal/service/custom_group_usage_rollup_test.go new file mode 100644 index 0000000000..35c568a456 --- /dev/null +++ b/backend/internal/service/custom_group_usage_rollup_test.go @@ -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()) +} diff --git a/backend/internal/service/dashboard_aggregation_service.go b/backend/internal/service/dashboard_aggregation_service.go index 5dfa02a824..6e351d75d2 100644 --- a/backend/internal/service/dashboard_aggregation_service.go +++ b/backend/internal/service/dashboard_aggregation_service.go @@ -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 diff --git a/backend/internal/service/dashboard_aggregation_service_test.go b/backend/internal/service/dashboard_aggregation_service_test.go index fbb671bb66..d164519330 100644 --- a/backend/internal/service/dashboard_aggregation_service_test.go +++ b/backend/internal/service/dashboard_aggregation_service_test.go @@ -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{ diff --git a/backend/internal/service/dashboard_service.go b/backend/internal/service/dashboard_service.go index d815ed6ef5..3c656b61f5 100644 --- a/backend/internal/service/dashboard_service.go +++ b/backend/internal/service/dashboard_service.go @@ -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 { diff --git a/backend/migrations/222_group_usage_daily_rollups.sql b/backend/migrations/222_group_usage_daily_rollups.sql new file mode 100644 index 0000000000..c16b69b4ae --- /dev/null +++ b/backend/migrations/222_group_usage_daily_rollups.sql @@ -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(); diff --git a/backend/migrations/223_group_usage_rollup_timezone.sql b/backend/migrations/223_group_usage_rollup_timezone.sql new file mode 100644 index 0000000000..8de2bdff97 --- /dev/null +++ b/backend/migrations/223_group_usage_rollup_timezone.sql @@ -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; +$$; diff --git a/backend/migrations/group_usage_rollup_migration_test.go b/backend/migrations/group_usage_rollup_migration_test.go new file mode 100644 index 0000000000..f2e9a5d0ad --- /dev/null +++ b/backend/migrations/group_usage_rollup_migration_test.go @@ -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") +} diff --git a/frontend/pnpm-lock.yaml b/frontend/pnpm-lock.yaml index e17969c9cc..4c86452cca 100644 --- a/frontend/pnpm-lock.yaml +++ b/frontend/pnpm-lock.yaml @@ -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 diff --git a/frontend/src/api/__tests__/admin.groups.usage-summary.spec.ts b/frontend/src/api/__tests__/admin.groups.usage-summary.spec.ts new file mode 100644 index 0000000000..c0b2c4bffc --- /dev/null +++ b/frontend/src/api/__tests__/admin.groups.usage-summary.spec.ts @@ -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') + }) +}) diff --git a/frontend/src/api/admin/groups.ts b/frontend/src/api/admin/groups.ts index 3828bf9d71..2c9c80f5e6 100644 --- a/frontend/src/api/admin/groups.ts +++ b/frontend/src/api/admin/groups.ts @@ -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 } diff --git a/frontend/src/i18n/locales/en/admin/overview.ts b/frontend/src/i18n/locales/en/admin/overview.ts index 1fcc7bf0e8..225865db9d 100644 --- a/frontend/src/i18n/locales/en/admin/overview.ts +++ b/frontend/src/i18n/locales/en/admin/overview.ts @@ -821,6 +821,7 @@ export default { userStatus: 'Status' }, usageToday: 'Today', + usageYesterday: 'Yesterday', usageTotal: 'Total', accountsAvailable: 'Avail:', accountsRateLimited: 'Limited:', diff --git a/frontend/src/i18n/locales/zh/admin/overview.ts b/frontend/src/i18n/locales/zh/admin/overview.ts index f4b02a9cab..200bb6c547 100644 --- a/frontend/src/i18n/locales/zh/admin/overview.ts +++ b/frontend/src/i18n/locales/zh/admin/overview.ts @@ -814,6 +814,7 @@ export default { userStatus: '状态' }, usageToday: '今日', + usageYesterday: '昨日', usageTotal: '累计', accountsAvailable: '可用:', accountsRateLimited: '限流:', diff --git a/frontend/src/views/admin/GroupsView.vue b/frontend/src/views/admin/GroupsView.vue index 5cb036f7bc..904b8926d4 100644 --- a/frontend/src/views/admin/GroupsView.vue +++ b/frontend/src/views/admin/GroupsView.vue @@ -333,6 +333,16 @@ }} +
+ {{ + t("admin.groups.usageYesterday") + }} + ${{ + formatCost(usageMap.get(row.id)?.yesterday_cost ?? 0) + }} +
{{ t("admin.groups.usageTotal") @@ -4894,6 +4904,7 @@ const groups = ref([]); 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(); 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, }); } diff --git a/frontend/src/views/admin/__tests__/GroupsView.columnSettings.spec.ts b/frontend/src/views/admin/__tests__/GroupsView.columnSettings.spec.ts index b69cc6bd52..8d4f8df74f 100644 --- a/frontend/src/views/admin/__tests__/GroupsView.columnSettings.spec.ts +++ b/frontend/src/views/admin/__tests__/GroupsView.columnSettings.spec.ts @@ -43,6 +43,9 @@ const messages: Record = { '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 = {
{{ columns.map((col) => col.key).join(',') }}
{{ data.map((row) => row.name).join(',') }}
+
+ +
`, } @@ -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')) + }) })