mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add Prometheus metrics for AI Governance cost control (#27490)
## Description Adds Prometheus metrics for AI budget cost control, emitted by the aibridged server under the `cost_control` subsystem (full names are prefixed `coder_ai_gateway_`). - `blocked_requests_total` (counter) — labels: `group_id` - `blocked_users` (gauge) — labels: `group_id` - `unpriced_requests_total` (counter) — labels: `provider`, `model` - `enforcement_duration_seconds` (histogram) — labels: `outcome` ## Changes - Add `GetOverBudgetUsersPerGroup` query (plus dbauthz/dbmetrics/dbmock wiring) to count over-budget users per effective group. - Add a background collector that refreshes the `blocked_users` gauge on an interval, started only when Prometheus is enabled. - Wire `Metrics` through the aibridged server, coderd API, `cli/server.go`, and the enterprise AI gateway handler; recording is nil-safe when metrics are unset. Closes https://linear.app/codercom/issue/AIGOV-296/add-prometheus-metrics-for-cost-control > [!NOTE] > Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
This commit is contained in:
Susana Ferreira
@@ -4321,6 +4321,14 @@ func (q *querier) GetOrganizationsWithPrebuildStatus(ctx context.Context, arg da
|
||||
return q.db.GetOrganizationsWithPrebuildStatus(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetOverBudgetUsersPerGroup(ctx context.Context, periodStart time.Time) ([]database.GetOverBudgetUsersPerGroupRow, error) {
|
||||
// Aggregates over-budget user counts per group for cost-control metrics.
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceGroup.All()); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetOverBudgetUsersPerGroup(ctx, periodStart)
|
||||
}
|
||||
|
||||
func (q *querier) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]database.ParameterSchema, error) {
|
||||
version, err := q.db.GetTemplateVersionByJobID(ctx, jobID)
|
||||
if err != nil {
|
||||
|
||||
@@ -7073,6 +7073,13 @@ func (s *MethodTestSuite) TestAIBridge() {
|
||||
check.Args(arg).Asserts(user, policy.ActionRead).Returns(row)
|
||||
}))
|
||||
|
||||
s.Run("GetOverBudgetUsersPerGroup", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
periodStart := time.Now().UTC().Truncate(24 * time.Hour)
|
||||
dbm.EXPECT().GetOverBudgetUsersPerGroup(gomock.Any(), periodStart).
|
||||
Return([]database.GetOverBudgetUsersPerGroupRow{}, nil).AnyTimes()
|
||||
check.Args(periodStart).Asserts(rbac.ResourceGroup.All(), policy.ActionRead)
|
||||
}))
|
||||
|
||||
s.Run("IncrementUserAIDailySpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
arg := database.IncrementUserAIDailySpendParams{
|
||||
UserID: uuid.New(),
|
||||
|
||||
+8
@@ -2609,6 +2609,14 @@ func (m queryMetricsStore) GetOrganizationsWithPrebuildStatus(ctx context.Contex
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetOverBudgetUsersPerGroup(ctx context.Context, periodStart time.Time) ([]database.GetOverBudgetUsersPerGroupRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetOverBudgetUsersPerGroup(ctx, periodStart)
|
||||
m.queryLatencies.WithLabelValues("GetOverBudgetUsersPerGroup").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetOverBudgetUsersPerGroup").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]database.ParameterSchema, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetParameterSchemasByJobID(ctx, jobID)
|
||||
|
||||
Generated
+15
@@ -4843,6 +4843,21 @@ func (mr *MockStoreMockRecorder) GetOrganizationsWithPrebuildStatus(ctx, arg any
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetOrganizationsWithPrebuildStatus", reflect.TypeOf((*MockStore)(nil).GetOrganizationsWithPrebuildStatus), ctx, arg)
|
||||
}
|
||||
|
||||
// GetOverBudgetUsersPerGroup mocks base method.
|
||||
func (m *MockStore) GetOverBudgetUsersPerGroup(ctx context.Context, periodStart time.Time) ([]database.GetOverBudgetUsersPerGroupRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetOverBudgetUsersPerGroup", ctx, periodStart)
|
||||
ret0, _ := ret[0].([]database.GetOverBudgetUsersPerGroupRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetOverBudgetUsersPerGroup indicates an expected call of GetOverBudgetUsersPerGroup.
|
||||
func (mr *MockStoreMockRecorder) GetOverBudgetUsersPerGroup(ctx, periodStart any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetOverBudgetUsersPerGroup", reflect.TypeOf((*MockStore)(nil).GetOverBudgetUsersPerGroup), ctx, periodStart)
|
||||
}
|
||||
|
||||
// GetParameterSchemasByJobID mocks base method.
|
||||
func (m *MockStore) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]database.ParameterSchema, error) {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
Generated
+5
@@ -690,6 +690,11 @@ type sqlcQuerier interface {
|
||||
// GetOrganizationsWithPrebuildStatus returns organizations with prebuilds configured and their
|
||||
// membership status for the prebuilds system user (org membership, group existence, group membership).
|
||||
GetOrganizationsWithPrebuildStatus(ctx context.Context, arg GetOrganizationsWithPrebuildStatusParams) ([]GetOrganizationsWithPrebuildStatusRow, error)
|
||||
// Returns, per effective group, the number of users at or over their spend
|
||||
// limit since period_start. Only users with an enforceable limit (override or
|
||||
// budgeted group) count, and the unlimited Everyone fallback does not.
|
||||
// TODO(AIGOV-527): unify effective group resolution in a single place.
|
||||
GetOverBudgetUsersPerGroup(ctx context.Context, periodStart time.Time) ([]GetOverBudgetUsersPerGroupRow, error)
|
||||
GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]ParameterSchema, error)
|
||||
GetPrebuildMetrics(ctx context.Context) ([]GetPrebuildMetricsRow, error)
|
||||
GetPrebuildsSettings(ctx context.Context) (string, error)
|
||||
|
||||
@@ -13901,6 +13901,197 @@ func TestGetHighestGroupAIBudgetByUser(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetOverBudgetUsersPerGroup(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
periodStart := dbtime.Now().UTC().Truncate(24 * time.Hour)
|
||||
|
||||
// seedSpendOnDay attributes micros of spend to (user, effectiveGroup) on a
|
||||
// specific day.
|
||||
seedSpendOnDay := func(t *testing.T, ctx context.Context, db database.Store, userID, effectiveGroupID uuid.UUID, day time.Time, micros int64) {
|
||||
t.Helper()
|
||||
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
|
||||
UserID: userID,
|
||||
EffectiveGroupID: effectiveGroupID,
|
||||
Day: day,
|
||||
CostMicros: micros,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// seedSpend attributes micros of spend to (user, effectiveGroup) within the
|
||||
// current period.
|
||||
seedSpend := func(t *testing.T, ctx context.Context, db database.Store, userID, effectiveGroupID uuid.UUID, micros int64) {
|
||||
t.Helper()
|
||||
seedSpendOnDay(t, ctx, db, userID, effectiveGroupID, periodStart, micros)
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
setup func(t *testing.T, ctx context.Context, db database.Store) []database.GetOverBudgetUsersPerGroupRow
|
||||
}{
|
||||
{
|
||||
// A user whose spend exceeds their group budget is counted.
|
||||
name: "OverBudgetCounted",
|
||||
setup: func(t *testing.T, ctx context.Context, db database.Store) []database.GetOverBudgetUsersPerGroupRow {
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: group.ID, SpendLimitMicros: 1_000_000})
|
||||
require.NoError(t, err)
|
||||
seedSpend(t, ctx, db, user.ID, group.ID, 1_500_000)
|
||||
return []database.GetOverBudgetUsersPerGroupRow{{GroupID: group.ID, OverBudgetUsers: 1}}
|
||||
},
|
||||
},
|
||||
{
|
||||
// A user under their group budget is not counted.
|
||||
name: "UnderBudgetNotCounted",
|
||||
setup: func(t *testing.T, ctx context.Context, db database.Store) []database.GetOverBudgetUsersPerGroupRow {
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: group.ID, SpendLimitMicros: 1_000_000})
|
||||
require.NoError(t, err)
|
||||
seedSpend(t, ctx, db, user.ID, group.ID, 500_000)
|
||||
return nil
|
||||
},
|
||||
},
|
||||
{
|
||||
// Spend exactly at the limit counts, since the check is inclusive.
|
||||
name: "AtLimitCounted",
|
||||
setup: func(t *testing.T, ctx context.Context, db database.Store) []database.GetOverBudgetUsersPerGroupRow {
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: group.ID, SpendLimitMicros: 1_000_000})
|
||||
require.NoError(t, err)
|
||||
seedSpend(t, ctx, db, user.ID, group.ID, 1_000_000)
|
||||
return []database.GetOverBudgetUsersPerGroupRow{{GroupID: group.ID, OverBudgetUsers: 1}}
|
||||
},
|
||||
},
|
||||
{
|
||||
// A zero limit blocks a user with no spend, since zero spend is at the
|
||||
// limit.
|
||||
name: "ZeroLimitCounted",
|
||||
setup: func(t *testing.T, ctx context.Context, db database.Store) []database.GetOverBudgetUsersPerGroupRow {
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: group.ID, SpendLimitMicros: 0})
|
||||
require.NoError(t, err)
|
||||
return []database.GetOverBudgetUsersPerGroupRow{{GroupID: group.ID, OverBudgetUsers: 1}}
|
||||
},
|
||||
},
|
||||
{
|
||||
// A per-user override overrides the group budget, both for the limit
|
||||
// and the group the spend is attributed to.
|
||||
name: "OverrideWins",
|
||||
setup: func(t *testing.T, ctx context.Context, db database.Store) []database.GetOverBudgetUsersPerGroupRow {
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
overrideGroup := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: overrideGroup.ID, UserID: user.ID})
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: group.ID, SpendLimitMicros: 5_000_000})
|
||||
require.NoError(t, err)
|
||||
_, err = db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{UserID: user.ID, GroupID: overrideGroup.ID, SpendLimitMicros: 1_000_000})
|
||||
require.NoError(t, err)
|
||||
// Over the override limit but under the group limit.
|
||||
seedSpend(t, ctx, db, user.ID, overrideGroup.ID, 1_500_000)
|
||||
return []database.GetOverBudgetUsersPerGroupRow{{GroupID: overrideGroup.ID, OverBudgetUsers: 1}}
|
||||
},
|
||||
},
|
||||
{
|
||||
// A user in multiple budgeted groups is attributed to their
|
||||
// highest-limit group.
|
||||
name: "HighestGroupWins",
|
||||
setup: func(t *testing.T, ctx context.Context, db database.Store) []database.GetOverBudgetUsersPerGroupRow {
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
lower := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
higher := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: lower.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: higher.ID, UserID: user.ID})
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: lower.ID, SpendLimitMicros: 1_000_000})
|
||||
require.NoError(t, err)
|
||||
_, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: higher.ID, SpendLimitMicros: 2_000_000})
|
||||
require.NoError(t, err)
|
||||
seedSpend(t, ctx, db, user.ID, higher.ID, 2_000_000)
|
||||
return []database.GetOverBudgetUsersPerGroupRow{{GroupID: higher.ID, OverBudgetUsers: 1}}
|
||||
},
|
||||
},
|
||||
{
|
||||
// A user with only the unlimited Everyone fallback is never counted.
|
||||
name: "EveryoneFallbackNotCounted",
|
||||
setup: func(t *testing.T, ctx context.Context, db database.Store) []database.GetOverBudgetUsersPerGroupRow {
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
|
||||
// Spend attributed to the Everyone group (id == organization_id).
|
||||
seedSpend(t, ctx, db, user.ID, org.ID, 9_000_000)
|
||||
return nil
|
||||
},
|
||||
},
|
||||
{
|
||||
// Multiple over-budget users in the same group are summed.
|
||||
name: "AggregatesUsersPerGroup",
|
||||
setup: func(t *testing.T, ctx context.Context, db database.Store) []database.GetOverBudgetUsersPerGroupRow {
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: group.ID, SpendLimitMicros: 1_000_000})
|
||||
require.NoError(t, err)
|
||||
for range 2 {
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
seedSpend(t, ctx, db, user.ID, group.ID, 2_000_000)
|
||||
}
|
||||
return []database.GetOverBudgetUsersPerGroupRow{{GroupID: group.ID, OverBudgetUsers: 2}}
|
||||
},
|
||||
},
|
||||
{
|
||||
// Spend on days before the period start is excluded, so a user whose
|
||||
// only over-limit spend predates the period is not counted.
|
||||
name: "SpendBeforePeriodNotCounted",
|
||||
setup: func(t *testing.T, ctx context.Context, db database.Store) []database.GetOverBudgetUsersPerGroupRow {
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
|
||||
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
|
||||
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{GroupID: group.ID, SpendLimitMicros: 1_000_000})
|
||||
require.NoError(t, err)
|
||||
seedSpendOnDay(t, ctx, db, user.ID, group.ID, periodStart.AddDate(0, 0, -1), 1_500_000)
|
||||
return nil
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
want := tt.setup(t, ctx, db)
|
||||
got, err := db.GetOverBudgetUsersPerGroup(ctx, periodStart)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetUserEveryoneFallbackGroup(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
Generated
+91
@@ -2819,6 +2819,97 @@ func (q *sqlQuerier) GetOrganizationGroupsAISpend(ctx context.Context, arg GetOr
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const getOverBudgetUsersPerGroup = `-- name: GetOverBudgetUsersPerGroup :many
|
||||
WITH budgeted_users AS (
|
||||
-- Users with an override or membership in a budgeted group.
|
||||
SELECT user_id FROM user_ai_budget_overrides
|
||||
UNION
|
||||
SELECT DISTINCT member.user_id
|
||||
FROM group_ai_budgets budget
|
||||
JOIN group_members_expanded member ON member.group_id = budget.group_id
|
||||
),
|
||||
user_highest_group AS (
|
||||
-- Per user, their highest-limit group ("highest" budget policy).
|
||||
SELECT DISTINCT ON (member.user_id)
|
||||
member.user_id,
|
||||
budget.group_id,
|
||||
budget.spend_limit_micros
|
||||
FROM group_ai_budgets budget
|
||||
JOIN group_members_expanded member ON member.group_id = budget.group_id
|
||||
JOIN organizations ON organizations.id = member.organization_id
|
||||
JOIN organization_members
|
||||
ON organization_members.user_id = member.user_id
|
||||
AND organization_members.organization_id = member.organization_id
|
||||
WHERE member.user_id IN (SELECT user_id FROM budgeted_users)
|
||||
AND organizations.deleted = false
|
||||
ORDER BY member.user_id, budget.spend_limit_micros DESC, organization_members.created_at ASC, budget.group_id ASC
|
||||
),
|
||||
effective AS (
|
||||
-- An override wins over the highest-limit group, and users with neither drop.
|
||||
SELECT
|
||||
budgeted_users.user_id,
|
||||
COALESCE(override.group_id, user_highest_group.group_id) AS effective_group_id,
|
||||
COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros
|
||||
FROM budgeted_users
|
||||
LEFT JOIN user_ai_budget_overrides override ON override.user_id = budgeted_users.user_id
|
||||
LEFT JOIN user_highest_group ON user_highest_group.user_id = budgeted_users.user_id
|
||||
WHERE COALESCE(override.group_id, user_highest_group.group_id) IS NOT NULL
|
||||
),
|
||||
user_spend AS (
|
||||
-- Each user's spend against their effective group since period_start.
|
||||
SELECT
|
||||
effective.user_id,
|
||||
effective.effective_group_id,
|
||||
effective.spend_limit_micros,
|
||||
COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros
|
||||
FROM effective
|
||||
LEFT JOIN ai_user_daily_spend spend
|
||||
ON spend.user_id = effective.user_id
|
||||
AND spend.effective_group_id = effective.effective_group_id
|
||||
AND spend.day >= (($1::timestamptz) AT TIME ZONE 'UTC')::date
|
||||
GROUP BY effective.user_id, effective.effective_group_id, effective.spend_limit_micros
|
||||
)
|
||||
SELECT
|
||||
effective_group_id AS group_id,
|
||||
COUNT(*)::BIGINT AS over_budget_users
|
||||
FROM user_spend
|
||||
WHERE current_spend_micros >= spend_limit_micros
|
||||
GROUP BY effective_group_id
|
||||
ORDER BY effective_group_id
|
||||
`
|
||||
|
||||
type GetOverBudgetUsersPerGroupRow struct {
|
||||
GroupID uuid.UUID `db:"group_id" json:"group_id"`
|
||||
OverBudgetUsers int64 `db:"over_budget_users" json:"over_budget_users"`
|
||||
}
|
||||
|
||||
// Returns, per effective group, the number of users at or over their spend
|
||||
// limit since period_start. Only users with an enforceable limit (override or
|
||||
// budgeted group) count, and the unlimited Everyone fallback does not.
|
||||
// TODO(AIGOV-527): unify effective group resolution in a single place.
|
||||
func (q *sqlQuerier) GetOverBudgetUsersPerGroup(ctx context.Context, periodStart time.Time) ([]GetOverBudgetUsersPerGroupRow, error) {
|
||||
rows, err := q.db.QueryContext(ctx, getOverBudgetUsersPerGroup, periodStart)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var items []GetOverBudgetUsersPerGroupRow
|
||||
for rows.Next() {
|
||||
var i GetOverBudgetUsersPerGroupRow
|
||||
if err := rows.Scan(&i.GroupID, &i.OverBudgetUsers); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, i)
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const getUserAIBudgetOverride = `-- name: GetUserAIBudgetOverride :one
|
||||
SELECT user_id, group_id, spend_limit_micros, created_at, updated_at
|
||||
FROM user_ai_budget_overrides
|
||||
|
||||
@@ -243,3 +243,65 @@ GROUP BY
|
||||
applied_budget.spend_limit_micros,
|
||||
applied_budget.limit_source
|
||||
ORDER BY effective.user_id;
|
||||
|
||||
-- name: GetOverBudgetUsersPerGroup :many
|
||||
-- Returns, per effective group, the number of users at or over their spend
|
||||
-- limit since period_start. Only users with an enforceable limit (override or
|
||||
-- budgeted group) count, and the unlimited Everyone fallback does not.
|
||||
-- TODO(AIGOV-527): unify effective group resolution in a single place.
|
||||
WITH budgeted_users AS (
|
||||
-- Users with an override or membership in a budgeted group.
|
||||
SELECT user_id FROM user_ai_budget_overrides
|
||||
UNION
|
||||
SELECT DISTINCT member.user_id
|
||||
FROM group_ai_budgets budget
|
||||
JOIN group_members_expanded member ON member.group_id = budget.group_id
|
||||
),
|
||||
user_highest_group AS (
|
||||
-- Per user, their highest-limit group ("highest" budget policy).
|
||||
SELECT DISTINCT ON (member.user_id)
|
||||
member.user_id,
|
||||
budget.group_id,
|
||||
budget.spend_limit_micros
|
||||
FROM group_ai_budgets budget
|
||||
JOIN group_members_expanded member ON member.group_id = budget.group_id
|
||||
JOIN organizations ON organizations.id = member.organization_id
|
||||
JOIN organization_members
|
||||
ON organization_members.user_id = member.user_id
|
||||
AND organization_members.organization_id = member.organization_id
|
||||
WHERE member.user_id IN (SELECT user_id FROM budgeted_users)
|
||||
AND organizations.deleted = false
|
||||
ORDER BY member.user_id, budget.spend_limit_micros DESC, organization_members.created_at ASC, budget.group_id ASC
|
||||
),
|
||||
effective AS (
|
||||
-- An override wins over the highest-limit group, and users with neither drop.
|
||||
SELECT
|
||||
budgeted_users.user_id,
|
||||
COALESCE(override.group_id, user_highest_group.group_id) AS effective_group_id,
|
||||
COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros
|
||||
FROM budgeted_users
|
||||
LEFT JOIN user_ai_budget_overrides override ON override.user_id = budgeted_users.user_id
|
||||
LEFT JOIN user_highest_group ON user_highest_group.user_id = budgeted_users.user_id
|
||||
WHERE COALESCE(override.group_id, user_highest_group.group_id) IS NOT NULL
|
||||
),
|
||||
user_spend AS (
|
||||
-- Each user's spend against their effective group since period_start.
|
||||
SELECT
|
||||
effective.user_id,
|
||||
effective.effective_group_id,
|
||||
effective.spend_limit_micros,
|
||||
COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros
|
||||
FROM effective
|
||||
LEFT JOIN ai_user_daily_spend spend
|
||||
ON spend.user_id = effective.user_id
|
||||
AND spend.effective_group_id = effective.effective_group_id
|
||||
AND spend.day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date
|
||||
GROUP BY effective.user_id, effective.effective_group_id, effective.spend_limit_micros
|
||||
)
|
||||
SELECT
|
||||
effective_group_id AS group_id,
|
||||
COUNT(*)::BIGINT AS over_budget_users
|
||||
FROM user_spend
|
||||
WHERE current_spend_micros >= spend_limit_micros
|
||||
GROUP BY effective_group_id
|
||||
ORDER BY effective_group_id;
|
||||
|
||||
Reference in New Issue
Block a user