server/graph/loaders/userloader.go

package loaders

import (
	"context"
	"git.aiterp.net/stufflog/server/database/repositories"
	"git.aiterp.net/stufflog/server/internal/xlerrors"
	"git.aiterp.net/stufflog/server/models"
	"time"
)

// go run github.com/vektah/dataloaden UserLoader string *git.aiterp.net/stufflog/server/models.User

var userLoaderCtxKey = "ctx.stufflog.UserLoader"

func NewUserLoaderContext(ctx context.Context, userRepo repositories.UserRepository) context.Context {
	return context.WithValue(ctx, userLoaderCtxKey, NewUserLoader(ctx, userRepo))
}

func UserLoaderFromContext(ctx context.Context) *UserLoader {
	return ctx.Value(userLoaderCtxKey).(*UserLoader)
}

func NewUserLoader(ctx context.Context, userRepo repositories.UserRepository) *UserLoader {
	return &UserLoader{
		wait:     2 * time.Millisecond,
		maxBatch: 100,
		fetch: func(keys []string) ([]*models.User, []error) {
			results := make([]*models.User, len(keys))
			errors := make([]error, len(keys))

			users, err := userRepo.List(ctx, models.UserFilter{UserIDs: keys})
			if err != nil {
				for i := range errors {
					errors[i] = err
				}

				return results, errors
			}

			userMap := make(map[string]*models.User, len(users))
			for _, user := range users {
				userMap[user.ID] = user
			}

			for i, id := range keys {
				user := userMap[id]
				if user != nil {
					results[i] = user
				} else {
					errors[i] = xlerrors.NotFound("User")
				}
			}

			return results, errors
		},
	}
}