[Fix] Support custom retry logic per method #1081
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -63,8 +63,7 @@ func (c *DatabricksClient) GetOAuthToken(ctx context.Context, authDetails string | |
|
|
||
| // Do sends an HTTP request against path. | ||
| func (c *DatabricksClient) Do(ctx context.Context, method, path string, | ||
| headers map[string]string, request, response any, visitors ...func(*http.Request) error) error { | ||
|
There was a problem hiding this comment. [optional] I would go with one parameter per line to echo the way function calls and struct declarations are made: func (c *DatabricksClient) Do(
ctx context.Context,
method string,
path string,
headers map[string]string,
request any,
response any,
visitors ...func(*http.Request) error
) error {There's a couple of similar patterns in the Go standard library but not many. One of the reason is that long lists of parameters are usually substituted with a struct (https://google.github.io/styleguide/go/best-practices#option-structure). I actually wanted to make that change for quite sometime but it didn't feel right sending one PR just for that. There was a problem hiding this comment. Yeah. This There was a problem hiding this comment. Agreed with not using a struct as this is out of scope for this PR. Let's at least format idiomatically then. Our Go style does not mandate a 80 char line length which makes the current formatting quite arbitrary. I'm fine with either having everything on a single line or one argument per line. |
||
| opts := []httpclient.DoOption{} | ||
| for _, v := range visitors { | ||
| opts = append(opts, httpclient.WithRequestVisitor(v)) | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,54 @@ | ||
| package config | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "io" | ||
| "net/http" | ||
| "strings" | ||
| "testing" | ||
|
|
||
| "github.com/databricks/databricks-sdk-go/httpclient" | ||
| "github.com/stretchr/testify/assert" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| type hc func(r *http.Request) (*http.Response, error) | ||
|
|
||
| func (cb hc) RoundTrip(r *http.Request) (*http.Response, error) { | ||
| return cb(r) | ||
| } | ||
|
|
||
| func (cb hc) SkipRetryOnIO() bool { | ||
| return true | ||
| } | ||
|
|
||
| func TestApiClient_RetriesGetPermissionsOnGatewayTimeout(t *testing.T) { | ||
| requestCount := 0 | ||
| c := &Config{ | ||
| HTTPTransport: hc(func(r *http.Request) (*http.Response, error) { | ||
| initialRequestCount := requestCount | ||
| requestCount++ | ||
| if initialRequestCount == 0 { | ||
| return &http.Response{ | ||
| Request: r, | ||
| StatusCode: http.StatusGatewayTimeout, | ||
| Body: io.NopCloser(strings.NewReader( | ||
| fmt.Sprintf(`{"error_code":"TEMPORARILY_UNAVAILABLE", "message":"The service at %s is taking too long to process your request. Please try again later or try a faster operation."}`, r.URL))), | ||
| }, nil | ||
| } | ||
| return &http.Response{ | ||
| Request: r, | ||
| StatusCode: http.StatusOK, | ||
| Body: io.NopCloser(strings.NewReader(`{"permissions": ["can_run_queries"]}`)), | ||
| }, nil | ||
| }), | ||
| } | ||
| client, err := c.NewApiClient() | ||
| require.NoError(t, err) | ||
| ctx := context.Background() | ||
| var res map[string][]string | ||
| err = client.Do(ctx, "GET", "/api/2.0/permissions/object/id", httpclient.WithResponseUnmarshal(&res)) | ||
| assert.NoError(t, err) | ||
| assert.Equal(t, map[string][]string{"permissions": {"can_run_queries"}}, res) | ||
| } | ||
|
Comment on lines
+1
to
+54
There was a problem hiding this comment. I'd recommend structuring the unit tests differently by having this test focused on how This test could look like the following (I did not verify that the code works): type mock struct {
MaxFails int // number of times the failed Response is returned
FailResponse *http.Response // response to return in case of fail
FailError error // error to return in case of fail
NumCalls int // total number of calls
}
func (m *mock) RoundTrip(r *http.Request) (*http.Response, error) {
m.NumCalls++
if m.NumCalls <= m.MaxFails {
return m.FailResponse, n.FailError
}
return &http.Response{
Request: r,
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader(`{}`)),
}, nil
}
func (m *mock) SkipRetryOnIO() bool {
return true
}
func TestApiClient_Do_retries(t *testing.T) {
testCases := []struct{
desc string
config *Config
errorRetrier ErrorRetrier
wantNumCalls int
} {
{
desc: "nil retrier",
mock: &mock{
MaxFails: 1,
FailResponse: &http.Response{StatusCode: http.StatusGatewayTimeout}
}
wantNumCalls: 1,
},
{
desc: "no retry",
mock: &mock{
MaxFails: 1,
FailResponse: &http.Response{StatusCode: http.StatusGatewayTimeout}
}
errorRetrier: func(context.Context, *http.Request, *common.ResponseWrapper, error) bool {
return false
},
wantNumCalls: 1,
},
{
desc: "retry 1 time",
mock: &mock{
MaxFails: 1,
FailResponse: &http.Response{StatusCode: http.StatusGatewayTimeout}
}
errorRetrier: func(context.Context, *http.Request, *common.ResponseWrapper, error) bool {
return true
},
wantNumCalls: 2,
},
{
desc: "retry 2 times",
mock: &mock{
MaxFails: 2,
FailResponse: &http.Response{StatusCode: http.StatusGatewayTimeout}
}
errorRetrier: func(_ context.Context, _ *http.Request, _ *common.ResponseWrapper, _ error) bool {
return true
},
wantNumCalls: 3,
},
{
desc: "retry 3 times",
mock: &mock{
MaxFails: 3,
FailResponse: &http.Response{StatusCode: http.StatusGatewayTimeout}
}
errorRetrier: func(_ context.Context, _ *http.Request, _ *common.ResponseWrapper, _ error) bool {
return true
},
wantNumCalls: 4,
},
}
func _, tc := range testCases {
t.Run(tc.desc, func(t *testing.T) {
cfg := &Config{HTTPTransport: tc.mock}
client, err := cfg.NewApiClient()
client.ErrorRetrier = tc.errorRetrier
err = client.Do(context.Background(), "GET", "test-path")
gotNumCalls = tc.mock.NumCalls
if gotNumCalls != tc.wantNumCalls {
t.Errorf("got %d calls, want %d", gotNumCalls, tc.wantNumCalls)
}
})
}
}Please feel free to ignore this comment if this is too much work or if the ApiClient cannot be instrumented that easily. There was a problem hiding this comment. It definitely can be instrumented this way, and this is a nice test case to use (I'll adapt it and include it in this PR). However, I did want to specifically test the get permissions pathway. Essentially, this tests that "the client returned by Config.GetApiClient() correctly implements retry on 504." I will add more test cases here though. There was a problem hiding this comment.
Sounds good to me as long as this complements the overall testing of the retry logic. |
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I can add this back in if desired, just a small formatting change.