MM-43143: Dataloaders for roles (#19936)

We use a dataloader to optimize roles loading.

https://mattermost.atlassian.net/browse/MM-43143

```release-note
NONE
```

Co-authored-by: Mattermod <mattermod@users.noreply.github.com>

api4/graphql.go

@@ -9,6 +9,7 @@ import (
 	"encoding/json"
 	"net/http"
 
+	"github.com/graph-gophers/dataloader/v6"
 	graphql "github.com/graph-gophers/graphql-go"
 	gqlerrors "github.com/graph-gophers/graphql-go/errors"
 	"github.com/mattermost/mattermost-server/v6/model"
@@ -58,7 +59,12 @@ func (api *API) InitGraphQL() error {
 }
 
 // Unique type to hold our context.
-type ctxKey struct{}
+type ctxKey int
+
+const (
+	webCtx         ctxKey = 0
+	rolesLoaderCtx ctxKey = 1
+)
 
 func (api *API) graphQL(c *Context, w http.ResponseWriter, r *http.Request) {
 	var response *graphql.Response
@@ -90,7 +96,10 @@ func (api *API) graphQL(c *Context, w http.ResponseWriter, r *http.Request) {
 
 	// Populate the context with required info.
 	reqCtx := r.Context()
-	reqCtx = context.WithValue(reqCtx, ctxKey{}, c)
+	reqCtx = context.WithValue(reqCtx, webCtx, c)
+
+	rolesLoader := dataloader.NewBatchedLoader(graphQLRolesLoader, dataloader.WithBatchCapacity(200))
+	reqCtx = context.WithValue(reqCtx, rolesLoaderCtx, rolesLoader)
 
 	response = api.schema.Exec(reqCtx,
 		params.Query,

api4/resolver.go

@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 
+	"github.com/graph-gophers/dataloader/v6"
 	"github.com/mattermost/mattermost-server/v6/app"
 	"github.com/mattermost/mattermost-server/v6/model"
 	"github.com/mattermost/mattermost-server/v6/web"
@@ -278,9 +279,18 @@ func (*resolver) ChannelMembers(ctx context.Context, args struct {
 // Kind of an anti-pattern, but there are lots of methods attached to *web.Context
 // so we use it for now.
 func getCtx(ctx context.Context) (*web.Context, error) {
-	c, ok := ctx.Value(ctxKey{}).(*web.Context)
+	c, ok := ctx.Value(webCtx).(*web.Context)
 	if !ok {
 		return nil, errors.New("no web.Context found in context")
 	}
 	return c, nil
 }
+
+// getRolesLoader returns the roles loader out of the context.
+func getRolesLoader(ctx context.Context) (*dataloader.Loader, error) {
+	l, ok := ctx.Value(rolesLoaderCtx).(*dataloader.Loader)
+	if !ok {
+		return nil, errors.New("no dataloader.Loader found in context")
+	}
+	return l, nil
+}

api4/resolver_channel_member.go

@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/graph-gophers/dataloader/v6"
 	"github.com/mattermost/mattermost-server/v6/model"
 	"github.com/mattermost/mattermost-server/v6/web"
 )
@@ -65,12 +66,24 @@ func (cm *channelMember) Channel(ctx context.Context) (*channel, error) {
 }
 
 func (cm *channelMember) Roles_(ctx context.Context) ([]*model.Role, error) {
-	c, err := getCtx(ctx)
+	loader, err := getRolesLoader(ctx)
 	if err != nil {
 		return nil, err
 	}
 
-	return getGraphQLRoles(c, strings.Fields(cm.Roles))
+	thunk := loader.LoadMany(ctx, dataloader.NewKeysFromStrings(strings.Fields(cm.Roles)))
+	results, errs := thunk()
+	// All errors are the same. We just return the first one.
+	if len(errs) > 0 && errs[0] != nil {
+		return nil, err
+	}
+
+	roles := make([]*model.Role, len(results))
+	for i, res := range results {
+		roles[i] = res.(*model.Role)
+	}
+
+	return roles, nil
 }
 
 func (cm *channelMember) Cursor() *string {
@@ -79,6 +92,28 @@ func (cm *channelMember) Cursor() *string {
 	return model.NewString(encoded)
 }
 
+func graphQLRolesLoader(ctx context.Context, keys dataloader.Keys) []*dataloader.Result {
+	stringKeys := keys.Keys()
+	result := make([]*dataloader.Result, len(stringKeys))
+
+	c, err := getCtx(ctx)
+	if err != nil {
+		result[0] = &dataloader.Result{Error: err}
+		return result
+	}
+
+	roles, err := getGraphQLRoles(c, stringKeys)
+	if err != nil {
+		result[0] = &dataloader.Result{Error: err}
+		return result
+	}
+
+	for i, role := range roles {
+		result[i] = &dataloader.Result{Data: role}
+	}
+	return result
+}
+
 func getGraphQLRoles(c *web.Context, roleNames []string) ([]*model.Role, error) {
 	cleanedRoleNames, valid := model.CleanRoleNames(roleNames)
 	if !valid {
@@ -91,6 +126,17 @@ func getGraphQLRoles(c *web.Context, roleNames []string) ([]*model.Role, error)
 		return nil, appErr
 	}
 
+	// The roles need to be in the exact same order as the input slice.
+	tmp := make(map[string]*model.Role)
+	for _, r := range roles {
+		tmp[r.Name] = r
+	}
+
+	// We reuse the same slice and just rewrite the roles.
+	for i, roleName := range roleNames {
+		roles[i] = tmp[roleName]
+	}
+
 	return roles, nil
 }
 

api4/resolver_team_member.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"strings"
 
+	"github.com/graph-gophers/dataloader/v6"
 	"github.com/mattermost/mattermost-server/v6/model"
 )
 
@@ -60,10 +61,22 @@ func (tm *teamMember) SidebarCategories(ctx context.Context) ([]*model.SidebarCa
 
 // match with api4.getRolesByNames
 func (tm *teamMember) Roles_(ctx context.Context) ([]*model.Role, error) {
-	c, err := getCtx(ctx)
+	loader, err := getRolesLoader(ctx)
 	if err != nil {
 		return nil, err
 	}
 
-	return getGraphQLRoles(c, strings.Fields(tm.Roles))
+	thunk := loader.LoadMany(ctx, dataloader.NewKeysFromStrings(strings.Fields(tm.Roles)))
+	results, errs := thunk()
+	// All errors are the same. We just return the first one.
+	if len(errs) > 0 && errs[0] != nil {
+		return nil, err
+	}
+
+	roles := make([]*model.Role, len(results))
+	for i, res := range results {
+		roles[i] = res.(*model.Role)
+	}
+
+	return roles, nil
 }

api4/resolver_test.go

@@ -143,3 +143,85 @@ func TestGraphQLChannelsLeft(t *testing.T) {
 		assert.Len(t, q.ChannelsLeft, 0)
 	})
 }
+
+func TestGraphQLRolesLoader(t *testing.T) {
+	os.Setenv("MM_FEATUREFLAGS_GRAPHQL", "true")
+	defer os.Unsetenv("MM_FEATUREFLAGS_GRAPHQL")
+
+	th := Setup(t).InitBasic()
+	defer th.TearDown()
+
+	var q struct {
+		User struct {
+			ID    string `json:"id"`
+			Roles []struct {
+				ID   string `json:"id"`
+				Name string `json:"Name"`
+			} `json:"roles"`
+		} `json:"user"`
+		ChannelMembers []struct {
+			MsgCount float64 `json:"msgCount"`
+			Roles    []struct {
+				ID   string `json:"id"`
+				Name string `json:"Name"`
+			} `json:"roles"`
+		} `json:"channelMembers"`
+		TeamMembers []struct {
+			SchemeUser bool `json:"schemeUser"`
+			Roles      []struct {
+				ID   string `json:"id"`
+				Name string `json:"Name"`
+			} `json:"roles"`
+		}
+	}
+
+	input := graphQLInput{
+		OperationName: "channelMembers",
+		Query: `
+	query channelMembers {
+	user(id: "me") {
+		id
+		username
+		roles {
+			id
+			name
+		}
+	}
+	channelMembers(userId: "me") {
+		msgCount
+		roles {
+			id
+			name
+		}
+	}
+	teamMembers(userId: "me") {
+		schemeUser
+		roles {
+			id
+			name
+		}
+	}
+	}
+	`,
+	}
+
+	resp, err := th.MakeGraphQLRequest(&input)
+	require.NoError(t, err)
+	require.Len(t, resp.Errors, 0)
+	require.NoError(t, json.Unmarshal(resp.Data, &q))
+
+	require.Len(t, q.User.Roles, 1)
+	assert.Equal(t, "system_user", q.User.Roles[0].Name)
+
+	require.Len(t, q.ChannelMembers, 5)
+	for _, cm := range q.ChannelMembers {
+		require.Len(t, cm.Roles, 1)
+		assert.Equal(t, "channel_user", cm.Roles[0].Name)
+	}
+
+	require.Len(t, q.TeamMembers, 1)
+	for _, tm := range q.TeamMembers {
+		require.Len(t, tm.Roles, 1)
+		assert.Equal(t, "team_user", tm.Roles[0].Name)
+	}
+}

api4/resolver_user.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"net/http"
 
+	"github.com/graph-gophers/dataloader/v6"
 	"github.com/mattermost/mattermost-server/v6/model"
 	"github.com/mattermost/mattermost-server/v6/web"
 )
@@ -66,17 +67,29 @@ func getGraphQLUser(ctx context.Context, id string) (*user, error) {
 
 // match with api4.getRolesByNames
 func (u *user) Roles(ctx context.Context) ([]*model.Role, error) {
-	c, err := getCtx(ctx)
-	if err != nil {
-		return nil, err
-	}
-
 	roleNames := u.GetRoles()
 	if len(roleNames) == 0 {
 		return nil, nil
 	}
 
-	return getGraphQLRoles(c, roleNames)
+	loader, err := getRolesLoader(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	thunk := loader.LoadMany(ctx, dataloader.NewKeysFromStrings(roleNames))
+	results, errs := thunk()
+	// All errors are the same. We just return the first one.
+	if len(errs) > 0 && errs[0] != nil {
+		return nil, err
+	}
+
+	roles := make([]*model.Role, len(results))
+	for i, res := range results {
+		roles[i] = res.(*model.Role)
+	}
+
+	return roles, nil
 }
 
 // match with api4.getPreferences