From 6b3a73161a74bdb194ec1601dfffba0774521c36 Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Thu, 12 Jun 2025 08:07:05 +0000 Subject: [PATCH 1/9] Allow users to view and modify group membership for the prebuilds system user so that they can configure quotas --- enterprise/coderd/groups.go | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/enterprise/coderd/groups.go b/enterprise/coderd/groups.go index cfe5d081271e3..8eebc384b5b8a 100644 --- a/enterprise/coderd/groups.go +++ b/enterprise/coderd/groups.go @@ -156,7 +156,7 @@ func (api *API) patchGroup(rw http.ResponseWriter, r *http.Request) { currentMembers, err := api.Database.GetGroupMembersByGroupID(ctx, database.GetGroupMembersByGroupIDParams{ GroupID: group.ID, - IncludeSystem: false, + IncludeSystem: true, }) if err != nil { httpapi.InternalServerError(rw, err) @@ -174,7 +174,7 @@ func (api *API) patchGroup(rw http.ResponseWriter, r *http.Request) { _, err := database.ExpectOne(api.Database.OrganizationMembers(ctx, database.OrganizationMembersParams{ OrganizationID: group.OrganizationID, UserID: uuid.MustParse(id), - IncludeSystem: false, + IncludeSystem: true, })) if errors.Is(err, sql.ErrNoRows) { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ @@ -290,7 +290,7 @@ func (api *API) patchGroup(rw http.ResponseWriter, r *http.Request) { patchedMembers, err := api.Database.GetGroupMembersByGroupID(ctx, database.GetGroupMembersByGroupIDParams{ GroupID: group.ID, - IncludeSystem: false, + IncludeSystem: true, }) if err != nil { httpapi.InternalServerError(rw, err) @@ -301,7 +301,7 @@ func (api *API) patchGroup(rw http.ResponseWriter, r *http.Request) { memberCount, err := api.Database.GetGroupMembersCountByGroupID(ctx, database.GetGroupMembersCountByGroupIDParams{ GroupID: group.ID, - IncludeSystem: false, + IncludeSystem: true, }) if err != nil { httpapi.InternalServerError(rw, err) @@ -347,7 +347,7 @@ func (api *API) deleteGroup(rw http.ResponseWriter, r *http.Request) { groupMembers, getMembersErr := api.Database.GetGroupMembersByGroupID(ctx, database.GetGroupMembersByGroupIDParams{ GroupID: group.ID, - IncludeSystem: false, + IncludeSystem: true, }) if getMembersErr != nil { httpapi.InternalServerError(rw, getMembersErr) @@ -401,7 +401,7 @@ func (api *API) group(rw http.ResponseWriter, r *http.Request) { users, err := api.Database.GetGroupMembersByGroupID(ctx, database.GetGroupMembersByGroupIDParams{ GroupID: group.ID, - IncludeSystem: false, + IncludeSystem: true, }) if err != nil && !errors.Is(err, sql.ErrNoRows) { httpapi.InternalServerError(rw, err) @@ -410,7 +410,7 @@ func (api *API) group(rw http.ResponseWriter, r *http.Request) { memberCount, err := api.Database.GetGroupMembersCountByGroupID(ctx, database.GetGroupMembersCountByGroupIDParams{ GroupID: group.ID, - IncludeSystem: false, + IncludeSystem: true, }) if err != nil { httpapi.InternalServerError(rw, err) @@ -506,7 +506,7 @@ func (api *API) groups(rw http.ResponseWriter, r *http.Request) { for _, group := range groups { members, err := api.Database.GetGroupMembersByGroupID(ctx, database.GetGroupMembersByGroupIDParams{ GroupID: group.Group.ID, - IncludeSystem: false, + IncludeSystem: true, }) if err != nil { httpapi.InternalServerError(rw, err) @@ -514,7 +514,7 @@ func (api *API) groups(rw http.ResponseWriter, r *http.Request) { } memberCount, err := api.Database.GetGroupMembersCountByGroupID(ctx, database.GetGroupMembersCountByGroupIDParams{ GroupID: group.Group.ID, - IncludeSystem: false, + IncludeSystem: true, }) if err != nil { httpapi.InternalServerError(rw, err) From 4b63826647fdc366c279dfbf9c328cf258eb3d68 Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Thu, 12 Jun 2025 09:05:07 +0000 Subject: [PATCH 2/9] Fix Tests --- enterprise/coderd/groups_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/enterprise/coderd/groups_test.go b/enterprise/coderd/groups_test.go index 028aa3328535f..6f19e585cc550 100644 --- a/enterprise/coderd/groups_test.go +++ b/enterprise/coderd/groups_test.go @@ -838,12 +838,12 @@ func TestGroup(t *testing.T) { // The 'Everyone' group always has an ID that matches the organization ID. group, err := userAdminClient.Group(ctx, user.OrganizationID) require.NoError(t, err) - require.Len(t, group.Members, 4) + require.Len(t, group.Members, 5) require.Equal(t, "Everyone", group.Name) require.Equal(t, user.OrganizationID, group.OrganizationID) require.Contains(t, group.Members, user1.ReducedUser) require.Contains(t, group.Members, user2.ReducedUser) - require.NotContains(t, group.Members, prebuildsUser.ReducedUser) + require.Contains(t, group.Members, prebuildsUser.ReducedUser) }) } From af7c7cd3016a2ba7daaf410fa838822ba33bcbf7 Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Tue, 15 Jul 2025 09:47:02 +0000 Subject: [PATCH 3/9] fix tests --- coderd/database/queries.sql.go | 18 +++++++++++++++--- .../database/queries/organizationmembers.sql | 6 ++++++ coderd/members.go | 3 ++- coderd/members_test.go | 5 +++-- enterprise/coderd/roles_test.go | 7 ++++--- 5 files changed, 30 insertions(+), 9 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 04ded71f1242a..1511d92305a75 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -5993,16 +5993,23 @@ WHERE organization_id = $1 ELSE true END + -- Filter by system type + AND CASE + WHEN $2::bool THEN TRUE + ELSE + is_system = false + END ORDER BY -- Deterministic and consistent ordering of all users. This is to ensure consistent pagination. - LOWER(username) ASC OFFSET $2 + LOWER(username) ASC OFFSET $3 LIMIT -- A null limit means "no limit", so 0 means return all - NULLIF($3 :: int, 0) + NULLIF($4 :: int, 0) ` type PaginatedOrganizationMembersParams struct { OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + IncludeSystem bool `db:"include_system" json:"include_system"` OffsetOpt int32 `db:"offset_opt" json:"offset_opt"` LimitOpt int32 `db:"limit_opt" json:"limit_opt"` } @@ -6018,7 +6025,12 @@ type PaginatedOrganizationMembersRow struct { } func (q *sqlQuerier) PaginatedOrganizationMembers(ctx context.Context, arg PaginatedOrganizationMembersParams) ([]PaginatedOrganizationMembersRow, error) { - rows, err := q.db.QueryContext(ctx, paginatedOrganizationMembers, arg.OrganizationID, arg.OffsetOpt, arg.LimitOpt) + rows, err := q.db.QueryContext(ctx, paginatedOrganizationMembers, + arg.OrganizationID, + arg.IncludeSystem, + arg.OffsetOpt, + arg.LimitOpt, + ) if err != nil { return nil, err } diff --git a/coderd/database/queries/organizationmembers.sql b/coderd/database/queries/organizationmembers.sql index 9d570bc1c49ee..692280415f1ea 100644 --- a/coderd/database/queries/organizationmembers.sql +++ b/coderd/database/queries/organizationmembers.sql @@ -89,6 +89,12 @@ WHERE organization_id = @organization_id ELSE true END + -- Filter by system type + AND CASE + WHEN @include_system::bool THEN TRUE + ELSE + is_system = false + END ORDER BY -- Deterministic and consistent ordering of all users. This is to ensure consistent pagination. LOWER(username) ASC OFFSET @offset_opt diff --git a/coderd/members.go b/coderd/members.go index 5a031fe7eab90..db4a81a31a8bb 100644 --- a/coderd/members.go +++ b/coderd/members.go @@ -161,7 +161,7 @@ func (api *API) listMembers(rw http.ResponseWriter, r *http.Request) { members, err := api.Database.OrganizationMembers(ctx, database.OrganizationMembersParams{ OrganizationID: organization.ID, UserID: uuid.Nil, - IncludeSystem: false, + IncludeSystem: true, }) if httpapi.Is404Error(err) { httpapi.ResourceNotFound(rw) @@ -203,6 +203,7 @@ func (api *API) paginatedMembers(rw http.ResponseWriter, r *http.Request) { paginatedMemberRows, err := api.Database.PaginatedOrganizationMembers(ctx, database.PaginatedOrganizationMembersParams{ OrganizationID: organization.ID, + IncludeSystem: true, // #nosec G115 - Pagination limits are small and fit in int32 LimitOpt: int32(paginationParams.Limit), // #nosec G115 - Pagination offsets are small and fit in int32 diff --git a/coderd/members_test.go b/coderd/members_test.go index bc892bb0679d4..5e7032ba20b1c 100644 --- a/coderd/members_test.go +++ b/coderd/members_test.go @@ -7,6 +7,7 @@ import ( "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/db2sdk" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/codersdk" @@ -62,9 +63,9 @@ func TestListMembers(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) members, err := client.OrganizationMembers(ctx, first.OrganizationID) require.NoError(t, err) - require.Len(t, members, 2) + require.Len(t, members, 3) require.ElementsMatch(t, - []uuid.UUID{first.UserID, user.ID}, + []uuid.UUID{first.UserID, user.ID, database.PrebuildsSystemUserID}, db2sdk.List(members, onlyIDs)) }) } diff --git a/enterprise/coderd/roles_test.go b/enterprise/coderd/roles_test.go index 70c432755f7fa..79dbf4ed2c98a 100644 --- a/enterprise/coderd/roles_test.go +++ b/enterprise/coderd/roles_test.go @@ -11,6 +11,7 @@ import ( "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/db2sdk" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/codersdk" @@ -360,9 +361,9 @@ func TestCustomOrganizationRole(t *testing.T) { // Verify members have the custom role originalMembers, err := orgAdmin.OrganizationMembers(ctx, first.OrganizationID) require.NoError(t, err) - require.Len(t, originalMembers, 5) // 3 members + org admin + owner + require.Len(t, originalMembers, 6) // 3 members + org admin + owner + prebuilds system user for _, member := range originalMembers { - if member.UserID == orgAdminUser.ID || member.UserID == first.UserID { + if member.UserID == orgAdminUser.ID || member.UserID == first.UserID || member.UserID == database.PrebuildsSystemUserID { continue } @@ -377,7 +378,7 @@ func TestCustomOrganizationRole(t *testing.T) { // Verify the role was removed from all members members, err := orgAdmin.OrganizationMembers(ctx, first.OrganizationID) require.NoError(t, err) - require.Len(t, members, 5) // 3 members + org admin + owner + require.Len(t, members, 6) // 3 members + org admin + owner + prebuilds system user for _, member := range members { require.False(t, slices.ContainsFunc(member.Roles, func(role codersdk.SlimRole) bool { return role.Name == customRoleIdentifier.Name From 1cd68e3df3411cf7fe241e8d2618b53973b556b5 Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Tue, 29 Jul 2025 07:33:17 +0000 Subject: [PATCH 4/9] create a prebuilds group in every org that needs it --- coderd/database/dbauthz/dbauthz.go | 8 + .../prebuilt-workspaces.md | 8 +- enterprise/coderd/prebuilds/membership.go | 71 +++-- .../coderd/prebuilds/membership_test.go | 242 ++++++++++++------ 4 files changed, 233 insertions(+), 96 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 55665b4381862..438cc7c30359c 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -467,6 +467,14 @@ var ( rbac.ResourceFile.Type: { policy.ActionRead, }, + rbac.ResourceGroup.Type: { + policy.ActionRead, + policy.ActionCreate, + policy.ActionUpdate, + }, + rbac.ResourceGroupMember.Type: { + policy.ActionRead, + }, }), }, }), diff --git a/docs/admin/templates/extending-templates/prebuilt-workspaces.md b/docs/admin/templates/extending-templates/prebuilt-workspaces.md index 8e61687ce0f01..70161f687ba76 100644 --- a/docs/admin/templates/extending-templates/prebuilt-workspaces.md +++ b/docs/admin/templates/extending-templates/prebuilt-workspaces.md @@ -235,12 +235,18 @@ The system always maintains the desired number of prebuilt workspaces for the ac ### Managing resource quotas -Prebuilt workspaces can be used in conjunction with [resource quotas](../../users/quotas.md). +To help prevent unexpected infrastructure costs, prebuilt workspaces can be used in conjunction with [resource quotas](../../users/quotas.md). Because unclaimed prebuilt workspaces are owned by the `prebuilds` user, you can: 1. Configure quotas for any group that includes this user. 1. Set appropriate limits to balance prebuilt workspace availability with resource constraints. +When prebuilt workspaces are configured for an organization, Coder creates a "prebuilds" group in that organization and adds the prebuilds user to it. This group has a default quota allowance of 0, which you should adjust based on your needs: + +- **Set a quota allowance** on the "prebuilds" group to control how many prebuilt workspaces can be provisioned +- **Monitor usage** to ensure the quota is appropriate for your desired number of prebuilt instances +- **Adjust as needed** based on your template costs and desired prebuilt workspace pool size + If a quota is exceeded, the prebuilt workspace will fail provisioning the same way other workspaces do. ### Template configuration best practices diff --git a/enterprise/coderd/prebuilds/membership.go b/enterprise/coderd/prebuilds/membership.go index 079711bcbcc49..9b41d68787ea1 100644 --- a/enterprise/coderd/prebuilds/membership.go +++ b/enterprise/coderd/prebuilds/membership.go @@ -44,37 +44,74 @@ func (s StoreMembershipReconciler) ReconcileAll(ctx context.Context, userID uuid return xerrors.Errorf("determine prebuild organization membership: %w", err) } - systemUserMemberships := make(map[uuid.UUID]struct{}, 0) + orgMemberShips := make(map[uuid.UUID]struct{}, 0) defaultOrg, err := s.store.GetDefaultOrganization(ctx) if err != nil { return xerrors.Errorf("get default organization: %w", err) } - systemUserMemberships[defaultOrg.ID] = struct{}{} + orgMemberShips[defaultOrg.ID] = struct{}{} for _, o := range organizationMemberships { - systemUserMemberships[o.ID] = struct{}{} + orgMemberShips[o.ID] = struct{}{} } var membershipInsertionErrors error for _, preset := range presets { - _, alreadyMember := systemUserMemberships[preset.OrganizationID] - if alreadyMember { - continue + _, alreadyOrgMember := orgMemberShips[preset.OrganizationID] + if !alreadyOrgMember { + // Add the organization to our list of memberships regardless of potential failure below + // to avoid a retry that will probably be doomed anyway. + orgMemberShips[preset.OrganizationID] = struct{}{} + + // Insert the missing membership + _, err = s.store.InsertOrganizationMember(ctx, database.InsertOrganizationMemberParams{ + OrganizationID: preset.OrganizationID, + UserID: userID, + CreatedAt: s.clock.Now(), + UpdatedAt: s.clock.Now(), + Roles: []string{}, + }) + if err != nil { + membershipInsertionErrors = errors.Join(membershipInsertionErrors, xerrors.Errorf("insert membership for prebuilt workspaces: %w", err)) + continue + } } - // Add the organization to our list of memberships regardless of potential failure below - // to avoid a retry that will probably be doomed anyway. - systemUserMemberships[preset.OrganizationID] = struct{}{} - // Insert the missing membership - _, err = s.store.InsertOrganizationMember(ctx, database.InsertOrganizationMemberParams{ + // Create a "prebuilds" group in the organization and add the system user to it + // This group will have a quota of 0 by default, which users can adjust based on their needs + prebuildsGroup, err := s.store.InsertGroup(ctx, database.InsertGroupParams{ + ID: uuid.New(), + Name: "prebuilds", + DisplayName: "Prebuilds", OrganizationID: preset.OrganizationID, - UserID: userID, - CreatedAt: s.clock.Now(), - UpdatedAt: s.clock.Now(), - Roles: []string{}, + AvatarURL: "", + QuotaAllowance: 0, // Default quota of 0, users should set this based on their needs + }) + if err != nil { + // If the group already exists, try to get it + if !database.IsUniqueViolation(err) { + membershipInsertionErrors = errors.Join(membershipInsertionErrors, xerrors.Errorf("create prebuilds group: %w", err)) + continue + } + prebuildsGroup, err = s.store.GetGroupByOrgAndName(ctx, database.GetGroupByOrgAndNameParams{ + OrganizationID: preset.OrganizationID, + Name: "prebuilds", + }) + if err != nil { + membershipInsertionErrors = errors.Join(membershipInsertionErrors, xerrors.Errorf("get existing prebuilds group: %w", err)) + continue + } + } + + // Add the system user to the prebuilds group + err = s.store.InsertGroupMember(ctx, database.InsertGroupMemberParams{ + GroupID: prebuildsGroup.ID, + UserID: userID, }) if err != nil { - membershipInsertionErrors = errors.Join(membershipInsertionErrors, xerrors.Errorf("insert membership for prebuilt workspaces: %w", err)) - continue + // Ignore unique violation errors as the user might already be in the group + if !database.IsUniqueViolation(err) { + membershipInsertionErrors = errors.Join(membershipInsertionErrors, xerrors.Errorf("add system user to prebuilds group: %w", err)) + } } } return membershipInsertionErrors diff --git a/enterprise/coderd/prebuilds/membership_test.go b/enterprise/coderd/prebuilds/membership_test.go index 82d2abf92a4d8..06fe4af576fc6 100644 --- a/enterprise/coderd/prebuilds/membership_test.go +++ b/enterprise/coderd/prebuilds/membership_test.go @@ -1,18 +1,23 @@ package prebuilds_test import ( - "context" + "database/sql" + "errors" "testing" "github.com/google/uuid" "github.com/stretchr/testify/require" + "tailscale.com/types/ptr" "github.com/coder/quartz" + "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/enterprise/coderd/prebuilds" + "github.com/coder/coder/v2/testutil" ) // TestReconcileAll verifies that StoreMembershipReconciler correctly updates membership @@ -20,7 +25,7 @@ import ( func TestReconcileAll(t *testing.T) { t.Parallel() - ctx := context.Background() + ctx := dbauthz.AsPrebuildsOrchestrator(testutil.Context(t, testutil.WaitLong)) clock := quartz.NewMock(t) // Helper to build a minimal Preset row belonging to a given org. @@ -32,87 +37,168 @@ func TestReconcileAll(t *testing.T) { } tests := []struct { - name string - includePreset bool - preExistingMembership bool + name string + includePreset []bool + preExistingOrgMembership []bool + preExistingGroup []bool + preExistingGroupMembership []bool + // Expected outcomes + expectOrgMembershipExists *bool + expectGroupExists *bool + expectUserInGroup *bool }{ - // The StoreMembershipReconciler acts based on the provided agplprebuilds.GlobalSnapshot. - // These test cases must therefore trust any valid snapshot, so the only relevant functional test cases are: - - // No presets to act on and the prebuilds user does not belong to any organizations. - // Reconciliation should be a no-op - {name: "no presets, no memberships", includePreset: false, preExistingMembership: false}, - // If we have a preset that requires prebuilds, but the prebuilds user is not a member of - // that organization, then we should add the membership. - {name: "preset, but no membership", includePreset: true, preExistingMembership: false}, - // If the prebuilds system user is already a member of the organization to which a preset belongs, - // then reconciliation should be a no-op: - {name: "preset, but already a member", includePreset: true, preExistingMembership: true}, - // If the prebuilds system user is a member of an organization that doesn't have need any prebuilds, - // then it must have required prebuilds in the past. The membership is not currently necessary, but - // the reconciler won't remove it, because there's little cost to keeping it and prebuilds might be - // enabled again. - {name: "member, but no presets", includePreset: false, preExistingMembership: true}, + { + name: "if there are no presets, membership reconciliation is a no-op", + includePreset: []bool{false}, + preExistingOrgMembership: []bool{true, false}, + preExistingGroup: []bool{true, false}, + preExistingGroupMembership: []bool{true, false}, + expectOrgMembershipExists: ptr.To(false), + expectGroupExists: ptr.To(false), + }, + { + name: "if there is a preset, then we should enforce org and group membership in all cases", + includePreset: []bool{true}, + preExistingOrgMembership: []bool{true, false}, + preExistingGroup: []bool{true, false}, + preExistingGroupMembership: []bool{true, false}, + expectOrgMembershipExists: ptr.To(true), + expectGroupExists: ptr.To(true), + expectUserInGroup: ptr.To(true), + }, } for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - - db, _ := dbtestutil.NewDB(t) - - defaultOrg, err := db.GetDefaultOrganization(ctx) - require.NoError(t, err) - - // introduce an unrelated organization to ensure that the membership reconciler don't interfere with it. - unrelatedOrg := dbgen.Organization(t, db, database.Organization{}) - targetOrg := dbgen.Organization(t, db, database.Organization{}) - - if !dbtestutil.WillUsePostgres() { - // dbmem doesn't ensure membership to the default organization - dbgen.OrganizationMember(t, db, database.OrganizationMember{ - OrganizationID: defaultOrg.ID, - UserID: database.PrebuildsSystemUserID, - }) - } - - dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: unrelatedOrg.ID, UserID: database.PrebuildsSystemUserID}) - if tc.preExistingMembership { - // System user already a member of both orgs. - dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: targetOrg.ID, UserID: database.PrebuildsSystemUserID}) + tc := tc + for _, includePreset := range tc.includePreset { + includePreset := includePreset + for _, preExistingOrgMembership := range tc.preExistingOrgMembership { + preExistingOrgMembership := preExistingOrgMembership + for _, preExistingGroup := range tc.preExistingGroup { + preExistingGroup := preExistingGroup + for _, preExistingGroupMembership := range tc.preExistingGroupMembership { + preExistingGroupMembership := preExistingGroupMembership + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + _, db := coderdtest.NewWithDatabase(t, nil) + + defaultOrg, err := db.GetDefaultOrganization(ctx) + require.NoError(t, err) + + // introduce an unrelated organization to ensure that the membership reconciler don't interfere with it. + unrelatedOrg := dbgen.Organization(t, db, database.Organization{}) + targetOrg := dbgen.Organization(t, db, database.Organization{}) + + if !dbtestutil.WillUsePostgres() { + // dbmem doesn't ensure membership to the default organization + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + OrganizationID: defaultOrg.ID, + UserID: database.PrebuildsSystemUserID, + }) + } + + // Ensure membership to unrelated org. + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: unrelatedOrg.ID, UserID: database.PrebuildsSystemUserID}) + + if preExistingOrgMembership { + // System user already a member of both orgs. + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: targetOrg.ID, UserID: database.PrebuildsSystemUserID}) + } + + // Create pre-existing prebuilds group if required by test case + var prebuildsGroup database.Group + if preExistingGroup { + prebuildsGroup = dbgen.Group(t, db, database.Group{ + Name: "prebuilds", + DisplayName: "Prebuilds", + OrganizationID: targetOrg.ID, + QuotaAllowance: 0, + }) + + // Add the system user to the group if preExistingGroupMembership is true + if preExistingGroupMembership { + dbgen.GroupMember(t, db, database.GroupMemberTable{ + GroupID: prebuildsGroup.ID, + UserID: database.PrebuildsSystemUserID, + }) + } + } + + presets := []database.GetTemplatePresetsWithPrebuildsRow{newPresetRow(unrelatedOrg.ID)} + if includePreset { + presets = append(presets, newPresetRow(targetOrg.ID)) + } + + // Verify memberships before reconciliation. + preReconcileMemberships, err := db.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{ + UserID: database.PrebuildsSystemUserID, + }) + require.NoError(t, err) + expectedMembershipsBefore := []uuid.UUID{defaultOrg.ID, unrelatedOrg.ID} + if preExistingOrgMembership { + expectedMembershipsBefore = append(expectedMembershipsBefore, targetOrg.ID) + } + require.ElementsMatch(t, expectedMembershipsBefore, extractOrgIDs(preReconcileMemberships)) + + // Reconcile + reconciler := prebuilds.NewStoreMembershipReconciler(db, clock) + require.NoError(t, reconciler.ReconcileAll(ctx, database.PrebuildsSystemUserID, presets)) + + // Verify memberships after reconciliation. + postReconcileMemberships, err := db.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{ + UserID: database.PrebuildsSystemUserID, + }) + require.NoError(t, err) + expectedMembershipsAfter := expectedMembershipsBefore + if !preExistingOrgMembership && tc.expectOrgMembershipExists != nil && *tc.expectOrgMembershipExists { + expectedMembershipsAfter = append(expectedMembershipsAfter, targetOrg.ID) + } + require.ElementsMatch(t, expectedMembershipsAfter, extractOrgIDs(postReconcileMemberships)) + + // Verify prebuilds group behavior based on expected outcomes + prebuildsGroup, err = db.GetGroupByOrgAndName(ctx, database.GetGroupByOrgAndNameParams{ + OrganizationID: targetOrg.ID, + Name: "prebuilds", + }) + if tc.expectGroupExists != nil && *tc.expectGroupExists { + require.NoError(t, err) + require.Equal(t, "prebuilds", prebuildsGroup.Name) + require.Equal(t, "Prebuilds", prebuildsGroup.DisplayName) + require.Equal(t, int32(0), prebuildsGroup.QuotaAllowance) // Default quota should be 0 + + if tc.expectUserInGroup != nil && *tc.expectUserInGroup { + // Check that the system user is a member of the prebuilds group + groupMembers, err := db.GetGroupMembersByGroupID(ctx, database.GetGroupMembersByGroupIDParams{ + GroupID: prebuildsGroup.ID, + IncludeSystem: true, + }) + require.NoError(t, err) + require.Len(t, groupMembers, 1) + require.Equal(t, database.PrebuildsSystemUserID, groupMembers[0].UserID) + } + + if tc.expectUserInGroup != nil && !*tc.expectUserInGroup { + // Check that the system user is NOT a member of the prebuilds group + groupMembers, err := db.GetGroupMembersByGroupID(ctx, database.GetGroupMembersByGroupIDParams{ + GroupID: prebuildsGroup.ID, + IncludeSystem: true, + }) + require.NoError(t, err) + require.Len(t, groupMembers, 0) + } + } + + if !preExistingGroup && tc.expectGroupExists != nil && !*tc.expectGroupExists { + // Verify that no prebuilds group exists + require.Error(t, err) + require.True(t, errors.Is(err, sql.ErrNoRows)) + } + }) + } + } } - - presets := []database.GetTemplatePresetsWithPrebuildsRow{newPresetRow(unrelatedOrg.ID)} - if tc.includePreset { - presets = append(presets, newPresetRow(targetOrg.ID)) - } - - // Verify memberships before reconciliation. - preReconcileMemberships, err := db.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{ - UserID: database.PrebuildsSystemUserID, - }) - require.NoError(t, err) - expectedMembershipsBefore := []uuid.UUID{defaultOrg.ID, unrelatedOrg.ID} - if tc.preExistingMembership { - expectedMembershipsBefore = append(expectedMembershipsBefore, targetOrg.ID) - } - require.ElementsMatch(t, expectedMembershipsBefore, extractOrgIDs(preReconcileMemberships)) - - // Reconcile - reconciler := prebuilds.NewStoreMembershipReconciler(db, clock) - require.NoError(t, reconciler.ReconcileAll(ctx, database.PrebuildsSystemUserID, presets)) - - // Verify memberships after reconciliation. - postReconcileMemberships, err := db.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{ - UserID: database.PrebuildsSystemUserID, - }) - require.NoError(t, err) - expectedMembershipsAfter := expectedMembershipsBefore - if !tc.preExistingMembership && tc.includePreset { - expectedMembershipsAfter = append(expectedMembershipsAfter, targetOrg.ID) - } - require.ElementsMatch(t, expectedMembershipsAfter, extractOrgIDs(postReconcileMemberships)) - }) + } } } From 93c46e91a89006c96a0cbf5cbe377e7ab257c350 Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Tue, 29 Jul 2025 10:04:38 +0000 Subject: [PATCH 5/9] add lint rule --- enterprise/coderd/prebuilds/membership_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/enterprise/coderd/prebuilds/membership_test.go b/enterprise/coderd/prebuilds/membership_test.go index 06fe4af576fc6..396bb55751242 100644 --- a/enterprise/coderd/prebuilds/membership_test.go +++ b/enterprise/coderd/prebuilds/membership_test.go @@ -25,6 +25,7 @@ import ( func TestReconcileAll(t *testing.T) { t.Parallel() + // nolint:gocritic // Reconciliation happens as prebuilds system user, not a human user. ctx := dbauthz.AsPrebuildsOrchestrator(testutil.Context(t, testutil.WaitLong)) clock := quartz.NewMock(t) From b3728d3e3f4eb5423f425baf9f3d972d338922d2 Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Wed, 30 Jul 2025 10:25:37 +0000 Subject: [PATCH 6/9] add quotas tests to verify prebuilds behaviour --- coderd/database/queries.sql.go | 42 +-- coderd/database/queries/groupmembers.sql | 18 +- .../database/queries/organizationmembers.sql | 12 +- coderd/database/queries/users.sql | 12 +- enterprise/coderd/workspacequota_test.go | 260 ++++++++++++++++++ 5 files changed, 278 insertions(+), 66 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 2b7780f05592a..a091c82955bba 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2248,11 +2248,7 @@ func (q *sqlQuerier) DeleteGroupMemberFromGroup(ctx context.Context, arg DeleteG const getGroupMembers = `-- name: GetGroupMembers :many SELECT user_id, user_email, user_username, user_hashed_password, user_created_at, user_updated_at, user_status, user_rbac_roles, user_login_type, user_avatar_url, user_deleted, user_last_seen_at, user_quiet_hours_schedule, user_name, user_github_com_user_id, user_is_system, organization_id, group_name, group_id FROM group_members_expanded -WHERE CASE - WHEN $1::bool THEN TRUE - ELSE - user_is_system = false - END +WHERE ($1::bool OR NOT user_is_system) ` func (q *sqlQuerier) GetGroupMembers(ctx context.Context, includeSystem bool) ([]GroupMember, error) { @@ -2303,11 +2299,7 @@ SELECT user_id, user_email, user_username, user_hashed_password, user_created_at FROM group_members_expanded WHERE group_id = $1 -- Filter by system type - AND CASE - WHEN $2::bool THEN TRUE - ELSE - user_is_system = false - END + AND ($2::bool OR NOT user_is_system) ` type GetGroupMembersByGroupIDParams struct { @@ -2363,11 +2355,7 @@ SELECT COUNT(*) FROM group_members_expanded WHERE group_id = $1 -- Filter by system type - AND CASE - WHEN $2::bool THEN TRUE - ELSE - user_is_system = false - END + AND ($2::bool OR NOT user_is_system) ` type GetGroupMembersCountByGroupIDParams struct { @@ -6485,11 +6473,7 @@ WHERE ELSE true END -- Filter by system type - AND CASE - WHEN $3::bool THEN TRUE - ELSE - is_system = false - END + AND ($3::bool OR NOT is_system) ` type OrganizationMembersParams struct { @@ -6562,11 +6546,7 @@ WHERE ELSE true END -- Filter by system type - AND CASE - WHEN $2::bool THEN TRUE - ELSE - is_system = false - END + AND ($2::bool OR NOT is_system) ORDER BY -- Deterministic and consistent ordering of all users. This is to ensure consistent pagination. LOWER(username) ASC OFFSET $3 @@ -13716,7 +13696,7 @@ func (q *sqlQuerier) UpdateUserLinkedID(ctx context.Context, arg UpdateUserLinke const allUserIDs = `-- name: AllUserIDs :many SELECT DISTINCT id FROM USERS - WHERE CASE WHEN $1::bool THEN TRUE ELSE is_system = false END + WHERE ($1::bool OR NOT is_system) ` // AllUserIDs returns all UserIDs regardless of user status or deletion. @@ -13750,7 +13730,7 @@ FROM users WHERE status = 'active'::user_status AND deleted = false - AND CASE WHEN $1::bool THEN TRUE ELSE is_system = false END + AND ($1::bool OR NOT is_system) ` func (q *sqlQuerier) GetActiveUserCount(ctx context.Context, includeSystem bool) (int64, error) { @@ -13914,7 +13894,7 @@ FROM users WHERE deleted = false - AND CASE WHEN $1::bool THEN TRUE ELSE is_system = false END + AND ($1::bool OR NOT is_system) ` func (q *sqlQuerier) GetUserCount(ctx context.Context, includeSystem bool) (int64, error) { @@ -14031,11 +14011,7 @@ WHERE created_at >= $8 ELSE true END - AND CASE - WHEN $9::bool THEN TRUE - ELSE - is_system = false - END + AND ($9::bool OR NOT is_system) AND CASE WHEN $10 :: bigint != 0 THEN github_com_user_id = $10 diff --git a/coderd/database/queries/groupmembers.sql b/coderd/database/queries/groupmembers.sql index 7de8dbe4e4523..2296e871939dc 100644 --- a/coderd/database/queries/groupmembers.sql +++ b/coderd/database/queries/groupmembers.sql @@ -1,21 +1,13 @@ -- name: GetGroupMembers :many SELECT * FROM group_members_expanded -WHERE CASE - WHEN @include_system::bool THEN TRUE - ELSE - user_is_system = false - END; +WHERE (@include_system::bool OR NOT user_is_system); -- name: GetGroupMembersByGroupID :many SELECT * FROM group_members_expanded WHERE group_id = @group_id -- Filter by system type - AND CASE - WHEN @include_system::bool THEN TRUE - ELSE - user_is_system = false - END; + AND (@include_system::bool OR NOT user_is_system); -- name: GetGroupMembersCountByGroupID :one -- Returns the total count of members in a group. Shows the total @@ -25,11 +17,7 @@ SELECT COUNT(*) FROM group_members_expanded WHERE group_id = @group_id -- Filter by system type - AND CASE - WHEN @include_system::bool THEN TRUE - ELSE - user_is_system = false - END; + AND (@include_system::bool OR NOT user_is_system); -- InsertUserGroupsByName adds a user to all provided groups, if they exist. -- name: InsertUserGroupsByName :exec diff --git a/coderd/database/queries/organizationmembers.sql b/coderd/database/queries/organizationmembers.sql index 692280415f1ea..688a6ddd69799 100644 --- a/coderd/database/queries/organizationmembers.sql +++ b/coderd/database/queries/organizationmembers.sql @@ -24,11 +24,7 @@ WHERE ELSE true END -- Filter by system type - AND CASE - WHEN @include_system::bool THEN TRUE - ELSE - is_system = false - END; + AND (@include_system::bool OR NOT is_system); -- name: InsertOrganizationMember :one INSERT INTO @@ -90,11 +86,7 @@ WHERE ELSE true END -- Filter by system type - AND CASE - WHEN @include_system::bool THEN TRUE - ELSE - is_system = false - END + AND (@include_system::bool OR NOT is_system) ORDER BY -- Deterministic and consistent ordering of all users. This is to ensure consistent pagination. LOWER(username) ASC OFFSET @offset_opt diff --git a/coderd/database/queries/users.sql b/coderd/database/queries/users.sql index eece2f96512ea..033a1abf8ed98 100644 --- a/coderd/database/queries/users.sql +++ b/coderd/database/queries/users.sql @@ -49,7 +49,7 @@ FROM users WHERE deleted = false - AND CASE WHEN @include_system::bool THEN TRUE ELSE is_system = false END; + AND (@include_system::bool OR NOT is_system); -- name: GetActiveUserCount :one SELECT @@ -58,7 +58,7 @@ FROM users WHERE status = 'active'::user_status AND deleted = false - AND CASE WHEN @include_system::bool THEN TRUE ELSE is_system = false END; + AND (@include_system::bool OR NOT is_system); -- name: InsertUser :one INSERT INTO @@ -250,11 +250,7 @@ WHERE created_at >= @created_after ELSE true END - AND CASE - WHEN @include_system::bool THEN TRUE - ELSE - is_system = false - END + AND (@include_system::bool OR NOT is_system) AND CASE WHEN @github_com_user_id :: bigint != 0 THEN github_com_user_id = @github_com_user_id @@ -364,7 +360,7 @@ RETURNING id, email, username, last_seen_at; -- AllUserIDs returns all UserIDs regardless of user status or deletion. -- name: AllUserIDs :many SELECT DISTINCT id FROM USERS - WHERE CASE WHEN @include_system::bool THEN TRUE ELSE is_system = false END; + WHERE (@include_system::bool OR NOT is_system); -- name: UpdateUserHashedOneTimePasscode :exec UPDATE diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index f49e135ad55b3..b40b9b7c8abfd 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -395,6 +395,266 @@ func TestWorkspaceQuota(t *testing.T) { verifyQuotaUser(ctx, t, client, second.Org.ID.String(), user.ID.String(), consumed, 35) }) + + // ZeroQuota tests that a user with a zero quota allowance can't create a workspace. + // Although relevant for all users, this test ensures that the prebuilds system user + // cannot create workspaces in an organization for which it has exhausted its quota. + t.Run("ZeroQuota", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) + defer cancel() + + // Create a client with no quota allowance + client, _, api, user := coderdenttest.NewWithAPI(t, &coderdenttest.Options{ + UserWorkspaceQuota: 0, // Set user workspace quota to 0 + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + }, + }, + }) + coderdtest.NewProvisionerDaemon(t, api.AGPL) + + // Verify initial quota is 0 + verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 0) + + // Create a template with a workspace that costs 1 credit + authToken := uuid.NewString() + version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{ + Parse: echo.ParseComplete, + ProvisionApply: []*proto.Response{{ + Type: &proto.Response_Apply{ + Apply: &proto.ApplyComplete{ + Resources: []*proto.Resource{{ + Name: "example", + Type: "aws_instance", + DailyCost: 1, + Agents: []*proto.Agent{{ + Id: uuid.NewString(), + Name: "example", + Auth: &proto.Agent_Token{ + Token: authToken, + }, + }}, + }}, + }, + }, + }}, + }) + coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) + template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID) + + // Attempt to create a workspace with zero quota - should fail + workspace := coderdtest.CreateWorkspace(t, client, template.ID) + build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID) + + // Verify the build failed due to quota + require.Equal(t, codersdk.WorkspaceStatusFailed, build.Status) + require.Contains(t, build.Job.Error, "quota") + + // Verify quota consumption remains at 0 + verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 0) + + // Test with a template that has zero cost - should pass + versionZeroCost := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{ + Parse: echo.ParseComplete, + ProvisionApply: []*proto.Response{{ + Type: &proto.Response_Apply{ + Apply: &proto.ApplyComplete{ + Resources: []*proto.Resource{{ + Name: "example", + Type: "aws_instance", + DailyCost: 0, // Zero cost workspace + Agents: []*proto.Agent{{ + Id: uuid.NewString(), + Name: "example", + Auth: &proto.Agent_Token{ + Token: uuid.NewString(), + }, + }}, + }}, + }, + }, + }}, + }) + coderdtest.AwaitTemplateVersionJobCompleted(t, client, versionZeroCost.ID) + templateZeroCost := coderdtest.CreateTemplate(t, client, user.OrganizationID, versionZeroCost.ID) + + // Even with zero cost, should pass + workspaceZeroCost := coderdtest.CreateWorkspace(t, client, templateZeroCost.ID) + buildZeroCost := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspaceZeroCost.LatestBuild.ID) + + // Verify the build failed due to quota + require.Equal(t, codersdk.WorkspaceStatusRunning, buildZeroCost.Status) + require.Empty(t, buildZeroCost.Job.Error) + + // Verify quota consumption remains at 0 + verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 0) + }) + + // MultiOrg tests that a user can create workspaces in multiple organizations + // as long as they have enough quota in each organization. Specifically, + // in exhausted quota in one organization does not affect the ability to + // create workspaces in other organizations. This test is relevant to all users + // but is particularly relevant for the prebuilds system user. + t.Run("MultiOrg", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + + // Create a setup with multiple organizations + owner, _, api, first := coderdenttest.NewWithAPI(t, &coderdenttest.Options{ + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureMultipleOrganizations: 1, + codersdk.FeatureExternalProvisionerDaemons: 1, + }, + }, + }) + coderdtest.NewProvisionerDaemon(t, api.AGPL) + + // Create a second organization + second := coderdenttest.CreateOrganization(t, owner, coderdenttest.CreateOrganizationOptions{ + IncludeProvisionerDaemon: true, + }) + + // Create a user that will be a member of both organizations + user, _ := coderdtest.CreateAnotherUser(t, owner, first.OrganizationID, rbac.ScopedRoleOrgMember(second.ID)) + + // Set up quota allowances for both organizations + // First org: 2 credits total + _, err := owner.PatchGroup(ctx, first.OrganizationID, codersdk.PatchGroupRequest{ + QuotaAllowance: ptr.Ref(2), + }) + require.NoError(t, err) + + // Second org: 3 credits total + _, err = owner.PatchGroup(ctx, second.ID, codersdk.PatchGroupRequest{ + QuotaAllowance: ptr.Ref(3), + }) + require.NoError(t, err) + + // Verify initial quotas + verifyQuota(ctx, t, user, first.OrganizationID.String(), 0, 2) + verifyQuota(ctx, t, user, second.ID.String(), 0, 3) + + // Create templates for both organizations + authToken := uuid.NewString() + version1 := coderdtest.CreateTemplateVersion(t, owner, first.OrganizationID, &echo.Responses{ + Parse: echo.ParseComplete, + ProvisionApply: []*proto.Response{{ + Type: &proto.Response_Apply{ + Apply: &proto.ApplyComplete{ + Resources: []*proto.Resource{{ + Name: "example", + Type: "aws_instance", + DailyCost: 1, + Agents: []*proto.Agent{{ + Id: uuid.NewString(), + Name: "example", + Auth: &proto.Agent_Token{ + Token: authToken, + }, + }}, + }}, + }, + }, + }}, + }) + coderdtest.AwaitTemplateVersionJobCompleted(t, owner, version1.ID) + template1 := coderdtest.CreateTemplate(t, owner, first.OrganizationID, version1.ID) + + version2 := coderdtest.CreateTemplateVersion(t, owner, second.ID, &echo.Responses{ + Parse: echo.ParseComplete, + ProvisionApply: []*proto.Response{{ + Type: &proto.Response_Apply{ + Apply: &proto.ApplyComplete{ + Resources: []*proto.Resource{{ + Name: "example", + Type: "aws_instance", + DailyCost: 1, + Agents: []*proto.Agent{{ + Id: uuid.NewString(), + Name: "example", + Auth: &proto.Agent_Token{ + Token: uuid.NewString(), + }, + }}, + }}, + }, + }, + }}, + }) + coderdtest.AwaitTemplateVersionJobCompleted(t, owner, version2.ID) + template2 := coderdtest.CreateTemplate(t, owner, second.ID, version2.ID) + + // Exhaust quota in the first organization by creating 2 workspaces + var workspaces1 []codersdk.Workspace + for i := 0; i < 2; i++ { + workspace := coderdtest.CreateWorkspace(t, user, template1.ID) + build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, user, workspace.LatestBuild.ID) + require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status) + workspaces1 = append(workspaces1, workspace) + } + + // Verify first org quota is exhausted + verifyQuota(ctx, t, user, first.OrganizationID.String(), 2, 2) + + // Try to create another workspace in the first org - should fail + workspace := coderdtest.CreateWorkspace(t, user, template1.ID) + build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, user, workspace.LatestBuild.ID) + require.Equal(t, codersdk.WorkspaceStatusFailed, build.Status) + require.Contains(t, build.Job.Error, "quota") + + // Verify first org quota consumption didn't increase + verifyQuota(ctx, t, user, first.OrganizationID.String(), 2, 2) + + // Verify second org quota is still available + verifyQuota(ctx, t, user, second.ID.String(), 0, 3) + + // Create workspaces in the second organization - should succeed + for i := 0; i < 3; i++ { + workspace := coderdtest.CreateWorkspace(t, user, template2.ID) + build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, user, workspace.LatestBuild.ID) + require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status) + } + + // Verify second org quota is now exhausted + verifyQuota(ctx, t, user, second.ID.String(), 3, 3) + + // Try to create another workspace in the second org - should fail + workspace = coderdtest.CreateWorkspace(t, user, template2.ID) + build = coderdtest.AwaitWorkspaceBuildJobCompleted(t, user, workspace.LatestBuild.ID) + require.Equal(t, codersdk.WorkspaceStatusFailed, build.Status) + require.Contains(t, build.Job.Error, "quota") + + // Verify second org quota consumption didn't increase + verifyQuota(ctx, t, user, second.ID.String(), 3, 3) + + // Verify first org quota is still exhausted + verifyQuota(ctx, t, user, first.OrganizationID.String(), 2, 2) + + // Delete one workspace from the first org to free up quota + build = coderdtest.CreateWorkspaceBuild(t, user, workspaces1[0], database.WorkspaceTransitionDelete) + build = coderdtest.AwaitWorkspaceBuildJobCompleted(t, user, build.ID) + require.Equal(t, codersdk.WorkspaceStatusDeleted, build.Status) + + // Verify first org quota is now available again + verifyQuota(ctx, t, user, first.OrganizationID.String(), 1, 2) + + // Create a workspace in the first org - should succeed + workspace = coderdtest.CreateWorkspace(t, user, template1.ID) + build = coderdtest.AwaitWorkspaceBuildJobCompleted(t, user, workspace.LatestBuild.ID) + require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status) + + // Verify first org quota is exhausted again + verifyQuota(ctx, t, user, first.OrganizationID.String(), 2, 2) + + // Verify second org quota remains exhausted + verifyQuota(ctx, t, user, second.ID.String(), 3, 3) + }) } // nolint:paralleltest,tparallel // Tests must run serially From efdfe9fa5df2198543c78fbd11f7e1f99e3fea30 Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Wed, 30 Jul 2025 10:37:36 +0000 Subject: [PATCH 7/9] fix unit test context --- enterprise/coderd/prebuilds/membership_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/enterprise/coderd/prebuilds/membership_test.go b/enterprise/coderd/prebuilds/membership_test.go index 396bb55751242..61928d0f1e3b1 100644 --- a/enterprise/coderd/prebuilds/membership_test.go +++ b/enterprise/coderd/prebuilds/membership_test.go @@ -25,8 +25,6 @@ import ( func TestReconcileAll(t *testing.T) { t.Parallel() - // nolint:gocritic // Reconciliation happens as prebuilds system user, not a human user. - ctx := dbauthz.AsPrebuildsOrchestrator(testutil.Context(t, testutil.WaitLong)) clock := quartz.NewMock(t) // Helper to build a minimal Preset row belonging to a given org. @@ -82,6 +80,8 @@ func TestReconcileAll(t *testing.T) { t.Run(tc.name, func(t *testing.T) { t.Parallel() + // nolint:gocritic // Reconciliation happens as prebuilds system user, not a human user. + ctx := dbauthz.AsPrebuildsOrchestrator(testutil.Context(t, testutil.WaitLong)) _, db := coderdtest.NewWithDatabase(t, nil) defaultOrg, err := db.GetDefaultOrganization(ctx) From 8204ec2c30f581e57e665a368013207b7972ba11 Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Mon, 4 Aug 2025 10:32:41 +0000 Subject: [PATCH 8/9] revert broken query changes --- coderd/database/queries/groupmembers.sql | 18 +++++++++++++++--- .../database/queries/organizationmembers.sql | 8 ++++++-- coderd/database/queries/users.sql | 12 ++++++++---- enterprise/coderd/prebuilds/membership.go | 10 +++++----- enterprise/coderd/workspacequota_test.go | 3 +-- 5 files changed, 35 insertions(+), 16 deletions(-) diff --git a/coderd/database/queries/groupmembers.sql b/coderd/database/queries/groupmembers.sql index 2296e871939dc..7de8dbe4e4523 100644 --- a/coderd/database/queries/groupmembers.sql +++ b/coderd/database/queries/groupmembers.sql @@ -1,13 +1,21 @@ -- name: GetGroupMembers :many SELECT * FROM group_members_expanded -WHERE (@include_system::bool OR NOT user_is_system); +WHERE CASE + WHEN @include_system::bool THEN TRUE + ELSE + user_is_system = false + END; -- name: GetGroupMembersByGroupID :many SELECT * FROM group_members_expanded WHERE group_id = @group_id -- Filter by system type - AND (@include_system::bool OR NOT user_is_system); + AND CASE + WHEN @include_system::bool THEN TRUE + ELSE + user_is_system = false + END; -- name: GetGroupMembersCountByGroupID :one -- Returns the total count of members in a group. Shows the total @@ -17,7 +25,11 @@ SELECT COUNT(*) FROM group_members_expanded WHERE group_id = @group_id -- Filter by system type - AND (@include_system::bool OR NOT user_is_system); + AND CASE + WHEN @include_system::bool THEN TRUE + ELSE + user_is_system = false + END; -- InsertUserGroupsByName adds a user to all provided groups, if they exist. -- name: InsertUserGroupsByName :exec diff --git a/coderd/database/queries/organizationmembers.sql b/coderd/database/queries/organizationmembers.sql index 688a6ddd69799..1c0af011776e3 100644 --- a/coderd/database/queries/organizationmembers.sql +++ b/coderd/database/queries/organizationmembers.sql @@ -24,7 +24,11 @@ WHERE ELSE true END -- Filter by system type - AND (@include_system::bool OR NOT is_system); + AND CASE + WHEN @include_system::bool THEN TRUE + ELSE + is_system = false + END; -- name: InsertOrganizationMember :one INSERT INTO @@ -86,7 +90,7 @@ WHERE ELSE true END -- Filter by system type - AND (@include_system::bool OR NOT is_system) + AND CASE WHEN @include_system::bool THEN TRUE ELSE is_system = false END ORDER BY -- Deterministic and consistent ordering of all users. This is to ensure consistent pagination. LOWER(username) ASC OFFSET @offset_opt diff --git a/coderd/database/queries/users.sql b/coderd/database/queries/users.sql index 033a1abf8ed98..eece2f96512ea 100644 --- a/coderd/database/queries/users.sql +++ b/coderd/database/queries/users.sql @@ -49,7 +49,7 @@ FROM users WHERE deleted = false - AND (@include_system::bool OR NOT is_system); + AND CASE WHEN @include_system::bool THEN TRUE ELSE is_system = false END; -- name: GetActiveUserCount :one SELECT @@ -58,7 +58,7 @@ FROM users WHERE status = 'active'::user_status AND deleted = false - AND (@include_system::bool OR NOT is_system); + AND CASE WHEN @include_system::bool THEN TRUE ELSE is_system = false END; -- name: InsertUser :one INSERT INTO @@ -250,7 +250,11 @@ WHERE created_at >= @created_after ELSE true END - AND (@include_system::bool OR NOT is_system) + AND CASE + WHEN @include_system::bool THEN TRUE + ELSE + is_system = false + END AND CASE WHEN @github_com_user_id :: bigint != 0 THEN github_com_user_id = @github_com_user_id @@ -360,7 +364,7 @@ RETURNING id, email, username, last_seen_at; -- AllUserIDs returns all UserIDs regardless of user status or deletion. -- name: AllUserIDs :many SELECT DISTINCT id FROM USERS - WHERE (@include_system::bool OR NOT is_system); + WHERE CASE WHEN @include_system::bool THEN TRUE ELSE is_system = false END; -- name: UpdateUserHashedOneTimePasscode :exec UPDATE diff --git a/enterprise/coderd/prebuilds/membership.go b/enterprise/coderd/prebuilds/membership.go index 9b41d68787ea1..4d062f1cd029a 100644 --- a/enterprise/coderd/prebuilds/membership.go +++ b/enterprise/coderd/prebuilds/membership.go @@ -44,23 +44,23 @@ func (s StoreMembershipReconciler) ReconcileAll(ctx context.Context, userID uuid return xerrors.Errorf("determine prebuild organization membership: %w", err) } - orgMemberShips := make(map[uuid.UUID]struct{}, 0) + orgMemberships := make(map[uuid.UUID]struct{}, 0) defaultOrg, err := s.store.GetDefaultOrganization(ctx) if err != nil { return xerrors.Errorf("get default organization: %w", err) } - orgMemberShips[defaultOrg.ID] = struct{}{} + orgMemberships[defaultOrg.ID] = struct{}{} for _, o := range organizationMemberships { - orgMemberShips[o.ID] = struct{}{} + orgMemberships[o.ID] = struct{}{} } var membershipInsertionErrors error for _, preset := range presets { - _, alreadyOrgMember := orgMemberShips[preset.OrganizationID] + _, alreadyOrgMember := orgMemberships[preset.OrganizationID] if !alreadyOrgMember { // Add the organization to our list of memberships regardless of potential failure below // to avoid a retry that will probably be doomed anyway. - orgMemberShips[preset.OrganizationID] = struct{}{} + orgMemberships[preset.OrganizationID] = struct{}{} // Insert the missing membership _, err = s.store.InsertOrganizationMember(ctx, database.InsertOrganizationMemberParams{ diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index b40b9b7c8abfd..c6a891b6ce12b 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -481,11 +481,10 @@ func TestWorkspaceQuota(t *testing.T) { coderdtest.AwaitTemplateVersionJobCompleted(t, client, versionZeroCost.ID) templateZeroCost := coderdtest.CreateTemplate(t, client, user.OrganizationID, versionZeroCost.ID) - // Even with zero cost, should pass + // Workspace with zero cost should pass workspaceZeroCost := coderdtest.CreateWorkspace(t, client, templateZeroCost.ID) buildZeroCost := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspaceZeroCost.LatestBuild.ID) - // Verify the build failed due to quota require.Equal(t, codersdk.WorkspaceStatusRunning, buildZeroCost.Status) require.Empty(t, buildZeroCost.Job.Error) From 92760d18511e787c3e1df6ecc9fbbe606e573e66 Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Mon, 4 Aug 2025 13:02:31 +0000 Subject: [PATCH 9/9] make gen --- coderd/database/dbauthz/dbauthz.go | 2 + coderd/database/queries.sql.go | 38 ++++++++++++++----- enterprise/coderd/prebuilds/membership.go | 24 ++++++++---- .../coderd/prebuilds/membership_test.go | 10 ++--- 4 files changed, 53 insertions(+), 21 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index dfd4d7982b0d7..cb41df3fce994 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -485,6 +485,8 @@ var ( rbac.ResourceFile.Type: { policy.ActionRead, }, + // Needs to be able to add the prebuilds system user to the "prebuilds" group in each organization that needs prebuilt workspaces + // so that prebuilt workspaces can be scheduled and owned in those organizations. rbac.ResourceGroup.Type: { policy.ActionRead, policy.ActionCreate, diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index a091c82955bba..d590630773e55 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2248,7 +2248,11 @@ func (q *sqlQuerier) DeleteGroupMemberFromGroup(ctx context.Context, arg DeleteG const getGroupMembers = `-- name: GetGroupMembers :many SELECT user_id, user_email, user_username, user_hashed_password, user_created_at, user_updated_at, user_status, user_rbac_roles, user_login_type, user_avatar_url, user_deleted, user_last_seen_at, user_quiet_hours_schedule, user_name, user_github_com_user_id, user_is_system, organization_id, group_name, group_id FROM group_members_expanded -WHERE ($1::bool OR NOT user_is_system) +WHERE CASE + WHEN $1::bool THEN TRUE + ELSE + user_is_system = false + END ` func (q *sqlQuerier) GetGroupMembers(ctx context.Context, includeSystem bool) ([]GroupMember, error) { @@ -2299,7 +2303,11 @@ SELECT user_id, user_email, user_username, user_hashed_password, user_created_at FROM group_members_expanded WHERE group_id = $1 -- Filter by system type - AND ($2::bool OR NOT user_is_system) + AND CASE + WHEN $2::bool THEN TRUE + ELSE + user_is_system = false + END ` type GetGroupMembersByGroupIDParams struct { @@ -2355,7 +2363,11 @@ SELECT COUNT(*) FROM group_members_expanded WHERE group_id = $1 -- Filter by system type - AND ($2::bool OR NOT user_is_system) + AND CASE + WHEN $2::bool THEN TRUE + ELSE + user_is_system = false + END ` type GetGroupMembersCountByGroupIDParams struct { @@ -6473,7 +6485,11 @@ WHERE ELSE true END -- Filter by system type - AND ($3::bool OR NOT is_system) + AND CASE + WHEN $3::bool THEN TRUE + ELSE + is_system = false + END ` type OrganizationMembersParams struct { @@ -6546,7 +6562,7 @@ WHERE ELSE true END -- Filter by system type - AND ($2::bool OR NOT is_system) + AND CASE WHEN $2::bool THEN TRUE ELSE is_system = false END ORDER BY -- Deterministic and consistent ordering of all users. This is to ensure consistent pagination. LOWER(username) ASC OFFSET $3 @@ -13696,7 +13712,7 @@ func (q *sqlQuerier) UpdateUserLinkedID(ctx context.Context, arg UpdateUserLinke const allUserIDs = `-- name: AllUserIDs :many SELECT DISTINCT id FROM USERS - WHERE ($1::bool OR NOT is_system) + WHERE CASE WHEN $1::bool THEN TRUE ELSE is_system = false END ` // AllUserIDs returns all UserIDs regardless of user status or deletion. @@ -13730,7 +13746,7 @@ FROM users WHERE status = 'active'::user_status AND deleted = false - AND ($1::bool OR NOT is_system) + AND CASE WHEN $1::bool THEN TRUE ELSE is_system = false END ` func (q *sqlQuerier) GetActiveUserCount(ctx context.Context, includeSystem bool) (int64, error) { @@ -13894,7 +13910,7 @@ FROM users WHERE deleted = false - AND ($1::bool OR NOT is_system) + AND CASE WHEN $1::bool THEN TRUE ELSE is_system = false END ` func (q *sqlQuerier) GetUserCount(ctx context.Context, includeSystem bool) (int64, error) { @@ -14011,7 +14027,11 @@ WHERE created_at >= $8 ELSE true END - AND ($9::bool OR NOT is_system) + AND CASE + WHEN $9::bool THEN TRUE + ELSE + is_system = false + END AND CASE WHEN $10 :: bigint != 0 THEN github_com_user_id = $10 diff --git a/enterprise/coderd/prebuilds/membership.go b/enterprise/coderd/prebuilds/membership.go index 4d062f1cd029a..03328c2012534 100644 --- a/enterprise/coderd/prebuilds/membership.go +++ b/enterprise/coderd/prebuilds/membership.go @@ -12,6 +12,11 @@ import ( "github.com/coder/quartz" ) +const ( + PrebuiltWorkspacesGroupName = "coder_prebuilt_workspaces" + PrebuiltWorkspacesGroupDisplayName = "Prebuilt Workspaces" +) + // StoreMembershipReconciler encapsulates the responsibility of ensuring that the prebuilds system user is a member of all // organizations for which prebuilt workspaces are requested. This is necessary because our data model requires that such // prebuilt workspaces belong to a member of the organization of their eventual claimant. @@ -27,11 +32,16 @@ func NewStoreMembershipReconciler(store database.Store, clock quartz.Clock) Stor } } -// ReconcileAll compares the current membership of a user to the membership required in order to create prebuilt workspaces. -// If the user in question is not yet a member of an organization that needs prebuilt workspaces, ReconcileAll will create -// the membership required. +// ReconcileAll compares the current organization and group memberships of a user to the memberships required +// in order to create prebuilt workspaces. If the user in question is not yet a member of an organization that +// needs prebuilt workspaces, ReconcileAll will create the membership required. +// +// To facilitate quota management, ReconcileAll will ensure: +// * the existence of a group (defined by PrebuiltWorkspacesGroupName) in each organization that needs prebuilt workspaces +// * that the prebuilds system user belongs to the group in each organization that needs prebuilt workspaces +// * that the group has a quota of 0 by default, which users can adjust based on their needs. // -// This method does not have an opinion on transaction or lock management. These responsibilities are left to the caller. +// ReconcileAll does not have an opinion on transaction or lock management. These responsibilities are left to the caller. func (s StoreMembershipReconciler) ReconcileAll(ctx context.Context, userID uuid.UUID, presets []database.GetTemplatePresetsWithPrebuildsRow) error { organizationMemberships, err := s.store.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{ UserID: userID, @@ -80,8 +90,8 @@ func (s StoreMembershipReconciler) ReconcileAll(ctx context.Context, userID uuid // This group will have a quota of 0 by default, which users can adjust based on their needs prebuildsGroup, err := s.store.InsertGroup(ctx, database.InsertGroupParams{ ID: uuid.New(), - Name: "prebuilds", - DisplayName: "Prebuilds", + Name: PrebuiltWorkspacesGroupName, + DisplayName: PrebuiltWorkspacesGroupDisplayName, OrganizationID: preset.OrganizationID, AvatarURL: "", QuotaAllowance: 0, // Default quota of 0, users should set this based on their needs @@ -94,7 +104,7 @@ func (s StoreMembershipReconciler) ReconcileAll(ctx context.Context, userID uuid } prebuildsGroup, err = s.store.GetGroupByOrgAndName(ctx, database.GetGroupByOrgAndNameParams{ OrganizationID: preset.OrganizationID, - Name: "prebuilds", + Name: PrebuiltWorkspacesGroupName, }) if err != nil { membershipInsertionErrors = errors.Join(membershipInsertionErrors, xerrors.Errorf("get existing prebuilds group: %w", err)) diff --git a/enterprise/coderd/prebuilds/membership_test.go b/enterprise/coderd/prebuilds/membership_test.go index 61928d0f1e3b1..34742b28b1f66 100644 --- a/enterprise/coderd/prebuilds/membership_test.go +++ b/enterprise/coderd/prebuilds/membership_test.go @@ -111,8 +111,8 @@ func TestReconcileAll(t *testing.T) { var prebuildsGroup database.Group if preExistingGroup { prebuildsGroup = dbgen.Group(t, db, database.Group{ - Name: "prebuilds", - DisplayName: "Prebuilds", + Name: prebuilds.PrebuiltWorkspacesGroupName, + DisplayName: prebuilds.PrebuiltWorkspacesGroupDisplayName, OrganizationID: targetOrg.ID, QuotaAllowance: 0, }) @@ -160,12 +160,12 @@ func TestReconcileAll(t *testing.T) { // Verify prebuilds group behavior based on expected outcomes prebuildsGroup, err = db.GetGroupByOrgAndName(ctx, database.GetGroupByOrgAndNameParams{ OrganizationID: targetOrg.ID, - Name: "prebuilds", + Name: prebuilds.PrebuiltWorkspacesGroupName, }) if tc.expectGroupExists != nil && *tc.expectGroupExists { require.NoError(t, err) - require.Equal(t, "prebuilds", prebuildsGroup.Name) - require.Equal(t, "Prebuilds", prebuildsGroup.DisplayName) + require.Equal(t, prebuilds.PrebuiltWorkspacesGroupName, prebuildsGroup.Name) + require.Equal(t, prebuilds.PrebuiltWorkspacesGroupDisplayName, prebuildsGroup.DisplayName) require.Equal(t, int32(0), prebuildsGroup.QuotaAllowance) // Default quota should be 0 if tc.expectUserInGroup != nil && *tc.expectUserInGroup {