GraphQL API and utilities for the rpdata project
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

80 lines
2.1 KiB

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*2)),
"Character.nick": dataloader.NewBatchedLoader(characterNickBatch, dataloader.WithWait(time.Millisecond*2)),
"Channel.name": dataloader.NewBatchedLoader(channelNameBatch, dataloader.WithWait(time.Millisecond*2)),
},
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()
}