feat: add GET /groups/{group}/members/ai/spend (#27130)

## Description

Adds `GET /api/v2/groups/{group}/members/ai/spend?user_ids=...` (also available org-scoped at `/api/v2/organizations/{org}/groups/{groupName}/members/ai/spend`) to return per-member AI spend attributed to a group, along with each member's effective budget group and the applied spend limit when the queried group is their effective budget source.

In the UI, this endpoint is used alongside the existing `/api/v2/groups/{group}/members` endpoint. AI spend data is kept separate from that endpoint so that:

- Different concepts stay on different endpoints: identity (group members) vs. cost control (spend). Cost control is an additional feature layered on top of groups/orgs.
- Callers that don't need spend information don't pay for its computation.

UI flow:

1. Request `/api/v2/groups/{group}/members` → returns the group's members.
2. Request `/api/v2/groups/{group}/members/ai/spend?user_ids=...` with the IDs from step 1.

**Note:** Only current members of the queried group are returned. `spend_limit_micros` and `limit_source` are populated only when the queried group is the member's effective budget source (its own limit or a user override). `effective_group_id` is null when the member's budget resolves to a group in another organization, since an organization is treated as a tenant boundary.

<img width="2880" height="1904" alt="image" src="https://github.com/user-attachments/assets/33ed395d-d1a3-4b46-bb04-c8d3f41c8886" />

## Changes

- Add `codersdk.GroupMembersAISpend` and `GroupMemberAISpend` types, reusing the shared `AISpendPeriodWindow`.
- Add `GetGroupMembersAISpend` SQL query with a dbauthz per-row filter that mirrors `GET /api/v2/groups/{group}/members`.
- Add handler and routes under `/groups/{group}/members/ai/spend` (and the org-scoped alias) with a required `user_ids` query param (cap 100). Callers with more than 100 members are expected to batch across multiple requests.
- Add codersdk client method.
- Tests: dbauthz, raw SQL, endpoint, and role-access.

Closes https://linear.app/codercom/issue/AIGOV-471/backend-group-members-endpoint-with-members-spend

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
This commit is contained in:
Susana Ferreira
2026-07-20 13:04:37 +01:00
committed by GitHub
parent 2adc8f5272
commit a9fdf87a2f
20 changed files with 2009 additions and 0 deletions
+17
View File
@@ -1469,6 +1469,23 @@ func OrganizationGroupAISpend(row database.GetOrganizationGroupsAISpendRow) code
return group
}
func GroupMemberAISpend(row database.GetGroupMembersAISpendRow) codersdk.GroupMemberAISpend {
member := codersdk.GroupMemberAISpend{
UserID: row.UserID,
GroupSpendMicros: row.GroupSpendMicros,
}
if row.EffectiveGroupID.Valid {
member.EffectiveGroupID = &row.EffectiveGroupID.UUID
}
if row.SpendLimitMicros.Valid {
member.GroupBudget = &codersdk.AIGroupBudget{
SpendLimitMicros: row.SpendLimitMicros.Int64,
LimitSource: codersdk.AIBudgetLimitSource(row.LimitSource.String),
}
}
return member
}
func InvalidatedPresets(invalidatedPresets []database.UpdatePresetsLastInvalidatedAtRow) []codersdk.InvalidatedPreset {
var presets []codersdk.InvalidatedPreset
for _, p := range invalidatedPresets {
+4
View File
@@ -3883,6 +3883,10 @@ func (q *querier) GetGroupMembers(ctx context.Context, includeSystem bool) ([]da
return q.db.GetGroupMembers(ctx, includeSystem)
}
func (q *querier) GetGroupMembersAISpend(ctx context.Context, arg database.GetGroupMembersAISpendParams) ([]database.GetGroupMembersAISpendRow, error) {
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetGroupMembersAISpend)(ctx, arg)
}
func (q *querier) GetGroupMembersByGroupID(ctx context.Context, arg database.GetGroupMembersByGroupIDParams) ([]database.GroupMember, error) {
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetGroupMembersByGroupID)(ctx, arg)
}
+16
View File
@@ -6954,6 +6954,22 @@ func (s *MethodTestSuite) TestAIBridge() {
Returns([]database.GetOrganizationGroupsAISpendRow{row1, row2})
}))
s.Run("GetGroupMembersAISpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
group := testutil.Fake(s.T(), faker, database.Group{})
row1 := testutil.Fake(s.T(), faker, database.GetGroupMembersAISpendRow{OrganizationID: group.OrganizationID})
row2 := testutil.Fake(s.T(), faker, database.GetGroupMembersAISpendRow{OrganizationID: group.OrganizationID})
arg := database.GetGroupMembersAISpendParams{
GroupID: group.ID,
UserIds: []uuid.UUID{row1.UserID, row2.UserID},
PeriodStart: time.Now().UTC().Truncate(24 * time.Hour),
}
dbm.EXPECT().GetGroupMembersAISpend(gomock.Any(), arg).
Return([]database.GetGroupMembersAISpendRow{row1, row2}, nil).AnyTimes()
check.Args(arg).
Asserts(row1, policy.ActionRead, row2, policy.ActionRead).
Returns([]database.GetGroupMembersAISpendRow{row1, row2})
}))
s.Run("GetGroupAIBudget", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
g := testutil.Fake(s.T(), faker, database.Group{})
b := testutil.Fake(s.T(), faker, database.GroupAIBudget{GroupID: g.ID})
+8
View File
@@ -2161,6 +2161,14 @@ func (m queryMetricsStore) GetGroupMembers(ctx context.Context, includeSystem bo
return r0, r1
}
func (m queryMetricsStore) GetGroupMembersAISpend(ctx context.Context, arg database.GetGroupMembersAISpendParams) ([]database.GetGroupMembersAISpendRow, error) {
start := time.Now()
r0, r1 := m.s.GetGroupMembersAISpend(ctx, arg)
m.queryLatencies.WithLabelValues("GetGroupMembersAISpend").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetGroupMembersAISpend").Inc()
return r0, r1
}
func (m queryMetricsStore) GetGroupMembersByGroupID(ctx context.Context, arg database.GetGroupMembersByGroupIDParams) ([]database.GroupMember, error) {
start := time.Now()
r0, r1 := m.s.GetGroupMembersByGroupID(ctx, arg)
+15
View File
@@ -4003,6 +4003,21 @@ func (mr *MockStoreMockRecorder) GetGroupMembers(ctx, includeSystem any) *gomock
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetGroupMembers", reflect.TypeOf((*MockStore)(nil).GetGroupMembers), ctx, includeSystem)
}
// GetGroupMembersAISpend mocks base method.
func (m *MockStore) GetGroupMembersAISpend(ctx context.Context, arg database.GetGroupMembersAISpendParams) ([]database.GetGroupMembersAISpendRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetGroupMembersAISpend", ctx, arg)
ret0, _ := ret[0].([]database.GetGroupMembersAISpendRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetGroupMembersAISpend indicates an expected call of GetGroupMembersAISpend.
func (mr *MockStoreMockRecorder) GetGroupMembersAISpend(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetGroupMembersAISpend", reflect.TypeOf((*MockStore)(nil).GetGroupMembersAISpend), ctx, arg)
}
// GetGroupMembersByGroupID mocks base method.
func (m *MockStore) GetGroupMembersByGroupID(ctx context.Context, arg database.GetGroupMembersByGroupIDParams) ([]database.GroupMember, error) {
m.ctrl.T.Helper()
+4
View File
@@ -470,6 +470,10 @@ func (gm GetGroupMembersByGroupIDPaginatedRow) RBACObject() rbac.Object {
return rbac.ResourceGroupMember.WithID(gm.UserID).InOrg(gm.OrganizationID).WithOwner(gm.UserID.String())
}
func (r GetGroupMembersAISpendRow) RBACObject() rbac.Object {
return rbac.ResourceGroupMember.WithID(r.UserID).InOrg(r.OrganizationID).WithOwner(r.UserID.String())
}
// PrebuiltWorkspaceResource defines the interface for types that can be identified as prebuilt workspaces
// and converted to their corresponding prebuilt workspace RBAC object.
type PrebuiltWorkspaceResource interface {
+12
View File
@@ -577,6 +577,18 @@ type sqlcQuerier interface {
GetGroupByID(ctx context.Context, id uuid.UUID) (Group, error)
GetGroupByOrgAndName(ctx context.Context, arg GetGroupByOrgAndNameParams) (Group, error)
GetGroupMembers(ctx context.Context, includeSystem bool) ([]GroupMember, error)
// Returns each user's AI spend attributed to the queried group, on or after
// period_start until NOW. Only current members of the queried group are
// returned. spend_limit_micros and limit_source are populated only when the
// queried group is the user's effective budget source. The effective_group_id
// is null when the user has no configured budget or when the effective group
// belongs to a different organization than the queried group.
// The period_start parameter is normalized to its UTC calendar day.
// TODO(AIGOV-527): unify effective group resolution in a single place.
// Spend is aggregated for the queried group, not the user's effective group.
// A LEFT JOIN leaves spend_limit_micros and limit_source null for users
// whose effective budget source is not the queried group.
GetGroupMembersAISpend(ctx context.Context, arg GetGroupMembersAISpendParams) ([]GetGroupMembersAISpendRow, error)
GetGroupMembersByGroupID(ctx context.Context, arg GetGroupMembersByGroupIDParams) ([]GroupMember, error)
GetGroupMembersByGroupIDPaginated(ctx context.Context, arg GetGroupMembersByGroupIDPaginatedParams) ([]GetGroupMembersByGroupIDPaginatedRow, error)
// Returns the total count of members in a group. Shows the total
+586
View File
@@ -13106,6 +13106,592 @@ func TestGetOrganizationGroupsAISpend(t *testing.T) {
})
}
func TestGetGroupMembersAISpend(t *testing.T) {
t.Parallel()
// Use fixed dates to keep the test deterministic.
monthStart := time.Date(2024, 6, 1, 0, 0, 0, 0, time.UTC)
now := monthStart.AddDate(0, 0, 14) // 2024-06-15
prevMonthLastDay := monthStart.AddDate(0, 0, -1) // 2024-05-31
tests := []struct {
name string
groupLimit int64
overrideLimit int64
spend int64
wantEffectiveGroup bool
wantLimit sql.NullInt64
wantSource sql.NullString
wantSpend int64
}{
{
name: "NoBudgetNoSpend",
wantEffectiveGroup: false,
wantLimit: sql.NullInt64{},
wantSource: sql.NullString{},
wantSpend: 0,
},
{
name: "GroupBudget",
groupLimit: 1_000_000,
wantEffectiveGroup: true,
wantLimit: sql.NullInt64{Int64: 1_000_000, Valid: true},
wantSource: sql.NullString{String: "group", Valid: true},
wantSpend: 0,
},
{
name: "OverrideBudget",
overrideLimit: 500_000,
wantEffectiveGroup: true,
wantLimit: sql.NullInt64{Int64: 500_000, Valid: true},
wantSource: sql.NullString{String: "user_override", Valid: true},
wantSpend: 0,
},
{
name: "NoBudgetWithSpend",
spend: 250,
wantEffectiveGroup: false,
wantLimit: sql.NullInt64{},
wantSource: sql.NullString{},
wantSpend: 250,
},
{
name: "BudgetWithSpend",
groupLimit: 1_000_000,
spend: 250,
wantEffectiveGroup: true,
wantLimit: sql.NullInt64{Int64: 1_000_000, Valid: true},
wantSource: sql.NullString{String: "group", Valid: true},
wantSpend: 250,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: a member of the queried group, optionally with a group
// budget, a user override, and seeded spend.
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{UserID: user.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
if tt.groupLimit > 0 {
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
GroupID: group.ID,
SpendLimitMicros: tt.groupLimit,
})
require.NoError(t, err)
}
if tt.overrideLimit > 0 {
_, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{
UserID: user.ID,
GroupID: group.ID,
SpendLimitMicros: tt.overrideLimit,
})
require.NoError(t, err)
}
if tt.spend > 0 {
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: tt.spend,
})
require.NoError(t, err)
}
// When: querying spend for the user.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: group.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: one row with the expected effective group, limit, source, and spend.
require.Len(t, got, 1)
require.Equal(t, user.ID, got[0].UserID)
require.Equal(t, org.ID, got[0].OrganizationID)
if tt.wantEffectiveGroup {
require.Equal(t, uuid.NullUUID{UUID: group.ID, Valid: true}, got[0].EffectiveGroupID)
} else {
require.False(t, got[0].EffectiveGroupID.Valid, "expected no effective group")
}
require.Equal(t, tt.wantLimit, got[0].SpendLimitMicros)
require.Equal(t, tt.wantSource, got[0].LimitSource)
require.Equal(t, tt.wantSpend, got[0].GroupSpendMicros)
})
}
t.Run("MultipleMembers", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: two group members with different spend attributed to the group.
userA := dbgen.User(t, db, database.User{})
userB := 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{UserID: userA.ID, OrganizationID: org.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: userB.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: userA.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: userB.ID})
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: userA.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: 100,
})
require.NoError(t, err)
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: userB.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: 250,
})
require.NoError(t, err)
// When: querying spend for both users.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: group.ID,
UserIds: []uuid.UUID{userA.ID, userB.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: both users are returned with their own aggregate spend.
require.Len(t, got, 2)
byID := make(map[uuid.UUID]database.GetGroupMembersAISpendRow, len(got))
for _, row := range got {
byID[row.UserID] = row
}
require.Equal(t, int64(100), byID[userA.ID].GroupSpendMicros)
require.Equal(t, int64(250), byID[userB.ID].GroupSpendMicros)
for _, row := range got {
require.False(t, row.EffectiveGroupID.Valid)
require.False(t, row.SpendLimitMicros.Valid)
require.False(t, row.LimitSource.Valid)
}
})
t.Run("AggregatesSpendAcrossDays", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: a group member with spend on multiple days in the period.
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{UserID: user.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
days := []time.Time{monthStart, monthStart.AddDate(0, 0, 7), monthStart.AddDate(0, 0, 14)}
for i, day := range days {
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID, EffectiveGroupID: group.ID, Day: day, CostMicros: int64((i + 1) * 100),
})
require.NoError(t, err)
}
// When: querying the group's member spend.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: group.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: per-user spend is summed across all days in the period.
require.Len(t, got, 1)
require.False(t, got[0].EffectiveGroupID.Valid)
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(600), got[0].GroupSpendMicros)
})
t.Run("OverrideWins", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: a member of the queried group who is also in two other
// budgeted groups, one being the natural highest-limit and the other
// set via a user override.
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
overrideTarget := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
highestLimit := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: overrideTarget.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: highestLimit.ID, UserID: user.ID})
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
GroupID: overrideTarget.ID,
SpendLimitMicros: 1_000_000,
})
require.NoError(t, err)
_, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
GroupID: highestLimit.ID,
SpendLimitMicros: 5_000_000,
})
require.NoError(t, err)
_, err = db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{
UserID: user.ID,
GroupID: overrideTarget.ID,
SpendLimitMicros: 500_000,
})
require.NoError(t, err)
// When: querying spend for the user.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: queried.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: the override target wins over the highest-limit group.
require.Len(t, got, 1)
require.Equal(t, uuid.NullUUID{UUID: overrideTarget.ID, Valid: true}, got[0].EffectiveGroupID)
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(0), got[0].GroupSpendMicros)
})
t.Run("EqualBudgetTieBreak", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: a member of the queried group who is in two same-org groups
// with identical spend limits.
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
groupA := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "aaa-tie-group"})
groupB := dbgen.Group(t, db, database.Group{OrganizationID: org.ID, Name: "bbb-tie-group"})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupA.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: groupB.ID, UserID: user.ID})
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
GroupID: groupA.ID,
SpendLimitMicros: 1_000_000,
})
require.NoError(t, err)
_, err = db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
GroupID: groupB.ID,
SpendLimitMicros: 1_000_000,
})
require.NoError(t, err)
// When: querying spend for the user.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: queried.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: the tie is broken by group name ascending, so groupA wins.
require.Len(t, got, 1)
require.Equal(t, uuid.NullUUID{UUID: groupA.ID, Valid: true}, got[0].EffectiveGroupID)
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(0), got[0].GroupSpendMicros)
})
t.Run("EveryoneGroupCounts", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: a member of the queried group whose only budgeted group is
// the org's implicit Everyone group.
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID})
// The Everyone group has ID equal to the organization ID and must be
// inserted explicitly for this test's FK constraint on group_ai_budgets.
//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: 1_000_000,
})
require.NoError(t, err)
// When: querying spend for the user.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: queried.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: effective_group_id resolves to the Everyone group.
require.Len(t, got, 1)
require.Equal(t, uuid.NullUUID{UUID: org.ID, Valid: true}, got[0].EffectiveGroupID)
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(0), got[0].GroupSpendMicros)
})
t.Run("SpendWithDifferentEffectiveGroup", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: a group member with spend attributed to the queried group,
// whose current effective group is a different same-org group.
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
other := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: queried.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: other.ID, UserID: user.ID})
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
GroupID: other.ID,
SpendLimitMicros: 1_000_000,
})
require.NoError(t, err)
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID, EffectiveGroupID: queried.ID, Day: now, CostMicros: 250,
})
require.NoError(t, err)
// When: querying the queried group's spend for the user.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: queried.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: queried-group spend is returned, effective_group_id is the
// other group, and the limit and source are null because the queried
// group is not the effective source.
require.Len(t, got, 1)
require.Equal(t, uuid.NullUUID{UUID: other.ID, Valid: true}, got[0].EffectiveGroupID)
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(250), got[0].GroupSpendMicros)
})
t.Run("ExcludesOtherGroupSpend", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: a group member with spend attributed to a different group in the same org.
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
otherGroup := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID, EffectiveGroupID: otherGroup.ID, Day: now, CostMicros: 500,
})
require.NoError(t, err)
// When: querying spend for the queried group.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: group.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: spend attributed to the other group is not counted.
require.Len(t, got, 1)
require.False(t, got[0].EffectiveGroupID.Valid)
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(0), got[0].GroupSpendMicros)
})
t.Run("ExcludesNonMembers", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: a member of the queried group and a non-member in the same org.
member := dbgen.User(t, db, database.User{})
nonMember := 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{UserID: member.ID, OrganizationID: org.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: nonMember.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: member.ID})
// When: querying with both user IDs.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: group.ID,
UserIds: []uuid.UUID{member.ID, nonMember.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: only the group member is returned.
require.Len(t, got, 1)
require.Equal(t, member.ID, got[0].UserID)
require.False(t, got[0].EffectiveGroupID.Valid)
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(0), got[0].GroupSpendMicros)
})
t.Run("HidesFormerMember", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: a user who has historical spend attributed to the queried
// group but is not currently a member of it.
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
queried := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID, EffectiveGroupID: queried.ID, Day: now, CostMicros: 500,
})
require.NoError(t, err)
// When: querying the queried group's spend for the user.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: queried.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: the user is filtered out and their historical spend is not returned.
require.Empty(t, got)
})
t.Run("CrossOrgEffectiveGroupMasked", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: a member of the queried group whose highest-limit budget
// group is in a different org.
user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
otherOrg := dbgen.Organization(t, db, database.Organization{})
group := dbgen.Group(t, db, database.Group{OrganizationID: org.ID})
otherOrgGroup := dbgen.Group(t, db, database.Group{OrganizationID: otherOrg.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: otherOrg.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: otherOrgGroup.ID, UserID: user.ID})
_, err := db.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
GroupID: otherOrgGroup.ID,
SpendLimitMicros: 9_999_999,
})
require.NoError(t, err)
// Seed spend attributed to the queried group so we can assert it is
// still returned even when the effective group is masked.
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID, EffectiveGroupID: group.ID, Day: now, CostMicros: 250,
})
require.NoError(t, err)
// When: querying spend for the user in the queried group's org.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: group.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: effective_group_id is masked to NULL, the highest-limit group is cross-org.
// The queried-group spend is still returned.
require.Len(t, got, 1)
require.False(t, got[0].EffectiveGroupID.Valid, "cross-org effective group must be masked")
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(250), got[0].GroupSpendMicros)
})
t.Run("ExcludesSpendBeforePeriodStart", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: a group member with spend both in the prior period and in the current period.
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{UserID: user.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID, EffectiveGroupID: group.ID, Day: prevMonthLastDay, CostMicros: 999,
})
require.NoError(t, err)
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID, EffectiveGroupID: group.ID, Day: monthStart, CostMicros: 25,
})
require.NoError(t, err)
// When: querying since monthStart.
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: group.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: monthStart,
})
require.NoError(t, err)
// Then: only current-period spend is aggregated.
require.Len(t, got, 1)
require.False(t, got[0].EffectiveGroupID.Valid)
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(25), got[0].GroupSpendMicros)
})
t.Run("NormalizesNonUTCPeriodStart", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitShort)
// Given: a group member with spend on the prior UTC day and on the first day of the current UTC month.
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{UserID: user.ID, OrganizationID: org.ID})
dbgen.GroupMember(t, db, database.GroupMemberTable{GroupID: group.ID, UserID: user.ID})
_, err := db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID, EffectiveGroupID: group.ID, Day: prevMonthLastDay, CostMicros: 999,
})
require.NoError(t, err)
_, err = db.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: user.ID, EffectiveGroupID: group.ID, Day: monthStart, CostMicros: 25,
})
require.NoError(t, err)
// When: querying with a non-UTC period_start that normalizes to June 1 UTC.
// 2024-05-31 23:00 in UTC-5 is 2024-06-01 04:00 UTC.
localLate := time.Date(2024, 5, 31, 23, 0, 0, 0, time.FixedZone("UTC-5", -5*3600))
got, err := db.GetGroupMembersAISpend(ctx, database.GetGroupMembersAISpendParams{
GroupID: group.ID,
UserIds: []uuid.UUID{user.ID},
PeriodStart: localLate,
})
require.NoError(t, err)
// Then: the prior UTC day's spend is excluded from the aggregate.
require.Len(t, got, 1)
require.False(t, got[0].EffectiveGroupID.Valid)
require.False(t, got[0].SpendLimitMicros.Valid)
require.False(t, got[0].LimitSource.Valid)
require.Equal(t, int64(25), got[0].GroupSpendMicros,
"sum must exclude prevMonthLastDay row after normalization")
})
}
func TestChatPinOrderQueries(t *testing.T) {
t.Parallel()
if testing.Short() {
+131
View File
@@ -2533,6 +2533,137 @@ func (q *sqlQuerier) GetGroupAIBudget(ctx context.Context, groupID uuid.UUID) (G
return i, err
}
const getGroupMembersAISpend = `-- name: GetGroupMembersAISpend :many
WITH queried_group AS (
-- The queried group's org, used to detect cross-org effective groups.
SELECT organization_id
FROM groups
WHERE id = $1
),
filtered_users AS (
-- Users from @user_ids that are members of the queried group. Uses
-- group_members_expanded so the implicit Everyone group counts.
SELECT DISTINCT user_id
FROM group_members_expanded
WHERE group_id = $1
AND user_id = ANY($3::uuid[])
),
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
WHERE member.user_id IN (SELECT user_id FROM filtered_users)
ORDER BY member.user_id, budget.spend_limit_micros DESC, member.group_name ASC, budget.group_id ASC
),
effective AS (
-- Effective budget per user: a per-user override wins over the
-- highest-limit group.
SELECT
filtered_users.user_id,
COALESCE(override.group_id, user_highest_group.group_id) AS raw_effective_group_id,
COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros,
(CASE
WHEN override.group_id IS NOT NULL THEN 'user_override'
WHEN user_highest_group.group_id IS NOT NULL THEN 'group'
END)::text AS limit_source
FROM filtered_users
LEFT JOIN user_ai_budget_overrides override ON override.user_id = filtered_users.user_id
LEFT JOIN user_highest_group ON user_highest_group.user_id = filtered_users.user_id
),
applied_budget AS (
-- The limit and source only for users whose effective budget source is the
-- queried group.
SELECT user_id, spend_limit_micros, limit_source
FROM effective
WHERE raw_effective_group_id = $1
)
SELECT
effective.user_id,
queried_group.organization_id,
effective_group.id AS effective_group_id,
applied_budget.spend_limit_micros,
applied_budget.limit_source,
COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS group_spend_micros
FROM effective
CROSS JOIN queried_group
LEFT JOIN groups effective_group
ON effective_group.id = effective.raw_effective_group_id
AND effective_group.organization_id = queried_group.organization_id
LEFT JOIN applied_budget ON applied_budget.user_id = effective.user_id
LEFT JOIN ai_user_daily_spend spend
ON spend.user_id = effective.user_id
AND spend.effective_group_id = $1
AND spend.day >= (($2::timestamptz) AT TIME ZONE 'UTC')::date
GROUP BY
effective.user_id,
queried_group.organization_id,
effective_group.id,
applied_budget.spend_limit_micros,
applied_budget.limit_source
ORDER BY effective.user_id
`
type GetGroupMembersAISpendParams struct {
GroupID uuid.UUID `db:"group_id" json:"group_id"`
PeriodStart time.Time `db:"period_start" json:"period_start"`
UserIds []uuid.UUID `db:"user_ids" json:"user_ids"`
}
type GetGroupMembersAISpendRow struct {
UserID uuid.UUID `db:"user_id" json:"user_id"`
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
EffectiveGroupID uuid.NullUUID `db:"effective_group_id" json:"effective_group_id"`
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
LimitSource sql.NullString `db:"limit_source" json:"limit_source"`
GroupSpendMicros int64 `db:"group_spend_micros" json:"group_spend_micros"`
}
// Returns each user's AI spend attributed to the queried group, on or after
// period_start until NOW. Only current members of the queried group are
// returned. spend_limit_micros and limit_source are populated only when the
// queried group is the user's effective budget source. The effective_group_id
// is null when the user has no configured budget or when the effective group
// belongs to a different organization than the queried group.
// The period_start parameter is normalized to its UTC calendar day.
// TODO(AIGOV-527): unify effective group resolution in a single place.
// Spend is aggregated for the queried group, not the user's effective group.
// A LEFT JOIN leaves spend_limit_micros and limit_source null for users
// whose effective budget source is not the queried group.
func (q *sqlQuerier) GetGroupMembersAISpend(ctx context.Context, arg GetGroupMembersAISpendParams) ([]GetGroupMembersAISpendRow, error) {
rows, err := q.db.QueryContext(ctx, getGroupMembersAISpend, arg.GroupID, arg.PeriodStart, pq.Array(arg.UserIds))
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetGroupMembersAISpendRow
for rows.Next() {
var i GetGroupMembersAISpendRow
if err := rows.Scan(
&i.UserID,
&i.OrganizationID,
&i.EffectiveGroupID,
&i.SpendLimitMicros,
&i.LimitSource,
&i.GroupSpendMicros,
); 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 getHighestGroupAIBudgetByUser = `-- name: GetHighestGroupAIBudgetByUser :one
SELECT
gaib.group_id,
+85
View File
@@ -121,3 +121,88 @@ 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;
-- name: GetGroupMembersAISpend :many
-- Returns each user's AI spend attributed to the queried group, on or after
-- period_start until NOW. Only current members of the queried group are
-- returned. spend_limit_micros and limit_source are populated only when the
-- queried group is the user's effective budget source. The effective_group_id
-- is null when the user has no configured budget or when the effective group
-- belongs to a different organization than the queried group.
-- The period_start parameter is normalized to its UTC calendar day.
-- TODO(AIGOV-527): unify effective group resolution in a single place.
WITH queried_group AS (
-- The queried group's org, used to detect cross-org effective groups.
SELECT organization_id
FROM groups
WHERE id = @group_id
),
filtered_users AS (
-- Users from @user_ids that are members of the queried group. Uses
-- group_members_expanded so the implicit Everyone group counts.
SELECT DISTINCT user_id
FROM group_members_expanded
WHERE group_id = @group_id
AND user_id = ANY(@user_ids::uuid[])
),
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
WHERE member.user_id IN (SELECT user_id FROM filtered_users)
ORDER BY member.user_id, budget.spend_limit_micros DESC, member.group_name ASC, budget.group_id ASC
),
effective AS (
-- Effective budget per user: a per-user override wins over the
-- highest-limit group.
SELECT
filtered_users.user_id,
COALESCE(override.group_id, user_highest_group.group_id) AS raw_effective_group_id,
COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros,
(CASE
WHEN override.group_id IS NOT NULL THEN 'user_override'
WHEN user_highest_group.group_id IS NOT NULL THEN 'group'
END)::text AS limit_source
FROM filtered_users
LEFT JOIN user_ai_budget_overrides override ON override.user_id = filtered_users.user_id
LEFT JOIN user_highest_group ON user_highest_group.user_id = filtered_users.user_id
),
applied_budget AS (
-- The limit and source only for users whose effective budget source is the
-- queried group.
SELECT user_id, spend_limit_micros, limit_source
FROM effective
WHERE raw_effective_group_id = @group_id
)
-- Spend is aggregated for the queried group, not the user's effective group.
SELECT
effective.user_id,
queried_group.organization_id,
effective_group.id AS effective_group_id,
applied_budget.spend_limit_micros,
applied_budget.limit_source,
COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS group_spend_micros
FROM effective
CROSS JOIN queried_group
LEFT JOIN groups effective_group
ON effective_group.id = effective.raw_effective_group_id
AND effective_group.organization_id = queried_group.organization_id
-- A LEFT JOIN leaves spend_limit_micros and limit_source null for users
-- whose effective budget source is not the queried group.
LEFT JOIN applied_budget ON applied_budget.user_id = effective.user_id
LEFT JOIN ai_user_daily_spend spend
ON spend.user_id = effective.user_id
AND spend.effective_group_id = @group_id
AND spend.day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date
GROUP BY
effective.user_id,
queried_group.organization_id,
effective_group.id,
applied_budget.spend_limit_micros,
applied_budget.limit_source
ORDER BY effective.user_id;