diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 2081fbf5a8..95937a77d2 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2648,6 +2648,13 @@ func (q *querier) GetOrganizationsByUserID(ctx context.Context, userID database. return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetOrganizationsByUserID)(ctx, userID) } +func (q *querier) GetOrganizationsWithPrebuildStatus(ctx context.Context, arg database.GetOrganizationsWithPrebuildStatusParams) ([]database.GetOrganizationsWithPrebuildStatusRow, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceOrganization.All()); err != nil { + return nil, err + } + return q.db.GetOrganizationsWithPrebuildStatus(ctx, arg) +} + func (q *querier) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]database.ParameterSchema, error) { version, err := q.db.GetTemplateVersionByJobID(ctx, jobID) if err != nil { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 63226271f7..32c951fb5c 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -3759,6 +3759,14 @@ func (s *MethodTestSuite) TestPrebuilds() { dbm.EXPECT().GetPrebuildMetrics(gomock.Any()).Return([]database.GetPrebuildMetricsRow{}, nil).AnyTimes() check.Args().Asserts(rbac.ResourceWorkspace.All(), policy.ActionRead) })) + s.Run("GetOrganizationsWithPrebuildStatus", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + arg := database.GetOrganizationsWithPrebuildStatusParams{ + UserID: uuid.New(), + GroupName: "test", + } + dbm.EXPECT().GetOrganizationsWithPrebuildStatus(gomock.Any(), arg).Return([]database.GetOrganizationsWithPrebuildStatusRow{}, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceOrganization.All(), policy.ActionRead) + })) s.Run("GetPrebuildsSettings", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { dbm.EXPECT().GetPrebuildsSettings(gomock.Any()).Return("{}", nil).AnyTimes() check.Args().Asserts() diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 1b36a3fa98..252f6f9b5a 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -1243,6 +1243,13 @@ func (m queryMetricsStore) GetOrganizationsByUserID(ctx context.Context, userID return organizations, err } +func (m queryMetricsStore) GetOrganizationsWithPrebuildStatus(ctx context.Context, arg database.GetOrganizationsWithPrebuildStatusParams) ([]database.GetOrganizationsWithPrebuildStatusRow, error) { + start := time.Now() + r0, r1 := m.s.GetOrganizationsWithPrebuildStatus(ctx, arg) + m.queryLatencies.WithLabelValues("GetOrganizationsWithPrebuildStatus").Observe(time.Since(start).Seconds()) + return r0, r1 +} + func (m queryMetricsStore) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]database.ParameterSchema, error) { start := time.Now() schemas, err := m.s.GetParameterSchemasByJobID(ctx, jobID) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 0ecd0191cb..af89a987a3 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -2622,6 +2622,21 @@ func (mr *MockStoreMockRecorder) GetOrganizationsByUserID(ctx, arg any) *gomock. return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetOrganizationsByUserID", reflect.TypeOf((*MockStore)(nil).GetOrganizationsByUserID), ctx, arg) } +// GetOrganizationsWithPrebuildStatus mocks base method. +func (m *MockStore) GetOrganizationsWithPrebuildStatus(ctx context.Context, arg database.GetOrganizationsWithPrebuildStatusParams) ([]database.GetOrganizationsWithPrebuildStatusRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetOrganizationsWithPrebuildStatus", ctx, arg) + ret0, _ := ret[0].([]database.GetOrganizationsWithPrebuildStatusRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetOrganizationsWithPrebuildStatus indicates an expected call of GetOrganizationsWithPrebuildStatus. +func (mr *MockStoreMockRecorder) GetOrganizationsWithPrebuildStatus(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetOrganizationsWithPrebuildStatus", reflect.TypeOf((*MockStore)(nil).GetOrganizationsWithPrebuildStatus), ctx, arg) +} + // GetParameterSchemasByJobID mocks base method. func (m *MockStore) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]database.ParameterSchema, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 2aa1cb8650..b1a4509398 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -269,6 +269,9 @@ type sqlcQuerier interface { GetOrganizationResourceCountByID(ctx context.Context, organizationID uuid.UUID) (GetOrganizationResourceCountByIDRow, error) GetOrganizations(ctx context.Context, arg GetOrganizationsParams) ([]Organization, error) GetOrganizationsByUserID(ctx context.Context, arg GetOrganizationsByUserIDParams) ([]Organization, error) + // GetOrganizationsWithPrebuildStatus returns organizations with prebuilds configured and their + // membership status for the prebuilds system user (org membership, group existence, group membership). + GetOrganizationsWithPrebuildStatus(ctx context.Context, arg GetOrganizationsWithPrebuildStatusParams) ([]GetOrganizationsWithPrebuildStatusRow, error) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]ParameterSchema, error) GetPrebuildMetrics(ctx context.Context) ([]GetPrebuildMetricsRow, error) GetPrebuildsSettings(ctx context.Context) (string, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 2869316266..de5a4d0233 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -8285,6 +8285,93 @@ func (q *sqlQuerier) FindMatchingPresetID(ctx context.Context, arg FindMatchingP return template_version_preset_id, err } +const getOrganizationsWithPrebuildStatus = `-- name: GetOrganizationsWithPrebuildStatus :many +WITH orgs_with_prebuilds AS ( + -- Get unique organizations that have presets with prebuilds configured + SELECT DISTINCT o.id, o.name + FROM organizations o + INNER JOIN templates t ON t.organization_id = o.id + INNER JOIN template_versions tv ON tv.template_id = t.id + INNER JOIN template_version_presets tvp ON tvp.template_version_id = tv.id + WHERE tvp.desired_instances IS NOT NULL +), +prebuild_user_membership AS ( + -- Check if the user is a member of the organizations + SELECT om.organization_id + FROM organization_members om + INNER JOIN orgs_with_prebuilds owp ON owp.id = om.organization_id + WHERE om.user_id = $1::uuid +), +prebuild_groups AS ( + -- Check if the organizations have the prebuilds group + SELECT g.organization_id, g.id as group_id + FROM groups g + INNER JOIN orgs_with_prebuilds owp ON owp.id = g.organization_id + WHERE g.name = $2::text +), +prebuild_group_membership AS ( + -- Check if the user is in the prebuilds group + SELECT pg.organization_id + FROM prebuild_groups pg + INNER JOIN group_members gm ON gm.group_id = pg.group_id + WHERE gm.user_id = $1::uuid +) +SELECT + owp.id AS organization_id, + owp.name AS organization_name, + (pum.organization_id IS NOT NULL)::boolean AS has_prebuild_user, + pg.group_id AS prebuilds_group_id, + (pgm.organization_id IS NOT NULL)::boolean AS has_prebuild_user_in_group +FROM orgs_with_prebuilds owp +LEFT JOIN prebuild_groups pg ON pg.organization_id = owp.id +LEFT JOIN prebuild_user_membership pum ON pum.organization_id = owp.id +LEFT JOIN prebuild_group_membership pgm ON pgm.organization_id = owp.id +` + +type GetOrganizationsWithPrebuildStatusParams struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + GroupName string `db:"group_name" json:"group_name"` +} + +type GetOrganizationsWithPrebuildStatusRow struct { + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + OrganizationName string `db:"organization_name" json:"organization_name"` + HasPrebuildUser bool `db:"has_prebuild_user" json:"has_prebuild_user"` + PrebuildsGroupID uuid.NullUUID `db:"prebuilds_group_id" json:"prebuilds_group_id"` + HasPrebuildUserInGroup bool `db:"has_prebuild_user_in_group" json:"has_prebuild_user_in_group"` +} + +// GetOrganizationsWithPrebuildStatus returns organizations with prebuilds configured and their +// membership status for the prebuilds system user (org membership, group existence, group membership). +func (q *sqlQuerier) GetOrganizationsWithPrebuildStatus(ctx context.Context, arg GetOrganizationsWithPrebuildStatusParams) ([]GetOrganizationsWithPrebuildStatusRow, error) { + rows, err := q.db.QueryContext(ctx, getOrganizationsWithPrebuildStatus, arg.UserID, arg.GroupName) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetOrganizationsWithPrebuildStatusRow + for rows.Next() { + var i GetOrganizationsWithPrebuildStatusRow + if err := rows.Scan( + &i.OrganizationID, + &i.OrganizationName, + &i.HasPrebuildUser, + &i.PrebuildsGroupID, + &i.HasPrebuildUserInGroup, + ); 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 getPrebuildMetrics = `-- name: GetPrebuildMetrics :many SELECT t.name as template_name, diff --git a/coderd/database/queries/prebuilds.sql b/coderd/database/queries/prebuilds.sql index 7c060971ef..ae70593b26 100644 --- a/coderd/database/queries/prebuilds.sql +++ b/coderd/database/queries/prebuilds.sql @@ -327,3 +327,47 @@ SET FROM jobs_to_cancel WHERE provisioner_jobs.id = jobs_to_cancel.id RETURNING jobs_to_cancel.id, jobs_to_cancel.workspace_id, jobs_to_cancel.template_id, jobs_to_cancel.template_version_preset_id; + +-- name: GetOrganizationsWithPrebuildStatus :many +-- GetOrganizationsWithPrebuildStatus returns organizations with prebuilds configured and their +-- membership status for the prebuilds system user (org membership, group existence, group membership). +WITH orgs_with_prebuilds AS ( + -- Get unique organizations that have presets with prebuilds configured + SELECT DISTINCT o.id, o.name + FROM organizations o + INNER JOIN templates t ON t.organization_id = o.id + INNER JOIN template_versions tv ON tv.template_id = t.id + INNER JOIN template_version_presets tvp ON tvp.template_version_id = tv.id + WHERE tvp.desired_instances IS NOT NULL +), +prebuild_user_membership AS ( + -- Check if the user is a member of the organizations + SELECT om.organization_id + FROM organization_members om + INNER JOIN orgs_with_prebuilds owp ON owp.id = om.organization_id + WHERE om.user_id = @user_id::uuid +), +prebuild_groups AS ( + -- Check if the organizations have the prebuilds group + SELECT g.organization_id, g.id as group_id + FROM groups g + INNER JOIN orgs_with_prebuilds owp ON owp.id = g.organization_id + WHERE g.name = @group_name::text +), +prebuild_group_membership AS ( + -- Check if the user is in the prebuilds group + SELECT pg.organization_id + FROM prebuild_groups pg + INNER JOIN group_members gm ON gm.group_id = pg.group_id + WHERE gm.user_id = @user_id::uuid +) +SELECT + owp.id AS organization_id, + owp.name AS organization_name, + (pum.organization_id IS NOT NULL)::boolean AS has_prebuild_user, + pg.group_id AS prebuilds_group_id, + (pgm.organization_id IS NOT NULL)::boolean AS has_prebuild_user_in_group +FROM orgs_with_prebuilds owp +LEFT JOIN prebuild_groups pg ON pg.organization_id = owp.id +LEFT JOIN prebuild_user_membership pum ON pum.organization_id = owp.id +LEFT JOIN prebuild_group_membership pgm ON pgm.organization_id = owp.id; diff --git a/enterprise/coderd/prebuilds/membership.go b/enterprise/coderd/prebuilds/membership.go index f843d33f7f..9436f68737 100644 --- a/enterprise/coderd/prebuilds/membership.go +++ b/enterprise/coderd/prebuilds/membership.go @@ -2,12 +2,13 @@ package prebuilds import ( "context" - "database/sql" "errors" "github.com/google/uuid" "golang.org/x/xerrors" + "cdr.dev/slog" + "github.com/coder/coder/v2/coderd/database" "github.com/coder/quartz" ) @@ -21,114 +22,117 @@ const ( // 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. type StoreMembershipReconciler struct { - store database.Store - clock quartz.Clock + store database.Store + clock quartz.Clock + logger slog.Logger } -func NewStoreMembershipReconciler(store database.Store, clock quartz.Clock) StoreMembershipReconciler { +func NewStoreMembershipReconciler(store database.Store, clock quartz.Clock, logger slog.Logger) StoreMembershipReconciler { return StoreMembershipReconciler{ - store: store, - clock: clock, + store: store, + clock: clock, + logger: logger, } } -// 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. +// ReconcileAll ensures the prebuilds system user has the necessary memberships to create prebuilt workspaces. +// For each organization with prebuilds configured, it ensures: +// * The user is a member of the organization +// * A group exists with quota 0 +// * The user is a member of that group // -// 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. +// Unique constraint violations are safely ignored (concurrent creation). // // 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, - Deleted: sql.NullBool{ - Bool: false, - Valid: true, - }, +func (s StoreMembershipReconciler) ReconcileAll(ctx context.Context, userID uuid.UUID, groupName string) error { + orgStatuses, err := s.store.GetOrganizationsWithPrebuildStatus(ctx, database.GetOrganizationsWithPrebuildStatusParams{ + UserID: userID, + GroupName: groupName, }) if err != nil { - return xerrors.Errorf("determine prebuild organization membership: %w", err) - } - - 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{}{} - for _, o := range organizationMemberships { - orgMemberships[o.ID] = struct{}{} + return xerrors.Errorf("get organizations with prebuild status: %w", err) } var membershipInsertionErrors error - for _, preset := range presets { - _, 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{}{} + for _, orgStatus := range orgStatuses { + s.logger.Debug(ctx, "organization prebuild status", + slog.F("organization_id", orgStatus.OrganizationID), + slog.F("organization_name", orgStatus.OrganizationName), + slog.F("has_prebuild_user", orgStatus.HasPrebuildUser), + slog.F("has_prebuild_group", orgStatus.PrebuildsGroupID.Valid), + slog.F("has_prebuild_user_in_group", orgStatus.HasPrebuildUserInGroup)) - // Insert the missing membership + // Add user to org if needed + if !orgStatus.HasPrebuildUser { _, err = s.store.InsertOrganizationMember(ctx, database.InsertOrganizationMemberParams{ - OrganizationID: preset.OrganizationID, + OrganizationID: orgStatus.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)) + // Unique violation means organization membership was created after status check, safe to ignore. + if err != nil && !database.IsUniqueViolation(err) { + membershipInsertionErrors = errors.Join(membershipInsertionErrors, err) continue } + if err == nil { + s.logger.Info(ctx, "added prebuilds user to organization", + slog.F("organization_id", orgStatus.OrganizationID), + slog.F("organization_name", orgStatus.OrganizationName), + slog.F("prebuilds_user", userID.String())) + } } - // determine whether the org already has a prebuilds group - prebuildsGroupExists := true - prebuildsGroup, err := s.store.GetGroupByOrgAndName(ctx, database.GetGroupByOrgAndNameParams{ - OrganizationID: preset.OrganizationID, - Name: PrebuiltWorkspacesGroupName, - }) - if err != nil { - if !xerrors.Is(err, sql.ErrNoRows) { - membershipInsertionErrors = errors.Join(membershipInsertionErrors, xerrors.Errorf("get prebuilds group: %w", err)) - continue - } - prebuildsGroupExists = false - } - - // if the prebuilds group does not exist, create it - if !prebuildsGroupExists { - // 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{ + // Create group if it doesn't exist + var groupID uuid.UUID + if !orgStatus.PrebuildsGroupID.Valid { + // Group doesn't exist, create it + group, err := s.store.InsertGroup(ctx, database.InsertGroupParams{ ID: uuid.New(), Name: PrebuiltWorkspacesGroupName, DisplayName: PrebuiltWorkspacesGroupDisplayName, - OrganizationID: preset.OrganizationID, + OrganizationID: orgStatus.OrganizationID, AvatarURL: "", - QuotaAllowance: 0, // Default quota of 0, users should set this based on their needs + QuotaAllowance: 0, }) - if err != nil { - membershipInsertionErrors = errors.Join(membershipInsertionErrors, xerrors.Errorf("create prebuilds group: %w", err)) + // Unique violation means group was created after status check, safe to ignore. + if err != nil && !database.IsUniqueViolation(err) { + membershipInsertionErrors = errors.Join(membershipInsertionErrors, err) continue } + if err == nil { + s.logger.Info(ctx, "created prebuilds group in organization", + slog.F("organization_id", orgStatus.OrganizationID), + slog.F("organization_name", orgStatus.OrganizationName), + slog.F("prebuilds_group", group.ID.String())) + } + groupID = group.ID + } else { + // Group exists + groupID = orgStatus.PrebuildsGroupID.UUID } - // add the system user to the prebuilds group - err = s.store.InsertGroupMember(ctx, database.InsertGroupMemberParams{ - GroupID: prebuildsGroup.ID, - UserID: userID, - }) - if err != nil { - // 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)) + // Add user to group if needed + if !orgStatus.HasPrebuildUserInGroup { + err = s.store.InsertGroupMember(ctx, database.InsertGroupMemberParams{ + GroupID: groupID, + UserID: userID, + }) + // Unique violation means group membership was created after status check, safe to ignore. + if err != nil && !database.IsUniqueViolation(err) { + membershipInsertionErrors = errors.Join(membershipInsertionErrors, err) + continue + } + if err == nil { + s.logger.Info(ctx, "added prebuilds user to prebuilds group", + slog.F("organization_id", orgStatus.OrganizationID), + slog.F("organization_name", orgStatus.OrganizationName), + slog.F("prebuilds_user", userID.String()), + slog.F("prebuilds_group", groupID.String())) } } } + return membershipInsertionErrors } diff --git a/enterprise/coderd/prebuilds/membership_test.go b/enterprise/coderd/prebuilds/membership_test.go index 55d6557b12..fe4ec26259 100644 --- a/enterprise/coderd/prebuilds/membership_test.go +++ b/enterprise/coderd/prebuilds/membership_test.go @@ -7,16 +7,17 @@ import ( "github.com/google/uuid" "github.com/stretchr/testify/require" - "tailscale.com/types/ptr" - "github.com/coder/quartz" + "cdr.dev/slog/sloggers/slogtest" "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/dbfake" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/enterprise/coderd/prebuilds" "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" ) // TestReconcileAll verifies that StoreMembershipReconciler correctly updates membership @@ -26,169 +27,178 @@ func TestReconcileAll(t *testing.T) { clock := quartz.NewMock(t) - // Helper to build a minimal Preset row belonging to a given org. - newPresetRow := func(orgID uuid.UUID) database.GetTemplatePresetsWithPrebuildsRow { - return database.GetTemplatePresetsWithPrebuildsRow{ - ID: uuid.New(), - OrganizationID: orgID, - } - } - tests := []struct { name string - includePreset []bool + includePreset bool preExistingOrgMembership []bool preExistingGroup []bool preExistingGroupMembership []bool // Expected outcomes - expectOrgMembershipExists *bool - expectGroupExists *bool - expectUserInGroup *bool + expectOrgMembershipExists bool + expectGroupExists bool + expectUserInGroup bool }{ { name: "if there are no presets, membership reconciliation is a no-op", - includePreset: []bool{false}, + includePreset: false, preExistingOrgMembership: []bool{true, false}, preExistingGroup: []bool{true, false}, preExistingGroupMembership: []bool{true, false}, - expectOrgMembershipExists: ptr.To(false), - expectGroupExists: ptr.To(false), + expectOrgMembershipExists: false, + expectGroupExists: false, + expectUserInGroup: false, }, { name: "if there is a preset, then we should enforce org and group membership in all cases", - includePreset: []bool{true}, + includePreset: 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), + expectOrgMembershipExists: true, + expectGroupExists: true, + expectUserInGroup: true, }, } for _, tc := range tests { 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() + includePreset := tc.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() - // 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) + // nolint:gocritic // Reconciliation happens as prebuilds system user, not a human user. + ctx := dbauthz.AsPrebuildsOrchestrator(testutil.Context(t, testutil.WaitLong)) + client, db := coderdtest.NewWithDatabase(t, nil) + owner := coderdtest.CreateFirstUser(t, client) - defaultOrg, err := db.GetDefaultOrganization(ctx) - require.NoError(t, err) + defaultOrg, err := db.GetDefaultOrganization(ctx) + require.NoError(t, err) - // introduce an unrelated organization to ensure that the membership reconciler doesn't interfere with it. - unrelatedOrg := dbgen.Organization(t, db, database.Organization{}) - targetOrg := dbgen.Organization(t, db, database.Organization{}) + // Introduce an unrelated organization to ensure that the membership reconciler doesn't interfere with it. + unrelatedOrg := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: unrelatedOrg.ID, UserID: database.PrebuildsSystemUserID}) - // Ensure membership to unrelated org. - dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: unrelatedOrg.ID, UserID: database.PrebuildsSystemUserID}) + // Organization to test + targetOrg := dbgen.Organization(t, db, database.Organization{}) - if preExistingOrgMembership { - // System user already a member of both orgs. - dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: targetOrg.ID, UserID: database.PrebuildsSystemUserID}) - } + // Prebuilds system user is a member of the organization + if preExistingOrgMembership { + 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.PrebuiltWorkspacesGroupName, - DisplayName: prebuilds.PrebuiltWorkspacesGroupDisplayName, - 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, + // Organization has the prebuilds group + var prebuildsGroup database.Group + if preExistingGroup { + prebuildsGroup = dbgen.Group(t, db, database.Group{ Name: prebuilds.PrebuiltWorkspacesGroupName, + DisplayName: prebuilds.PrebuiltWorkspacesGroupDisplayName, + OrganizationID: targetOrg.ID, + QuotaAllowance: 0, }) - if tc.expectGroupExists != nil && *tc.expectGroupExists { - require.NoError(t, err) - 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 { - // 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 no preset exists, then we do not enforce group membership: - 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) - } + // Add the system user to the group if required by test case + if preExistingGroupMembership { + dbgen.GroupMember(t, db, database.GroupMemberTable{ + GroupID: prebuildsGroup.ID, + UserID: database.PrebuildsSystemUserID, + }) } + } - 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)) - } + // Setup unrelated org preset + dbfake.TemplateVersion(t, db).Seed(database.TemplateVersion{ + OrganizationID: unrelatedOrg.ID, + CreatedBy: owner.UserID, + }).Preset(database.TemplateVersionPreset{ + DesiredInstances: sql.NullInt32{ + Int32: 1, + Valid: true, + }, + }).Do() + + // Setup target org preset + dbfake.TemplateVersion(t, db).Seed(database.TemplateVersion{ + OrganizationID: targetOrg.ID, + CreatedBy: owner.UserID, + }).Preset(database.TemplateVersionPreset{ + DesiredInstances: sql.NullInt32{ + Int32: 0, + Valid: includePreset, + }, + }).Do() + + // 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, slogtest.Make(t, nil)) + require.NoError(t, reconciler.ReconcileAll(ctx, database.PrebuildsSystemUserID, prebuilds.PrebuiltWorkspacesGroupName)) + + // Verify memberships after reconciliation. + postReconcileMemberships, err := db.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{ + UserID: database.PrebuildsSystemUserID, + }) + require.NoError(t, err) + expectedMembershipsAfter := expectedMembershipsBefore + if !preExistingOrgMembership && 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.PrebuiltWorkspacesGroupName, + }) + if tc.expectGroupExists { + require.NoError(t, err) + 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 { + // 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 no preset exists, then we do not enforce group membership: + if !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 { + // Verify that no prebuilds group exists + require.Error(t, err) + require.True(t, errors.Is(err, sql.ErrNoRows)) + } + }) } } } diff --git a/enterprise/coderd/prebuilds/reconcile.go b/enterprise/coderd/prebuilds/reconcile.go index 9af8fbc4cd..f280436ea9 100644 --- a/enterprise/coderd/prebuilds/reconcile.go +++ b/enterprise/coderd/prebuilds/reconcile.go @@ -322,6 +322,12 @@ func (c *StoreReconciler) ReconcileAll(ctx context.Context) (stats prebuilds.Rec return nil } + membershipReconciler := NewStoreMembershipReconciler(c.store, c.clock, logger) + err = membershipReconciler.ReconcileAll(ctx, database.PrebuildsSystemUserID, PrebuiltWorkspacesGroupName) + if err != nil { + return xerrors.Errorf("reconcile prebuild membership: %w", err) + } + snapshot, err := c.SnapshotState(ctx, c.store) if err != nil { return xerrors.Errorf("determine current snapshot: %w", err) @@ -334,12 +340,6 @@ func (c *StoreReconciler) ReconcileAll(ctx context.Context) (stats prebuilds.Rec return nil } - membershipReconciler := NewStoreMembershipReconciler(c.store, c.clock) - err = membershipReconciler.ReconcileAll(ctx, database.PrebuildsSystemUserID, snapshot.Presets) - if err != nil { - return xerrors.Errorf("reconcile prebuild membership: %w", err) - } - var eg errgroup.Group // Reconcile presets in parallel. Each preset in its own goroutine. for _, preset := range snapshot.Presets {