Add global rate limit to HTTP calls

Fixes #156

Author: nfx

common/client.go

@@ -12,6 +12,8 @@ import (
 	"sync"
 	"time"
 
+	"golang.org/x/time/rate"
+
 	"github.com/hashicorp/go-retryablehttp"
 	"github.com/mitchellh/go-homedir"
 	"gopkg.in/ini.v1"
@@ -30,6 +32,8 @@ type DatabricksClient struct {
 	TimeoutSeconds     int
 	DebugTruncateBytes int
 	DebugHeaders       bool
+	RateLimit          int
+	rateLimiter        *rate.Limiter
 	httpClient         *retryablehttp.Client
 	authMutex          sync.Mutex
 	authVisitor        func(r *http.Request) error
@@ -177,6 +181,10 @@ func (c *DatabricksClient) configureHTTPCLient() {
 	if c.TimeoutSeconds == 0 {
 		c.TimeoutSeconds = 60
 	}
+	if c.RateLimit == 0 {
+		c.RateLimit = 1200
+	}
+	c.rateLimiter = rate.NewLimiter(rate.Every(1*time.Minute), c.RateLimit)
 	// Set up a retryable HTTP Client to handle cases where the service returns
 	// a transient error on initial creation
 	retryDelayDuration := 10 * time.Second

common/http.go

@@ -413,6 +413,9 @@ func (c *DatabricksClient) genericQuery(ctx context.Context, method, requestURL
 	if c.httpClient == nil {
 		return nil, fmt.Errorf("DatabricksClient is not configured")
 	}
+	if err = c.rateLimiter.Wait(ctx); err != nil {
+		return nil, err
+	}
 	requestBody, err := makeRequestBody(method, &requestURL, data, true)
 	if err != nil {
 		return nil, err

go.mod

@@ -17,5 +17,6 @@ require (
 	github.com/pkg/errors v0.9.1
 	github.com/smartystreets/goconvey v1.6.4 // indirect
 	github.com/stretchr/testify v1.6.1
+	golang.org/x/time v0.0.0-20191024005414-555d28b269f0
 	gopkg.in/ini.v1 v1.62.0
-)
+)

go.sum

@@ -468,6 +468,7 @@ golang.org/x/text v0.3.3 h1:cokOdA+Jmi5PJGXLlLllQSgYigAEfHXJAERHVMaCc2k=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
+golang.org/x/time v0.0.0-20191024005414-555d28b269f0 h1:/5xXl8Y5W96D+TtHSlonuFqGHIWVuyCkGJLwGh9JJFs=
 golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
@@ -616,4 +617,4 @@ honnef.co/go/tools v0.0.1-2020.1.3/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9
 honnef.co/go/tools v0.0.1-2020.1.4/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
 rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0=
-rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA=
+rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA=

provider/provider.go

@@ -212,6 +212,12 @@ func DatabricksProvider() *schema.Provider {
 				Description: "Debug HTTP headers of requests made by the provider. Default is false. Visible only when TF_LOG=DEBUG is set",
 				DefaultFunc: schema.EnvDefaultFunc("DATABRICKS_DEBUG_HEADERS", false),
 			},
+			"rate_limit": {
+				Optional:    true,
+				Type:        schema.TypeInt,
+				Description: "Maximum number of requests per minute made to Databricks REST API by Terraform.",
+				DefaultFunc: schema.EnvDefaultFunc("DATABRICKS_RATE_LIMIT", 1200),
+			},
 		},
 		ConfigureContextFunc: func(c context.Context, d *schema.ResourceData) (interface{}, diag.Diagnostics) {
 			pc := common.DatabricksClient{}
@@ -286,6 +292,9 @@ func DatabricksProvider() *schema.Provider {
 			if v, ok := d.GetOk("debug_truncate_bytes"); ok {
 				pc.DebugTruncateBytes = v.(int)
 			}
+			if v, ok := d.GetOk("rate_limit"); ok {
+				pc.RateLimit = v.(int)
+			}
 			if v, ok := d.GetOk("debug_headers"); ok {
 				pc.DebugHeaders = v.(bool)
 			}