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

  1. package loader
  2. import (
  3. "context"
  4. "errors"
  5. "sync"
  6. "time"
  7. "github.com/graph-gophers/dataloader"
  8. )
  9. var contextKey struct{}
  10. // ErrNotFound is returned in batches when one or more things weren't found. Usually harmless.
  11. var ErrNotFound = errors.New("not found")
  12. // A Loader is a collection of data loaders and functions to act on them. It's supposed to be
  13. // request-scoped, and will thus keep things cached indefinitely.
  14. type Loader struct {
  15. mutex sync.Mutex
  16. ctx context.Context
  17. loaders map[string]*dataloader.Loader
  18. primedKeys map[string]map[string]bool
  19. }
  20. // New initializes the loader.
  21. func New() *Loader {
  22. return &Loader{
  23. ctx: context.Background(),
  24. loaders: map[string]*dataloader.Loader{
  25. "Character.id": dataloader.NewBatchedLoader(characterIDBatch, dataloader.WithWait(time.Millisecond)),
  26. "Character.nick": dataloader.NewBatchedLoader(characterNickBatch, dataloader.WithWait(time.Millisecond)),
  27. "Channel.name": dataloader.NewBatchedLoader(channelNameBatch, dataloader.WithWait(time.Millisecond)),
  28. },
  29. primedKeys: make(map[string]map[string]bool),
  30. }
  31. }
  32. // FromContext gets the Loader from context.
  33. func FromContext(ctx context.Context) *Loader {
  34. value := ctx.Value(&contextKey)
  35. if value == nil {
  36. return nil
  37. }
  38. return value.(*Loader)
  39. }
  40. // ToContext gets a context with the loader as a value
  41. func (loader *Loader) ToContext(ctx context.Context) context.Context {
  42. loader.ctx = ctx
  43. return context.WithValue(ctx, &contextKey, loader)
  44. }
  45. func (loader *Loader) prime(key string, values []string) {
  46. loader.mutex.Lock()
  47. if loader.primedKeys[key] == nil {
  48. loader.primedKeys[key] = make(map[string]bool)
  49. }
  50. for _, value := range values {
  51. loader.primedKeys[key][value] = true
  52. }
  53. loader.mutex.Unlock()
  54. }
  55. func (loader *Loader) loadPrimed(key string) {
  56. loader.mutex.Lock()
  57. if len(loader.primedKeys[key]) > 0 {
  58. primedKeys := make([]string, 0, len(loader.primedKeys[key]))
  59. for key := range loader.primedKeys[key] {
  60. primedKeys = append(primedKeys, key)
  61. }
  62. loader.loaders[key].LoadMany(loader.ctx, dataloader.NewKeysFromStrings(primedKeys))
  63. loader.primedKeys[key] = nil
  64. }
  65. loader.mutex.Unlock()
  66. }