work on batching removal by name or id

coder · Emyrk · Sep 11, 2024 · Sep 3, 2024 · Sep 3, 2024 · Sep 3, 2024
commit 2a1769c7fdcd34b5c7a82f335d7f1e4f05b696ce
diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -682,6 +682,17 @@ func (q *FakeQuerier) getWorkspaceResourcesByJobIDNoLock(_ context.Context, jobI
 	return resources, nil
 }
 
+func (q *FakeQuerier) getGroupByNameNoLock(arg database.NameOrganizationPair) (database.Group, error) {
+	for _, group := range q.groups {
+		if group.OrganizationID == arg.OrganizationID &&
+			group.Name == arg.Name {
+			return group, nil
+		}
+	}
+
+	return database.Group{}, sql.ErrNoRows
+}
+
 func (q *FakeQuerier) getGroupByIDNoLock(_ context.Context, id uuid.UUID) (database.Group, error) {
 	for _, group := range q.groups {
 		if group.ID == id {
@@ -2613,14 +2624,10 @@ func (q *FakeQuerier) GetGroupByOrgAndName(_ context.Context, arg database.GetGr
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()
 
-	for _, group := range q.groups {
-		if group.OrganizationID == arg.OrganizationID &&
-			group.Name == arg.Name {
-			return group, nil
-		}
-	}
-
-	return database.Group{}, sql.ErrNoRows
+	return q.getGroupByNameNoLock(database.NameOrganizationPair{
+		Name:           arg.Name,
+		OrganizationID: arg.OrganizationID,
+	})
 }
 
 func (q *FakeQuerier) GetGroupMembers(ctx context.Context) ([]database.GroupMember, error) {
@@ -7648,14 +7655,24 @@ func (q *FakeQuerier) RemoveUserFromGroups(_ context.Context, arg database.Remov
 
 	removed := make([]uuid.UUID, 0)
 	q.data.groupMembers = slices.DeleteFunc(q.data.groupMembers, func(groupMember database.GroupMemberTable) bool {
+		// Delete all group members that match the arguments.
 		if groupMember.UserID != arg.UserID {
+			// Not the right user, ignore.
 			return false
 		}
-		if !slices.Contains(arg.GroupIds, groupMember.GroupID) {
-			return false
+
+		matchesByID := slices.Contains(arg.GroupIds, groupMember.GroupID)
+		matchesByName := slices.ContainsFunc(arg.GroupNames, func(name database.NameOrganizationPair) bool {
+			_, err := q.getGroupByNameNoLock(name)
+			return err == nil
+		})
+
+		if matchesByName || matchesByID {
+			removed = append(removed, groupMember.GroupID)
+			return true
 		}
-		removed = append(removed, groupMember.GroupID)
-		return true
+
+		return false
 	})
 
 	return removed, nil

diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/groupmembers.sql b/coderd/database/queries/groupmembers.sql
@@ -57,9 +57,19 @@ WHERE
 -- name: RemoveUserFromGroups :many
 DELETE FROM
 	group_members
+	USING groups
 WHERE
+	group_members.group_id = groups.id AND
 	user_id = @user_id AND
-	group_id = ANY(@group_ids :: uuid [])
+	(
+		CASE WHEN array_length(@group_names :: name_organization_pair[], 1) > 0  THEN
+			 -- Using 'coalesce' to avoid troubles with null literals being an empty string.
+			 (groups.name, coalesce(groups.organization_id, '00000000-0000-0000-0000-000000000000' ::uuid)) = ANY (@group_names::name_organization_pair[])
+		 ELSE false
+		END
+		OR
+		group_id = ANY (@group_ids :: uuid[])
+	)
 RETURNING group_id;
 
 -- name: InsertGroupMember :exec

diff --git a/coderd/idpsync/group.go b/coderd/idpsync/group.go
@@ -91,8 +91,8 @@ func (s AGPLIDPSync) SyncGroups(ctx context.Context, db database.Store, user dat
 		}
 
 		// collect all diffs to do 1 sql update for all orgs
-		groupsToAdd := make([]uuid.UUID, 0)
-		groupsToRemove := make([]uuid.UUID, 0)
+		groupIDsToAdd := make([]uuid.UUID, 0)
+		groupsToRemove := make([]ExpectedGroup, 0)
 		// For each org, determine which groups the user should land in
 		for orgID, settings := range orgSettings {
 			if settings.GroupField == "" {
@@ -112,7 +112,8 @@ func (s AGPLIDPSync) SyncGroups(ctx context.Context, db database.Store, user dat
 			}
 			// Everyone group is always implied.
 			expectedGroups = append(expectedGroups, ExpectedGroup{
-				GroupID: &orgID,
+				OrganizationID: orgID,
+				GroupID:        &orgID,
 			})
 
 			// Now we know what groups the user should be in for a given org,
@@ -121,8 +122,9 @@ func (s AGPLIDPSync) SyncGroups(ctx context.Context, db database.Store, user dat
 			existingGroups := userOrgs[orgID]
 			existingGroupsTyped := db2sdk.List(existingGroups, func(f database.GetGroupsRow) ExpectedGroup {
 				return ExpectedGroup{
-					GroupID:   &f.Group.ID,
-					GroupName: &f.Group.Name,
+					OrganizationID: orgID,
+					GroupID:        &f.Group.ID,
+					GroupName:      &f.Group.Name,
 				}
 			})
 			add, remove := slice.SymmetricDifferenceFunc(existingGroupsTyped, expectedGroups, func(a, b ExpectedGroup) bool {
@@ -144,52 +146,75 @@ func (s AGPLIDPSync) SyncGroups(ctx context.Context, db database.Store, user dat
 				return xerrors.Errorf("handle missing groups: %w", err)
 			}
 
-			for _, removeGroup := range remove {
-				// This should always be the case.
-				// TODO: make sure this is always the case
-				if removeGroup.GroupID != nil {
-					groupsToRemove = append(groupsToRemove, *removeGroup.GroupID)
-				}
-			}
+			groupsToRemove = append(groupsToRemove, remove...)
+			groupIDsToAdd = append(groupIDsToAdd, assignGroups...)
+		}
 
-			groupsToAdd = append(groupsToAdd, assignGroups...)
+		err = s.applyGroupDifference(ctx, tx, user, groupIDsToAdd, groupsToRemove)
+		if err != nil {
+			return xerrors.Errorf("apply group difference: %w", err)
 		}
 
-		assignedGroupIDs, err := tx.InsertUserGroupsByID(ctx, database.InsertUserGroupsByIDParams{
-			UserID:   user.ID,
-			GroupIds: groupsToAdd,
+		return nil
+	}, nil)
+
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (s AGPLIDPSync) applyGroupDifference(ctx context.Context, tx database.Store, user database.User, add []uuid.UUID, remove []ExpectedGroup) error {
+	// Always do group removal before group add. This way if there is an error,
+	// we error on the underprivileged side.
+	removeIDs := make([]uuid.UUID, 0)
+	removeNames := make([]database.NameOrganizationPair, 0)
+	for _, r := range remove {
+		if r.GroupID != nil {
+			removeIDs = append(removeIDs, *r.GroupID)
+		} else if r.GroupName != nil {
+			removeNames = append(removeNames, database.NameOrganizationPair{
+				Name:           *r.GroupName,
+				OrganizationID: r.OrganizationID,
+			})
+		}
+	}
+
+	// If there is something to remove, do it.
+	if len(removeIDs) > 0 || len(removeNames) > 0 {
+		removedGroupIDs, err := tx.RemoveUserFromGroups(ctx, database.RemoveUserFromGroupsParams{
+			UserID:     user.ID,
+			GroupNames: removeNames,
+			GroupIds:   removeIDs,
 		})
 		if err != nil {
-			return xerrors.Errorf("insert user into %d groups: %w", len(groupsToAdd), err)
+			return xerrors.Errorf("remove user from %d groups: %w", len(removeIDs), err)
 		}
-		if len(assignedGroupIDs) != len(groupsToAdd) {
-			s.Logger.Debug(ctx, "failed to assign all groups to user",
+		if len(removedGroupIDs) != len(removeIDs) {
+			s.Logger.Debug(ctx, "failed to remove user from all groups",
 				slog.F("user_id", user.ID),
-				slog.F("groups_assigned_count", len(assignedGroupIDs)),
-				slog.F("expected_count", len(groupsToAdd)),
+				slog.F("groups_removed_count", len(removedGroupIDs)),
+				slog.F("expected_count", len(removeIDs)),
 			)
 		}
+	}
 
-		removedGroupIDs, err := tx.RemoveUserFromGroups(ctx, database.RemoveUserFromGroupsParams{
+	if len(add) > 0 {
+		assignedGroupIDs, err := tx.InsertUserGroupsByID(ctx, database.InsertUserGroupsByIDParams{
 			UserID:   user.ID,
-			GroupIds: groupsToRemove,
+			GroupIds: add,
 		})
 		if err != nil {
-			return xerrors.Errorf("remove user from %d groups: %w", len(groupsToRemove), err)
+			return xerrors.Errorf("insert user into %d groups: %w", len(add), err)
 		}
-		if len(removedGroupIDs) != len(groupsToRemove) {
-			s.Logger.Debug(ctx, "failed to remove user from all groups",
+		if len(assignedGroupIDs) != len(add) {
+			s.Logger.Debug(ctx, "failed to assign all groups to user",
 				slog.F("user_id", user.ID),
-				slog.F("groups_removed_count", len(removedGroupIDs)),
-				slog.F("expected_count", len(groupsToRemove)),
+				slog.F("groups_assigned_count", len(assignedGroupIDs)),
+				slog.F("expected_count", len(add)),
 			)
 		}
-
-		return nil
-	}, nil)
-
-	if err != nil {
-		return err
 	}
 
 	return nil
@@ -226,8 +251,9 @@ func (s *GroupSyncSettings) Type() string {
 }
 
 type ExpectedGroup struct {
-	GroupID   *uuid.UUID
-	GroupName *string
+	OrganizationID uuid.UUID
+	GroupID        *uuid.UUID
+	GroupName      *string
 }
 
 // ParseClaims will take the merged claims from the IDP and return the groups
@@ -280,20 +306,28 @@ func (s GroupSyncSettings) ParseClaims(mergedClaims jwt.MapClaims) ([]ExpectedGr
 	return groups, nil
 }
 
+// HandleMissingGroups ensures all ExpectedGroups convert to uuids.
+// Groups can be referenced by name via legacy params or IDP group names.
+// These group names are converted to IDs for easier assignment.
+// Missing groups are created if AutoCreate is enabled.
+// TODO: Batching this would be better, as this is 1 or 2 db calls per organization.
 func (s GroupSyncSettings) HandleMissingGroups(ctx context.Context, tx database.Store, orgID uuid.UUID, add []ExpectedGroup) ([]uuid.UUID, error) {
 	if !s.AutoCreateMissingGroups {
-		// construct the list of groups to search by name to see if they exist.
+		// If we are not creating groups, then just construct a db lookup for
+		// all groups by name.
 		var lookups []string
 		filter := make([]uuid.UUID, 0)
 		for _, expected := range add {
 			if expected.GroupID != nil {
+				// Groups with IDs are easy!
 				filter = append(filter, *expected.GroupID)
 			} else if expected.GroupName != nil {
 				lookups = append(lookups, *expected.GroupName)
 			}
 		}
 
 		if len(lookups) > 0 {
+			// Do name lookups for all groups that are missing IDs.
 			newGroups, err := tx.GetGroups(ctx, database.GetGroupsParams{
 				OrganizationID: uuid.UUID{},
 				HasMemberID:    uuid.UUID{},

diff --git a/coderd/idpsync/group_test.go b/coderd/idpsync/group_test.go
@@ -208,9 +208,12 @@ func TestGroupSyncTable(t *testing.T) {
 				},
 				AutoCreateMissingGroups: true,
 			},
-			Groups: map[uuid.UUID]bool{},
+			Groups: map[uuid.UUID]bool{
+				ids.ID("lg-foo"): true,
+			},
 			GroupNames: map[string]bool{
 				"legacy-foo": false,
+				"extra":      true,
 			},
 			ExpectedGroupNames: []string{
 				"legacy-bar",