api/graphql/loader/loader.go

package loader

import (
	"context"
	"errors"
	"sync"
	"time"

	"github.com/graph-gophers/dataloader"
)

var contextKey struct{}

// ErrNotFound is returned in batches when one or more things weren't found. Usually harmless.
var ErrNotFound = errors.New("not found")

// A Loader is a collection of data loaders and functions to act on them. It's supposed to be
// request-scoped, and will thus keep things cached indefinitely.
type Loader struct {
	mutex   sync.Mutex
	ctx     context.Context
	loaders map[string]*dataloader.Loader

	primedKeys map[string]map[string]bool
}

// New initializes the loader.
func New() *Loader {
	return &Loader{
		ctx: context.Background(),
		loaders: map[string]*dataloader.Loader{
			"Character.id":   dataloader.NewBatchedLoader(characterIDBatch, dataloader.WithWait(time.Millisecond)),
			"Character.nick": dataloader.NewBatchedLoader(characterNickBatch, dataloader.WithWait(time.Millisecond)),
			"Channel.name":   dataloader.NewBatchedLoader(channelNameBatch, dataloader.WithWait(time.Millisecond)),
		},
		primedKeys: make(map[string]map[string]bool),
	}
}

// FromContext gets the Loader from context.
func FromContext(ctx context.Context) *Loader {
	value := ctx.Value(&contextKey)
	if value == nil {
		return nil
	}

	return value.(*Loader)
}

// ToContext gets a context with the loader as a value
func (loader *Loader) ToContext(ctx context.Context) context.Context {
	loader.ctx = ctx
	return context.WithValue(ctx, &contextKey, loader)
}

func (loader *Loader) prime(key string, values []string) {
	loader.mutex.Lock()
	if loader.primedKeys[key] == nil {
		loader.primedKeys[key] = make(map[string]bool)
	}

	for _, value := range values {
		loader.primedKeys[key][value] = true
	}
	loader.mutex.Unlock()
}

func (loader *Loader) loadPrimed(key string) {
	loader.mutex.Lock()
	if len(loader.primedKeys[key]) > 0 {
		primedKeys := make([]string, 0, len(loader.primedKeys[key]))
		for key := range loader.primedKeys[key] {
			primedKeys = append(primedKeys, key)
		}

		loader.loaders[key].LoadMany(loader.ctx, dataloader.NewKeysFromStrings(primedKeys))
		loader.primedKeys[key] = nil
	}
	loader.mutex.Unlock()
}