diff --git a/go.mod b/go.mod index 5ffff20c..f3d0b49b 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,7 @@ require ( github.com/sirupsen/logrus v1.9.4 github.com/stretchr/testify v1.11.1 golang.org/x/exp v0.0.0-20220414153411-bcd21879b8fd + golang.org/x/time v0.5.0 gopkg.in/tomb.v2 v2.0.0-20161208151619-d5d1b5820637 ) diff --git a/go.sum b/go.sum index d8729762..c89b844a 100644 --- a/go.sum +++ b/go.sum @@ -29,6 +29,8 @@ golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac= golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM= golang.org/x/sys v0.29.0 h1:TPYlXGxvx1MGTn2GiZDhnjPA9wZzZeGKHHmKhHYvgaU= golang.org/x/sys v0.29.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk= +golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/tomb.v2 v2.0.0-20161208151619-d5d1b5820637 h1:yiW+nvdHb9LVqSHQBXfZCieqV4fzYhNBql77zY0ykqs= diff --git a/rest/client.go b/rest/client.go index bf55280d..0b0c5059 100644 --- a/rest/client.go +++ b/rest/client.go @@ -5,18 +5,25 @@ import ( "fmt" "net/http" "os" - "time" "reflect" + "time" + + "golang.org/x/time/rate" "github.com/massive-com/client-go/v3/rest/gen" ) +const ( + rateLimiterSleepInterval = 100 * time.Millisecond +) + type Client struct { *gen.ClientWithResponses httpClient *http.Client apiKey string trace bool pagination bool + limiter *rate.Limiter } type Option func(*Client) @@ -29,6 +36,14 @@ func WithPagination(enabled bool) Option { return func(c *Client) { c.pagination = enabled } } +// WithLimiter sets a rate limiter for the request rate. +// limiter is checked before +func WithLimiter(limiter *rate.Limiter) Option { + return func(c *Client) { + c.limiter = limiter + } +} + // New is backward-compatible (no options = trace=false, pagination=true) func New(apiKey string) *Client { return NewWithOptions(apiKey) @@ -66,8 +81,9 @@ func NewWithOptions(apiKey string, opts ...Option) *Client { var err error c.ClientWithResponses, err = gen.NewClientWithResponses("https://api.massive.com", - gen.WithHTTPClient(c.httpClient), // ← THIS makes the FIRST request traced + gen.WithHTTPClient(c.httpClient), // ← THIS makes the FIRST request traced gen.WithRequestEditorFn(c.addHeaders), + gen.WithRequestEditorFn(c.rateLimit), ) if err != nil { panic(err) @@ -76,6 +92,19 @@ func NewWithOptions(apiKey string, opts ...Option) *Client { return c } +// rateLimit enforces the rate limit on the client that was +// set by WithLimiter during creation time. +func (c *Client) rateLimit(_ context.Context, _ *http.Request) error { + if c.limiter != nil { + for !c.limiter.Allow() { + time.Sleep(rateLimiterSleepInterval) + continue + } + } + + return nil +} + func (c *Client) addHeaders(_ context.Context, req *http.Request) error { req.Header.Set("Authorization", "Bearer "+c.apiKey) req.Header.Set("User-Agent", "massive-go-test") @@ -83,11 +112,11 @@ func (c *Client) addHeaders(_ context.Context, req *http.Request) error { } // === Pointer helpers === -func String(v string) *string { return &v } -func Int(v int) *int { return &v } -func Int64(v int64) *int64 { return &v } +func String(v string) *string { return &v } +func Int(v int) *int { return &v } +func Int64(v int64) *int64 { return &v } func Float64(v float64) *float64 { return &v } -func Bool(v bool) *bool { return &v } +func Bool(v bool) *bool { return &v } // Generic Ptr (used for everything else, including custom enums) func Ptr[T any](v T) *T { return &v } diff --git a/rest/iterator.go b/rest/iterator.go index 138a28f6..9b62bd97 100644 --- a/rest/iterator.go +++ b/rest/iterator.go @@ -57,6 +57,9 @@ func (it *Iterator) fetchNextPage(urlStr string) ([]map[string]any, *string, err if err != nil { return nil, nil, err } + + it.client.rateLimit(req.Context(), req) + if err := it.client.addHeaders(context.Background(), req); err != nil { return nil, nil, err }