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() }