diff --git a/common/client.go b/common/client.go index 4312ff0d80..3db30f0a50 100644 --- a/common/client.go +++ b/common/client.go @@ -12,11 +12,20 @@ import ( "sync" "time" + "golang.org/x/time/rate" + "github.com/hashicorp/go-retryablehttp" "github.com/mitchellh/go-homedir" "gopkg.in/ini.v1" ) +// Default settings +const ( + DefaultRateLimit = 1200 + DebugTruncateBytes = 96 + DefaultTimeoutSeconds = 60 +) + // DatabricksClient is the client struct that contains clients for all the services available on Databricks type DatabricksClient struct { Host string @@ -30,6 +39,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 @@ -41,7 +52,7 @@ func (c *DatabricksClient) Configure() error { c.configureHTTPCLient() c.AzureAuth.databricksClient = c if c.DebugTruncateBytes == 0 { - c.DebugTruncateBytes = 96 + c.DebugTruncateBytes = DebugTruncateBytes } return nil } @@ -175,8 +186,12 @@ func (c *DatabricksClient) encodeBasicAuth(username, password string) string { func (c *DatabricksClient) configureHTTPCLient() { if c.TimeoutSeconds == 0 { - c.TimeoutSeconds = 60 + c.TimeoutSeconds = DefaultTimeoutSeconds + } + if c.RateLimit == 0 { + c.RateLimit = DefaultRateLimit } + 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 diff --git a/common/http.go b/common/http.go index c18c7c32b4..b302dbbc23 100644 --- a/common/http.go +++ b/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 diff --git a/go.mod b/go.mod index 41c53be65b..6b74187fe6 100644 --- a/go.mod +++ b/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 -) \ No newline at end of file +) diff --git a/go.sum b/go.sum index dde1b01a37..1fdf0f3e0a 100644 --- a/go.sum +++ b/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= \ No newline at end of file +rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA= diff --git a/provider/provider.go b/provider/provider.go index 3066614f3c..db7acfda65 100644 --- a/provider/provider.go +++ b/provider/provider.go @@ -204,7 +204,7 @@ func DatabricksProvider() *schema.Provider { Optional: true, Type: schema.TypeInt, Description: "Truncate JSON fields in JSON above this limit. Default is 96. Visible only when TF_LOG=DEBUG is set", - DefaultFunc: schema.EnvDefaultFunc("DATABRICKS_DEBUG_TRUNCATE_BYTES", 96), + DefaultFunc: schema.EnvDefaultFunc("DATABRICKS_DEBUG_TRUNCATE_BYTES", common.DebugTruncateBytes), }, "debug_headers": { Optional: true, @@ -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", common.DefaultRateLimit), + }, }, 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) }