From f003f3f25743160fbf7b442c848731c45d420b80 Mon Sep 17 00:00:00 2001 From: Felipe Pontes Date: Mon, 8 Jul 2019 18:11:46 -0300 Subject: [PATCH] Add ability to specify output filename Refactor default filename to snake_case --- dataloaden.go | 17 +- example/filename/myloader.go | 224 ++++++++++++++++++ example/filename/user.go | 3 + .../{userloader_gen.go => user_loader_gen.go} | 0 ...loader_gen.go => user_slice_loader_gen.go} | 0 .../{userloader_gen.go => user_loader_gen.go} | 0 go.mod | 1 + go.sum | 2 + pkg/generator/generator.go | 20 +- 9 files changed, 260 insertions(+), 7 deletions(-) create mode 100644 example/filename/myloader.go create mode 100644 example/filename/user.go rename example/pkgname/{userloader_gen.go => user_loader_gen.go} (100%) rename example/slice/{usersliceloader_gen.go => user_slice_loader_gen.go} (100%) rename example/{userloader_gen.go => user_loader_gen.go} (100%) diff --git a/dataloaden.go b/dataloaden.go index 3419286..8a4c543 100644 --- a/dataloaden.go +++ b/dataloaden.go @@ -1,6 +1,7 @@ package main import ( + "flag" "fmt" "os" @@ -8,7 +9,12 @@ import ( ) func main() { - if len(os.Args) != 4 { + output := flag.String("output", "", "output filename") + + flag.Parse() + + args := flag.Args() + if len(args) != 3 { fmt.Println("usage: name keyType valueType") fmt.Println(" example:") fmt.Println(" dataloaden 'UserLoader int []*github.com/my/package.User'") @@ -20,8 +26,13 @@ func main() { fmt.Fprintln(os.Stderr, err.Error()) os.Exit(2) } - - if err := generator.Generate(os.Args[1], os.Args[2], os.Args[3], wd); err != nil { + if err := generator.Generate(&generator.GenerateInput{ + Name: args[0], + KeyType: args[1], + ValueType: args[2], + WorkingDir: wd, + Output: *output, + }); err != nil { fmt.Fprintln(os.Stderr, err.Error()) os.Exit(2) } diff --git a/example/filename/myloader.go b/example/filename/myloader.go new file mode 100644 index 0000000..42b04a2 --- /dev/null +++ b/example/filename/myloader.go @@ -0,0 +1,224 @@ +// Code generated by github.com/vektah/dataloaden, DO NOT EDIT. + +package filename + +import ( + "sync" + "time" + + "github.com/vektah/dataloaden/example" +) + +// UserLoaderConfig captures the config to create a new UserLoader +type UserLoaderConfig struct { + // Fetch is a method that provides the data for the loader + Fetch func(keys []string) ([]*example.User, []error) + + // Wait is how long wait before sending a batch + Wait time.Duration + + // MaxBatch will limit the maximum number of keys to send in one batch, 0 = not limit + MaxBatch int +} + +// NewUserLoader creates a new UserLoader given a fetch, wait, and maxBatch +func NewUserLoader(config UserLoaderConfig) *UserLoader { + return &UserLoader{ + fetch: config.Fetch, + wait: config.Wait, + maxBatch: config.MaxBatch, + } +} + +// UserLoader batches and caches requests +type UserLoader struct { + // this method provides the data for the loader + fetch func(keys []string) ([]*example.User, []error) + + // how long to done before sending a batch + wait time.Duration + + // this will limit the maximum number of keys to send in one batch, 0 = no limit + maxBatch int + + // INTERNAL + + // lazily created cache + cache map[string]*example.User + + // the current batch. keys will continue to be collected until timeout is hit, + // then everything will be sent to the fetch method and out to the listeners + batch *userLoaderBatch + + // mutex to prevent races + mu sync.Mutex +} + +type userLoaderBatch struct { + keys []string + data []*example.User + error []error + closing bool + done chan struct{} +} + +// Load a User by key, batching and caching will be applied automatically +func (l *UserLoader) Load(key string) (*example.User, error) { + return l.LoadThunk(key)() +} + +// LoadThunk returns a function that when called will block waiting for a User. +// This method should be used if you want one goroutine to make requests to many +// different data loaders without blocking until the thunk is called. +func (l *UserLoader) LoadThunk(key string) func() (*example.User, error) { + l.mu.Lock() + if it, ok := l.cache[key]; ok { + l.mu.Unlock() + return func() (*example.User, error) { + return it, nil + } + } + if l.batch == nil { + l.batch = &userLoaderBatch{done: make(chan struct{})} + } + batch := l.batch + pos := batch.keyIndex(l, key) + l.mu.Unlock() + + return func() (*example.User, error) { + <-batch.done + + var data *example.User + if pos < len(batch.data) { + data = batch.data[pos] + } + + var err error + // its convenient to be able to return a single error for everything + if len(batch.error) == 1 { + err = batch.error[0] + } else if batch.error != nil { + err = batch.error[pos] + } + + if err == nil { + l.mu.Lock() + l.unsafeSet(key, data) + l.mu.Unlock() + } + + return data, err + } +} + +// LoadAll fetches many keys at once. It will be broken into appropriate sized +// sub batches depending on how the loader is configured +func (l *UserLoader) LoadAll(keys []string) ([]*example.User, []error) { + results := make([]func() (*example.User, error), len(keys)) + + for i, key := range keys { + results[i] = l.LoadThunk(key) + } + + users := make([]*example.User, len(keys)) + errors := make([]error, len(keys)) + for i, thunk := range results { + users[i], errors[i] = thunk() + } + return users, errors +} + +// LoadAllThunk returns a function that when called will block waiting for a Users. +// This method should be used if you want one goroutine to make requests to many +// different data loaders without blocking until the thunk is called. +func (l *UserLoader) LoadAllThunk(keys []string) func() ([]*example.User, []error) { + results := make([]func() (*example.User, error), len(keys)) + for i, key := range keys { + results[i] = l.LoadThunk(key) + } + return func() ([]*example.User, []error) { + users := make([]*example.User, len(keys)) + errors := make([]error, len(keys)) + for i, thunk := range results { + users[i], errors[i] = thunk() + } + return users, errors + } +} + +// Prime the cache with the provided key and value. If the key already exists, no change is made +// and false is returned. +// (To forcefully prime the cache, clear the key first with loader.clear(key).prime(key, value).) +func (l *UserLoader) Prime(key string, value *example.User) bool { + l.mu.Lock() + var found bool + if _, found = l.cache[key]; !found { + // make a copy when writing to the cache, its easy to pass a pointer in from a loop var + // and end up with the whole cache pointing to the same value. + cpy := *value + l.unsafeSet(key, &cpy) + } + l.mu.Unlock() + return !found +} + +// Clear the value at key from the cache, if it exists +func (l *UserLoader) Clear(key string) { + l.mu.Lock() + delete(l.cache, key) + l.mu.Unlock() +} + +func (l *UserLoader) unsafeSet(key string, value *example.User) { + if l.cache == nil { + l.cache = map[string]*example.User{} + } + l.cache[key] = value +} + +// keyIndex will return the location of the key in the batch, if its not found +// it will add the key to the batch +func (b *userLoaderBatch) keyIndex(l *UserLoader, key string) int { + for i, existingKey := range b.keys { + if key == existingKey { + return i + } + } + + pos := len(b.keys) + b.keys = append(b.keys, key) + if pos == 0 { + go b.startTimer(l) + } + + if l.maxBatch != 0 && pos >= l.maxBatch-1 { + if !b.closing { + b.closing = true + l.batch = nil + go b.end(l) + } + } + + return pos +} + +func (b *userLoaderBatch) startTimer(l *UserLoader) { + time.Sleep(l.wait) + l.mu.Lock() + + // we must have hit a batch limit and are already finalizing this batch + if b.closing { + l.mu.Unlock() + return + } + + l.batch = nil + l.mu.Unlock() + + b.end(l) +} + +func (b *userLoaderBatch) end(l *UserLoader) { + b.data, b.error = l.fetch(b.keys) + close(b.done) +} diff --git a/example/filename/user.go b/example/filename/user.go new file mode 100644 index 0000000..34c8a5c --- /dev/null +++ b/example/filename/user.go @@ -0,0 +1,3 @@ +package filename + +//go:generate go run github.com/vektah/dataloaden -output=myloader.go UserLoader string *github.com/vektah/dataloaden/example.User diff --git a/example/pkgname/userloader_gen.go b/example/pkgname/user_loader_gen.go similarity index 100% rename from example/pkgname/userloader_gen.go rename to example/pkgname/user_loader_gen.go diff --git a/example/slice/usersliceloader_gen.go b/example/slice/user_slice_loader_gen.go similarity index 100% rename from example/slice/usersliceloader_gen.go rename to example/slice/user_slice_loader_gen.go diff --git a/example/userloader_gen.go b/example/user_loader_gen.go similarity index 100% rename from example/userloader_gen.go rename to example/user_loader_gen.go diff --git a/go.mod b/go.mod index ba56ca0..5af6a15 100644 --- a/go.mod +++ b/go.mod @@ -2,6 +2,7 @@ module github.com/vektah/dataloaden require ( github.com/davecgh/go-spew v1.1.0 // indirect + github.com/iancoleman/strcase v0.0.0-20190422225806-e506e3ef7365 github.com/pkg/errors v0.8.1 github.com/pmezard/go-difflib v1.0.0 // indirect github.com/stretchr/testify v1.2.1 diff --git a/go.sum b/go.sum index a350afb..13db6d4 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,7 @@ github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/iancoleman/strcase v0.0.0-20190422225806-e506e3ef7365 h1:ECW73yc9MY7935nNYXUkK7Dz17YuSUI9yqRqYS8aBww= +github.com/iancoleman/strcase v0.0.0-20190422225806-e506e3ef7365/go.mod h1:SK73tn/9oHe+/Y0h39VT4UCxmurVJkR5NA7kMEAOgSE= github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I= github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= diff --git a/pkg/generator/generator.go b/pkg/generator/generator.go index ff618e7..24a8a40 100644 --- a/pkg/generator/generator.go +++ b/pkg/generator/generator.go @@ -9,6 +9,7 @@ import ( "strings" "unicode" + "github.com/iancoleman/strcase" "github.com/pkg/errors" "golang.org/x/tools/go/packages" "golang.org/x/tools/imports" @@ -78,15 +79,26 @@ func parseType(str string) (*goType, error) { return t, nil } -func Generate(name string, keyType string, valueType string, wd string) error { - data, err := getData(name, keyType, valueType, wd) +type GenerateInput struct { + Name string + KeyType string + ValueType string + WorkingDir string + Output string +} + +func Generate(input *GenerateInput) error { + data, err := getData(input.Name, input.KeyType, input.ValueType, input.WorkingDir) if err != nil { return err } - filename := strings.ToLower(data.Name) + "_gen.go" + filename := input.Output + if filename == "" { + filename = strcase.ToSnake(data.Name) + "_gen.go" + } - if err := writeTemplate(filepath.Join(wd, filename), data); err != nil { + if err := writeTemplate(filepath.Join(input.WorkingDir, filename), data); err != nil { return err }