opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: report combined member limit in group AI spend (#27589)

Browse Source 
## Problem

The organization groups page showed each group's AI budget as the
group's per-member limit, so the total it displayed was effectively
group members × group budget. That ignores per-user budget overrides
charged to the group, so a group where one member has an override
reported a limit that doesn't match what its members can actually spend.

## Changes

- Add `total_spend_limit_micros` to the organization groups AI spend
payload, the combined budget of the members attributed to the group,
with each member's override replacing their share.
- Return `null` for the total when the group has no budget, since its
members spend without a cap.
- Both the organization groups and single group spend endpoints report
the new field, as they share the same query.
- Use the total as the denominator on the groups page AI budget column.

Depends on #27568
This commit is contained in:
Susana Ferreira
2026-07-29 09:16:37 +01:00
committed by GitHub
parent e96d8646e2
commit 0b4095085e
14 changed files with 705 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/db2sdk/db2sdk.go
+3
View File
@@ -1475,6 +1475,9 @@ func OrganizationGroupAISpend(row database.GetOrganizationGroupsAISpendRow) code
if row.SpendLimitMicros.Valid {
group.SpendLimitMicros = &row.SpendLimitMicros.Int64
}
if row.TotalSpendLimitMicros.Valid {
group.TotalSpendLimitMicros = &row.TotalSpendLimitMicros.Int64
}
return group
}
coderd/database/querier.go
Generated
+6 -2
View File
@@ -691,9 +691,13 @@ type sqlcQuerier interface {
GetOrganizationByID(ctx context.Context, id uuid.UUID) (Organization, error)
GetOrganizationByName(ctx context.Context, arg GetOrganizationByNameParams) (Organization, error)
// Returns AI spend limits and aggregate spend for groups in @group_ids that
// belong to @organization_id, on or after period_start until NOW. The spend
// limit is null when the group has no configured budget.
// belong to @organization_id, on or after period_start until NOW.
// spend_limit_micros is the per-member limit, null when the group has no budget.
// total_spend_limit_micros is the combined budget of the members attributed to
// the group, with each member's override replacing their share. It is null when
// the group has no budget.
// The period_start parameter is normalized to its UTC calendar day.
// TODO(AIGOV-527): unify effective group resolution in a single place.
GetOrganizationGroupsAISpend(ctx context.Context, arg GetOrganizationGroupsAISpendParams) ([]GetOrganizationGroupsAISpendRow, error)
GetOrganizationIDsByMemberIDs(ctx context.Context, ids []uuid.UUID) ([]GetOrganizationIDsByMemberIDsRow, error)
GetOrganizationResourceCountByID(ctx context.Context, organizationID uuid.UUID) (GetOrganizationResourceCountByIDRow, error)
coderd/database/querier_test.go
+301 -47
View File
@@ -13085,45 +13085,174 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) {
now := monthStart.AddDate(0, 0, 14) // 2024-06-15
prevMonthLastDay := monthStart.AddDate(0, 0, -1) // 2024-05-31
type seedRow struct {
// nullInt64 keeps the expectations below readable.
nullInt64 := func(v int64) sql.NullInt64 {
return sql.NullInt64{Int64: v, Valid: true}
}
type groupBudget struct {
group string
limit int64
}
type membership struct {
user string
group string
}
type override struct {
user string
group string
limit int64
}
type spendRow struct {
user string
group string
day time.Time
spend int64
}
type wantGroup struct {
spendLimit sql.NullInt64
totalLimit sql.NullInt64
spend int64
}
// A case declares its groups, members, overrides and spend by name. Every
// named group is created in one org and queried, and every named user is an
// org member.
tests := []struct {
name string
setBudget bool
spendLimit int64
rows []seedRow
wantCurrentSpend int64
name string
budgets []groupBudget
members []membership
overrides []override
spend []spendRow
want map[string]wantGroup
}{
{
name: "NoBudgetNoSpend",
wantCurrentSpend: 0,
name: "NoBudgetNoSpend",
want: map[string]wantGroup{
"eng": {spendLimit: sql.NullInt64{}, totalLimit: sql.NullInt64{}},
},
},
{
name: "ZeroLimitBudget",
setBudget: true,
spendLimit: 0,
wantCurrentSpend: 0,
name: "ZeroLimitBudget",
budgets: []groupBudget{{group: "eng", limit: 0}},
want: map[string]wantGroup{
"eng": {spendLimit: nullInt64(0), totalLimit: nullInt64(0)},
},
},
{
name: "BudgetZeroSpend",
setBudget: true,
spendLimit: 1_000_000,
wantCurrentSpend: 0,
// The group has no members, so nothing is attributed to it and the
// total is zero despite the budget.
name: "BudgetZeroSpend",
budgets: []groupBudget{{group: "eng", limit: 1_000_000}},
want: map[string]wantGroup{
"eng": {spendLimit: nullInt64(1_000_000), totalLimit: nullInt64(0)},
},
},
{
name: "BudgetWithSpend",
setBudget: true,
spendLimit: 1_000_000,
rows: []seedRow{{now, 250}},
wantCurrentSpend: 250,
// alice has spend attributed to the group but is currently not a member, so the
// total stays zero while the spend counts.
name: "BudgetWithSpend",
budgets: []groupBudget{{group: "eng", limit: 1_000_000}},
spend: []spendRow{{user: "alice", group: "eng", day: now, spend: 250}},
want: map[string]wantGroup{
"eng": {spendLimit: nullInt64(1_000_000), totalLimit: nullInt64(0), spend: 250},
},
},
{
name: "NoBudgetWithSpend",
rows: []seedRow{{now, 100}},
wantCurrentSpend: 100,
name: "NoBudgetWithSpend",
spend: []spendRow{{user: "alice", group: "eng", day: now, spend: 100}},
want: map[string]wantGroup{
"eng": {spendLimit: sql.NullInt64{}, totalLimit: sql.NullInt64{}, spend: 100},
},
},
{
name: "BudgetedTwoPlainMembers",
budgets: []groupBudget{{group: "eng", limit: 100}},
members: []membership{{user: "alice", group: "eng"}, {user: "bob", group: "eng"}},
want: map[string]wantGroup{
"eng": {spendLimit: nullInt64(100), totalLimit: nullInt64(200)},
},
},
{
name: "BudgetedPlainMemberPlusOverride",
budgets: []groupBudget{{group: "eng", limit: 100}},
members: []membership{{user: "alice", group: "eng"}, {user: "bob", group: "eng"}},
overrides: []override{{user: "alice", group: "eng", limit: 1000}},
// The override replaces its holder's share of the group limit.
want: map[string]wantGroup{
"eng": {spendLimit: nullInt64(100), totalLimit: nullInt64(1100)},
},
},
{
name: "BudgetedOnlyOverrideMember",
budgets: []groupBudget{{group: "eng", limit: 100}},
members: []membership{{user: "alice", group: "eng"}},
overrides: []override{{user: "alice", group: "eng", limit: 1000}},
want: map[string]wantGroup{
"eng": {spendLimit: nullInt64(100), totalLimit: nullInt64(1000)},
},
},
{
name: "BudgetedZeroLimitWithMembers",
budgets: []groupBudget{{group: "eng", limit: 0}},
members: []membership{{user: "alice", group: "eng"}, {user: "bob", group: "eng"}},
want: map[string]wantGroup{
"eng": {spendLimit: nullInt64(0), totalLimit: nullInt64(0)},
},
},
{
name: "UnbudgetedWithMembers",
members: []membership{{user: "alice", group: "eng"}, {user: "bob", group: "eng"}},
want: map[string]wantGroup{
"eng": {spendLimit: sql.NullInt64{}, totalLimit: sql.NullInt64{}},
},
},
{
name: "UnbudgetedWithOverride",
members: []membership{{user: "alice", group: "eng"}},
overrides: []override{{user: "alice", group: "eng", limit: 1000}},
// Members of a group with no budget spend without a cap, so the
// override does not make the group's total finite.
want: map[string]wantGroup{
"eng": {spendLimit: sql.NullInt64{}, totalLimit: sql.NullInt64{}},
},
},
{
name: "MemberResolvesToHigherBudgetGroup",
budgets: []groupBudget{
{group: "eng", limit: 100},
{group: "platform", limit: 500},
},
members: []membership{
{user: "alice", group: "eng"},
{user: "bob", group: "eng"},
{user: "alice", group: "platform"},
},
// alice is attributed to the higher-limit group, so only bob counts
// toward eng.
want: map[string]wantGroup{
"eng": {spendLimit: nullInt64(100), totalLimit: nullInt64(100)},
"platform": {spendLimit: nullInt64(500), totalLimit: nullInt64(500)},
},
},
{
name: "OverrideChargedToOtherGroup",
budgets: []groupBudget{
{group: "eng", limit: 100},
{group: "platform", limit: 50},
},
members: []membership{
{user: "alice", group: "eng"},
{user: "bob", group: "eng"},
{user: "alice", group: "platform"},
},
overrides: []override{{user: "alice", group: "platform", limit: 1000}},
// The override attributes alice to platform even though eng carries
// the higher group limit.
want: map[string]wantGroup{
"eng": {spendLimit: nullInt64(100), totalLimit: nullInt64(100)},
"platform": {spendLimit: nullInt64(50), totalLimit: nullInt64(1000)},
},
},
}
@@ -13133,46 +13262,92 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) {
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: an org with a single group, optionally with a budget and seeded spend.
user := dbgen.User(t, db, database.User{})
// Given: the groups, members, overrides and spend the case declares.
org := dbgen.Organization(t, db, database.Organization{})
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
if tt.setBudget {
groupIDs := make(map[string]uuid.UUID)
groupID := func(name string) uuid.UUID {
if id, ok := groupIDs[name]; ok {
return id
}
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
groupIDs[name] = group.ID
return group.ID
}
userIDs := make(map[string]uuid.UUID)
userID := func(name string) uuid.UUID {
if id, ok := userIDs[name]; ok {
return id
}
user := dbgen.User(t, db, database.User{})
dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: org.ID,
})
userIDs[name] = user.ID
return user.ID
}
for name := range tt.want {
groupID(name)
}
for _, b := range tt.budgets {
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
GroupID: group.ID,
SpendLimitMicros: tt.spendLimit,
GroupID: groupID(b.group),
SpendLimitMicros: b.limit,
})
require.NoError(t, err)
}
for _, r := range tt.rows {
for _, m := range tt.members {
dbgen.GroupMember(t, db, database.GroupMemberTable{
GroupID: groupID(m.group),
UserID: userID(m.user),
})
}
for _, o := range tt.overrides {
_, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{
UserID: userID(o.user),
GroupID: groupID(o.group),
SpendLimitMicros: o.limit,
})
require.NoError(t, err)
}
for _, s := range tt.spend {
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID,
EffectiveGroupID: group.ID,
Day: r.day,
CostMicros: r.spend,
UserID: userID(s.user),
EffectiveGroupID: groupID(s.group),
Day: s.day,
CostMicros: s.spend,
})
require.NoError(t, err)
}
// When: querying spend for the group since monthStart.
// When: querying every group the case declares.
queried := make([]uuid.UUID, 0, len(groupIDs))
for _, id := range groupIDs {
queried = append(queried, id)
}
got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{
OrganizationID: org.ID,
GroupIds: []uuid.UUID{group.ID},
GroupIds: queried,
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: one row is returned with the group's limit and spend.
require.Len(t, got, 1)
require.Equal(t, group.ID, got[0].GroupID)
require.Equal(t, org.ID, got[0].OrganizationID)
if tt.setBudget {
require.True(t, got[0].SpendLimitMicros.Valid, "expected configured budget")
require.Equal(t, tt.spendLimit, got[0].SpendLimitMicros.Int64, "spend_limit_micros")
} else {
require.False(t, got[0].SpendLimitMicros.Valid, "expected no configured budget")
// Then: each group reports its own limit, spend, and the combined
// limit of the members attributed to it.
require.Len(t, got, len(groupIDs))
byID := make(map[uuid.UUID]database.GetOrganizationGroupsAISpendRow, len(got))
for _, r := range got {
byID[r.GroupID] = r
}
for name, want := range tt.want {
row, ok := byID[groupIDs[name]]
require.True(t, ok, "group %q missing from response", name)
require.Equal(t, org.ID, row.OrganizationID)
require.Equal(t, want.spendLimit, row.SpendLimitMicros, "%s spend_limit_micros", name)
require.Equal(t, want.totalLimit, row.TotalSpendLimitMicros, "%s total_spend_limit_micros", name)
require.Equal(t, want.spend, row.CurrentSpendMicros, "%s current_spend_micros", name)
}
require.Equal(t, tt.wantCurrentSpend, got[0].CurrentSpendMicros)
})
}
@@ -13217,10 +13392,14 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) {
rowA, ok := byID[groupA.ID]
require.True(t, ok, "groupA missing from response")
require.Equal(t, sql.NullInt64{Int64: 1_000_000, Valid: true}, rowA.SpendLimitMicros)
// Neither group has members, so the budgeted group totals zero and the
// unbudgeted one is null.
require.Equal(t, sql.NullInt64{Int64: 0, Valid: true}, rowA.TotalSpendLimitMicros)
require.Equal(t, int64(250), rowA.CurrentSpendMicros)
rowB, ok := byID[groupB.ID]
require.True(t, ok, "groupB missing from response")
require.Equal(t, sql.NullInt64{}, rowB.SpendLimitMicros)
require.Equal(t, sql.NullInt64{}, rowB.TotalSpendLimitMicros)
require.Equal(t, int64(500), rowB.CurrentSpendMicros)
})
@@ -13381,6 +13560,81 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) {
require.Equal(t, int64(25), got[0].CurrentSpendMicros,
"sum must exclude prevMonthLastDay row after normalization")
})
t.Run("EveryoneGroupCountsOrgMembers", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: an org with two members whose implicit Everyone group carries the
// only budget.
org := dbgen.Organization(t, db, database.Organization{})
for range 2 {
user := dbgen.User(t, db, database.User{})
dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: org.ID,
})
}
// The Everyone group has ID equal to the organization ID and must be
// inserted explicitly for the group_ai_budgets FK constraint.
//nolint:gocritic // Requires system context.
_, err := db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), org.ID)
require.NoError(t, err)
_, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
GroupID: org.ID,
SpendLimitMicros: 100,
})
require.NoError(t, err)
// When: querying the Everyone group.
got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{
OrganizationID: org.ID,
GroupIds: []uuid.UUID{org.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: every org member counts toward the total.
require.Len(t, got, 1)
require.Equal(t, sql.NullInt64{Int64: 100, Valid: true}, got[0].SpendLimitMicros, "spend_limit_micros")
require.Equal(t, sql.NullInt64{Int64: 200, Valid: true}, got[0].TotalSpendLimitMicros, "total_spend_limit_micros")
require.Equal(t, int64(0), got[0].CurrentSpendMicros, "current_spend_micros")
})
t.Run("EveryoneGroupWithoutBudgetIsUnlimited", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: an org with two members and no budget anywhere, which is where
// uncapped users are attributed.
org := dbgen.Organization(t, db, database.Organization{})
for range 2 {
user := dbgen.User(t, db, database.User{})
dbgen.OrganizationMember(t, db, database.OrganizationMember{
UserID: user.ID,
OrganizationID: org.ID,
})
}
//nolint:gocritic // Requires system context.
_, err := db.InsertAllUsersGroup(dbauthz.AsSystemRestricted(ctx), org.ID)
require.NoError(t, err)
// When: querying the Everyone group.
got, err := db.GetOrganizationGroupsAISpend(ctx, database.GetOrganizationGroupsAISpendParams{
OrganizationID: org.ID,
GroupIds: []uuid.UUID{org.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: the total is null, so the group reads as unlimited.
require.Len(t, got, 1)
require.Equal(t, sql.NullInt64{}, got[0].SpendLimitMicros, "spend_limit_micros")
require.Equal(t, sql.NullInt64{}, got[0].TotalSpendLimitMicros, "total_spend_limit_micros")
require.Equal(t, int64(0), got[0].CurrentSpendMicros, "current_spend_micros")
})
}
func TestGetGroupMembersAISpend(t *testing.T) {
coderd/database/queries.sql.go
Generated
+99 -20
View File
@@ -2860,41 +2860,119 @@ func (q *sqlQuerier) GetHighestGroupAIBudgetByUser(ctx context.Context, userID u
}
const getOrganizationGroupsAISpend = `-- name: GetOrganizationGroupsAISpend :many
WITH queried_groups AS (
-- The requested groups that belong to the queried organization.
SELECT groups.id, groups.organization_id
FROM groups
WHERE groups.organization_id = $1
AND groups.id = ANY($2::uuid[])
),
candidate_users AS (
-- Members of the queried groups. Uses group_members_expanded so the implicit
-- Everyone group counts.
SELECT DISTINCT member.user_id
FROM group_members_expanded member
WHERE member.group_id IN (SELECT id FROM queried_groups)
),
user_highest_group AS (
-- Per user, the highest-limit group they belong to. Uses
-- group_members_expanded so the implicit Everyone group counts.
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 candidate_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 (
-- Effective budget group per user: an override wins over the highest-limit
-- group they belong to. Users with neither are left out, since the group they
-- fall back to has no budget and reports null.
SELECT
candidate_users.user_id,
COALESCE(override.group_id, user_highest_group.group_id) AS effective_group_id,
override.spend_limit_micros AS override_limit_micros
FROM candidate_users
LEFT JOIN user_ai_budget_overrides override ON override.user_id = candidate_users.user_id
LEFT JOIN user_highest_group ON user_highest_group.user_id = candidate_users.user_id
),
group_limits AS (
-- Per attributed group, how many members take the group's own limit and the
-- combined limit of those carrying an override.
SELECT
effective.effective_group_id AS group_id,
count(*) FILTER (WHERE effective.override_limit_micros IS NULL) AS plain_member_count,
COALESCE(SUM(effective.override_limit_micros), 0)::BIGINT AS override_limit_sum
FROM effective
WHERE effective.effective_group_id IS NOT NULL
GROUP BY effective.effective_group_id
),
group_totals AS (
-- Combined limit per budgeted group, counting members with no override at the
-- group's own limit and adding the overrides on top. Unbudgeted groups are
-- absent here, so the join below leaves their total null.
SELECT
queried_groups.id AS group_id,
(budget.spend_limit_micros * COALESCE(group_limits.plain_member_count, 0)
+ COALESCE(group_limits.override_limit_sum, 0))::BIGINT AS total_spend_limit_micros
FROM queried_groups
JOIN group_ai_budgets budget ON budget.group_id = queried_groups.id
LEFT JOIN group_limits ON group_limits.group_id = queried_groups.id
),
group_spend AS (
-- Spend per queried group over the period.
SELECT
spend.effective_group_id AS group_id,
COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros
FROM ai_user_daily_spend spend
WHERE spend.effective_group_id IN (SELECT id FROM queried_groups)
AND spend.day >= (($3::timestamptz) AT TIME ZONE 'UTC')::date
GROUP BY spend.effective_group_id
)
SELECT
groups.id AS group_id,
groups.organization_id AS organization_id,
queried_groups.id AS group_id,
queried_groups.organization_id AS organization_id,
budget.spend_limit_micros AS spend_limit_micros,
COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros
FROM groups
LEFT JOIN group_ai_budgets budget ON budget.group_id = groups.id
LEFT JOIN ai_user_daily_spend spend
ON spend.effective_group_id = groups.id
AND spend.day >= (($1::timestamptz) AT TIME ZONE 'UTC')::date
WHERE groups.organization_id = $2
AND groups.id = ANY($3::uuid[])
GROUP BY groups.id, budget.spend_limit_micros
ORDER BY groups.id
group_totals.total_spend_limit_micros AS total_spend_limit_micros,
COALESCE(group_spend.current_spend_micros, 0)::BIGINT AS current_spend_micros
FROM queried_groups
LEFT JOIN group_ai_budgets budget ON budget.group_id = queried_groups.id
LEFT JOIN group_totals ON group_totals.group_id = queried_groups.id
LEFT JOIN group_spend ON group_spend.group_id = queried_groups.id
ORDER BY queried_groups.id
`
type GetOrganizationGroupsAISpendParams struct {
PeriodStart time.Time `db:"period_start" json:"period_start"`
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
GroupIds []uuid.UUID `db:"group_ids" json:"group_ids"`
PeriodStart time.Time `db:"period_start" json:"period_start"`
}
type GetOrganizationGroupsAISpendRow struct {
GroupID uuid.UUID `db:"group_id" json:"group_id"`
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
CurrentSpendMicros int64 `db:"current_spend_micros" json:"current_spend_micros"`
GroupID uuid.UUID `db:"group_id" json:"group_id"`
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
TotalSpendLimitMicros sql.NullInt64 `db:"total_spend_limit_micros" json:"total_spend_limit_micros"`
CurrentSpendMicros int64 `db:"current_spend_micros" json:"current_spend_micros"`
}
// Returns AI spend limits and aggregate spend for groups in @group_ids that
// belong to @organization_id, on or after period_start until NOW. The spend
// limit is null when the group has no configured budget.
// belong to @organization_id, on or after period_start until NOW.
// spend_limit_micros is the per-member limit, null when the group has no budget.
// total_spend_limit_micros is the combined budget of the members attributed to
// the group, with each member's override replacing their share. It is null when
// the group has no budget.
// The period_start parameter is normalized to its UTC calendar day.
// TODO(AIGOV-527): unify effective group resolution in a single place.
func (q *sqlQuerier) GetOrganizationGroupsAISpend(ctx context.Context, arg GetOrganizationGroupsAISpendParams) ([]GetOrganizationGroupsAISpendRow, error) {
rows, err := q.db.QueryContext(ctx, getOrganizationGroupsAISpend, arg.PeriodStart, arg.OrganizationID, pq.Array(arg.GroupIds))
rows, err := q.db.QueryContext(ctx, getOrganizationGroupsAISpend, arg.OrganizationID, pq.Array(arg.GroupIds), arg.PeriodStart)
if err != nil {
return nil, err
}
@@ -2906,6 +2984,7 @@ func (q *sqlQuerier) GetOrganizationGroupsAISpend(ctx context.Context, arg GetOr
&i.GroupID,
&i.OrganizationID,
&i.SpendLimitMicros,
&i.TotalSpendLimitMicros,
&i.CurrentSpendMicros,
); err != nil {
return nil, err
coderd/database/queries/aicostcontrol.sql
+91 -14
View File
@@ -122,23 +122,100 @@ WHERE user_id = @user_id
-- name: GetOrganizationGroupsAISpend :many
-- Returns AI spend limits and aggregate spend for groups in @group_ids that
-- belong to @organization_id, on or after period_start until NOW. The spend
-- limit is null when the group has no configured budget.
-- belong to @organization_id, on or after period_start until NOW.
-- spend_limit_micros is the per-member limit, null when the group has no budget.
-- total_spend_limit_micros is the combined budget of the members attributed to
-- the group, with each member's override replacing their share. It is null when
-- the group has no budget.
-- The period_start parameter is normalized to its UTC calendar day.
-- TODO(AIGOV-527): unify effective group resolution in a single place.
WITH queried_groups AS (
-- The requested groups that belong to the queried organization.
SELECT groups.id, groups.organization_id
FROM groups
WHERE groups.organization_id = @organization_id
AND groups.id = ANY(@group_ids::uuid[])
),
candidate_users AS (
-- Members of the queried groups. Uses group_members_expanded so the implicit
-- Everyone group counts.
SELECT DISTINCT member.user_id
FROM group_members_expanded member
WHERE member.group_id IN (SELECT id FROM queried_groups)
),
user_highest_group AS (
-- Per user, the highest-limit group they belong to. Uses
-- group_members_expanded so the implicit Everyone group counts.
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 candidate_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 (
-- Effective budget group per user: an override wins over the highest-limit
-- group they belong to. Users with neither are left out, since the group they
-- fall back to has no budget and reports null.
SELECT
candidate_users.user_id,
COALESCE(override.group_id, user_highest_group.group_id) AS effective_group_id,
override.spend_limit_micros AS override_limit_micros
FROM candidate_users
LEFT JOIN user_ai_budget_overrides override ON override.user_id = candidate_users.user_id
LEFT JOIN user_highest_group ON user_highest_group.user_id = candidate_users.user_id
),
group_limits AS (
-- Per attributed group, how many members take the group's own limit and the
-- combined limit of those carrying an override.
SELECT
effective.effective_group_id AS group_id,
count(*) FILTER (WHERE effective.override_limit_micros IS NULL) AS plain_member_count,
COALESCE(SUM(effective.override_limit_micros), 0)::BIGINT AS override_limit_sum
FROM effective
WHERE effective.effective_group_id IS NOT NULL
GROUP BY effective.effective_group_id
),
group_totals AS (
-- Combined limit per budgeted group, counting members with no override at the
-- group's own limit and adding the overrides on top. Unbudgeted groups are
-- absent here, so the join below leaves their total null.
SELECT
queried_groups.id AS group_id,
(budget.spend_limit_micros * COALESCE(group_limits.plain_member_count, 0)
+ COALESCE(group_limits.override_limit_sum, 0))::BIGINT AS total_spend_limit_micros
FROM queried_groups
JOIN group_ai_budgets budget ON budget.group_id = queried_groups.id
LEFT JOIN group_limits ON group_limits.group_id = queried_groups.id
),
group_spend AS (
-- Spend per queried group over the period.
SELECT
spend.effective_group_id AS group_id,
COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros
FROM ai_user_daily_spend spend
WHERE spend.effective_group_id IN (SELECT id FROM queried_groups)
AND spend.day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date
GROUP BY spend.effective_group_id
)
SELECT
groups.id AS group_id,
groups.organization_id AS organization_id,
queried_groups.id AS group_id,
queried_groups.organization_id AS organization_id,
budget.spend_limit_micros AS spend_limit_micros,
COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros
FROM groups
LEFT JOIN group_ai_budgets budget ON budget.group_id = groups.id
LEFT JOIN ai_user_daily_spend spend
ON spend.effective_group_id = groups.id
AND spend.day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date
WHERE groups.organization_id = @organization_id
AND groups.id = ANY(@group_ids::uuid[])
GROUP BY groups.id, budget.spend_limit_micros
ORDER BY groups.id;
group_totals.total_spend_limit_micros AS total_spend_limit_micros,
COALESCE(group_spend.current_spend_micros, 0)::BIGINT AS current_spend_micros
FROM queried_groups
LEFT JOIN group_ai_budgets budget ON budget.group_id = queried_groups.id
LEFT JOIN group_totals ON group_totals.group_id = queried_groups.id
LEFT JOIN group_spend ON group_spend.group_id = queried_groups.id
ORDER BY queried_groups.id;
-- name: GetGroupMembersAISpend :many
-- Returns each user's AI spend attributed to the queried group, on or after