mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
Generated
+15
@@ -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()
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Generated
+12
@@ -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
|
||||
|
||||
@@ -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() {
|
||||
|
||||
Generated
+131
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user