Repository: snowflakedb/gosnowflake
Pull request #948 (2023-10-26 09:07:36)

SNOW-878073 Add MaxRetryCount as configurable parameter
### Description
Added MaxRetryCount as configurable parameter. It specifies how many times requests can be retried.

### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [ ] Created tests which fail without the change (if possible)
- [ ] All tests passing
- [ ] Extended the README / documentation, if necessary


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
> Merging [#948](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) (4d86d91) into [master](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/ce9ef59c09504e5bfb63b30060fb4612d556eee2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) (ce9ef59) will **decrease** coverage by `2.18%`.
> Report is 1 commits behind head on master.
> The diff coverage is `92.98%`.

```diff
@@            Coverage Diff             @@
##           master     #948      +/-   ##
==========================================
- Coverage   83.26%   81.09%   -2.18%     
==========================================
  Files          49       48       -1     
  Lines        9340     9339       -1     
==========================================
- Hits         7777     7573     -204     
- Misses       1563     1766     +203     
```


| [Files](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | Coverage Δ | |
|---|---|---|
| [async.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-YXN5bmMuZ28=) | `81.74% <100.00%> (ø)` | |
| [auth.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-YXV0aC5nbw==) | `80.48% <100.00%> (+0.10%)` | :arrow_up: |
| [authexternalbrowser.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-YXV0aGV4dGVybmFsYnJvd3Nlci5nbw==) | `57.38% <100.00%> (ø)` | |
| [authokta.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-YXV0aG9rdGEuZ28=) | `78.42% <100.00%> (ø)` | |
| [chunk\_downloader.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-Y2h1bmtfZG93bmxvYWRlci5nbw==) | `82.18% <100.00%> (-1.92%)` | :arrow_down: |
| [client.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-Y2xpZW50Lmdv) | `100.00% <100.00%> (ø)` | |
| [connection.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-Y29ubmVjdGlvbi5nbw==) | `83.51% <100.00%> (-0.90%)` | :arrow_down: |
| [file\_transfer\_agent.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-ZmlsZV90cmFuc2Zlcl9hZ2VudC5nbw==) | `73.27% <100.00%> (-1.84%)` | :arrow_down: |
| [heartbeat.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-aGVhcnRiZWF0Lmdv) | `69.84% <100.00%> (ø)` | |
| [monitoring.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-bW9uaXRvcmluZy5nbw==) | `83.58% <100.00%> (ø)` | |
| ... and [6 more](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | |

... and [5 files with indirect coverage changes](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/948/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb)


Patch:
diff --git a/auth.go b/auth.go
index c3894e4..9493459 100644
--- a/auth.go
+++ b/auth.go
@@ -226,7 +226,7 @@ func postAuth(
 
 	fullURL := sr.getFullURL(loginRequestPath, params)
 	logger.Infof("full URL: %v", fullURL)
-	resp, err := sr.FuncAuthPost(ctx, client, fullURL, headers, bodyCreator, timeout)
+	resp, err := sr.FuncAuthPost(ctx, client, fullURL, headers, bodyCreator, timeout, sr.MaxRetryCount)
 	if err != nil {
 		return nil, err
 	}
diff --git a/chunk_downloader.go b/chunk_downloader.go
index 6806e5e..626b5c4 100644
--- a/chunk_downloader.go
+++ b/chunk_downloader.go
@@ -264,7 +264,7 @@ func getChunk(
 	if err != nil {
 		return nil, err
 	}
-	return newRetryHTTP(ctx, sc.rest.Client, http.NewRequest, u, headers, timeout, sc.currentTimeProvider, sc.cfg).execute()
+	return newRetryHTTP(ctx, sc.rest.Client, http.NewRequest, u, headers, timeout, sc.rest.MaxRetryCount, sc.currentTimeProvider, sc.cfg).execute()
 }
 
 func (scd *snowflakeChunkDownloader) startArrowBatches() error {
@@ -638,7 +638,7 @@ func (f *httpStreamChunkFetcher) fetch(URL string, rows chan<- []*string) error
 	if err != nil {
 		return err
 	}
-	res, err := newRetryHTTP(context.Background(), f.client, http.NewRequest, fullURL, f.headers, 0, defaultTimeProvider, nil).execute()
+	res, err := newRetryHTTP(context.Background(), f.client, http.NewRequest, fullURL, f.headers, 0, 0, defaultTimeProvider, nil).execute()
 	if err != nil {
 		return err
 	}
diff --git a/connection.go b/connection.go
index 087ac72..5b39d14 100644
--- a/connection.go
+++ b/connection.go
@@ -789,6 +789,7 @@ func buildSnowflakeConn(ctx context.Context, config Config) (*snowflakeConn, err
 		TokenAccessor:       tokenAccessor,
 		LoginTimeout:        sc.cfg.LoginTimeout,
 		RequestTimeout:      sc.cfg.RequestTimeout,
+		MaxRetryCount:       sc.cfg.MaxRetryCount,
 		FuncPost:            postRestful,
 		FuncGet:             getRestful,
 		FuncAuthPost:        postAuthRestful,
diff --git a/dsn.go b/dsn.go
index b3a8212..1ccebba 100644
--- a/dsn.go
+++ b/dsn.go
@@ -19,12 +19,13 @@ import (
 )
 
 const (
-	defaultClientTimeout          = 300 * time.Second // Timeout for network round trip + read out http response
+	defaultClientTimeout          = 900 * time.Second // Timeout for network round trip + read out http response
 	defaultJWTClientTimeout       = 10 * time.Second  // Timeout for network round trip + read out http response but used for JWT auth
-	defaultLoginTimeout           = 60 * time.Second  // Timeout for retry for login EXCLUDING clientTimeout
+	defaultLoginTimeout           = 300 * time.Second // Timeout for retry for login EXCLUDING clientTimeout
 	defaultRequestTimeout         = 0 * time.Second   // Timeout for retry for request EXCLUDING clientTimeout
 	defaultJWTTimeout             = 60 * time.Second
 	defaultExternalBrowserTimeout = 120 * time.Second // Timeout for external browser login
+	defaultMaxRetryCount          = 7                 // specifies maximum number of subsequent retries
 	defaultDomain                 = ".snowflakecomputing.com"
 )
 
@@ -74,6 +75,7 @@ type Config struct {
 	ClientTimeout          time.Duration // Timeout for network round trip + read out http response
 	JWTClientTimeout       time.Duration // Timeout for network round trip + read out http response used when JWT token auth is taking place
 	ExternalBrowserTimeout time.Duration // Timeout for external browser login
+	MaxRetryCount          int           // Specifies how many times non-periodic HTTP request can be retried
 
 	Application  string           // application name.
 	InsecureMode bool             // driver doesn't check certificate revocation status
@@ -205,6 +207,9 @@ func DSN(cfg *Config) (dsn string, err error) {
 	if cfg.ExternalBrowserTimeout != defaultExternalBrowserTimeout {
 		params.Add("externalBrowserTimeout", strconv.FormatInt(int64(cfg.ExternalBrowserTimeout/time.Second), 10))
 	}
+	if cfg.MaxRetryCount != defaultMaxRetryCount {
+		params.Add("maxRetryCount", strconv.Itoa(cfg.MaxRetryCount))
+	}
 	if cfg.Application != clientType {
 		params.Add("application", cfg.Application)
 	}
@@ -471,6 +476,9 @@ func fillMissingConfigParameters(cfg *Config) error {
 	if cfg.ExternalBrowserTimeout == 0 {
 		cfg.ExternalBrowserTimeout = defaultExternalBrowserTimeout
 	}
+	if cfg.MaxRetryCount == 0 {
+		cfg.MaxRetryCount = defaultMaxRetryCount
+	}
 	if strings.Trim(cfg.Application, " ") == "" {
 		cfg.Application = clientType
 	}
@@ -642,6 +650,11 @@ func parseDSNParams(cfg *Config, params string) (err error) {
 			if err != nil {
 				return err
 			}
+		case "maxRetryCount":
+			cfg.MaxRetryCount, err = strconv.Atoi(value)
+			if err != nil {
+				return err
+			}
 		case "application":
 			cfg.Application = value
 		case "authenticator":
diff --git a/ocsp.go b/ocsp.go
index 297a999..8fc7cef 100644
--- a/ocsp.go
+++ b/ocsp.go
@@ -358,7 +358,7 @@ func checkOCSPCacheServer(
 	ocspS *ocspStatus) {
 	var respd map[string][]interface{}
 	headers := make(map[string]string)
-	res, err := newRetryHTTP(ctx, client, req, ocspServerHost, headers, totalTimeout, defaultTimeProvider, nil).execute()
+	res, err := newRetryHTTP(ctx, client, req, ocspServerHost, headers, totalTimeout, defaultMaxRetryCount, defaultTimeProvider, nil).execute()
 	if err != nil {
 		logger.Errorf("failed to get OCSP cache from OCSP Cache Server. %v", err)
 		return nil, &ocspStatus{
@@ -413,7 +413,7 @@ func retryOCSP(
 	}
 	res, err := newRetryHTTP(
 		ctx, client, req, ocspHost, headers,
-		totalTimeout*time.Duration(multiplier), defaultTimeProvider, nil).doPost().setBody(reqBody).execute()
+		totalTimeout*time.Duration(multiplier), defaultMaxRetryCount, defaultTimeProvider, nil).doPost().setBody(reqBody).execute()
 	if err != nil {
 		return ocspRes, ocspResBytes, &ocspStatus{
 			code: ocspFailedSubmit,
@@ -466,7 +466,7 @@ func fallbackRetryOCSPToGETRequest(
 		multiplier = 3 // up to 3 times for Fail Close mode
 	}
 	res, err := newRetryHTTP(ctx, client, req, ocspHost, headers,
-		totalTimeout*time.Duration(multiplier), defaultTimeProvider, nil).execute()
+		totalTimeout*time.Duration(multiplier), defaultMaxRetryCount, defaultTimeProvider, nil).execute()
 	if err != nil {
 		return ocspRes, ocspResBytes, &ocspStatus{
 			code: ocspFailedSubmit,
diff --git a/restful.go b/restful.go
index c94e286..777d94d 100644
--- a/restful.go
+++ b/restful.go
@@ -44,7 +44,7 @@ const (
 type (
 	funcGetType      func(context.Context, *snowflakeRestful, *url.URL, map[string]string, time.Duration) (*http.Response, error)
 	funcPostType     func(context.Context, *snowflakeRestful, *url.URL, map[string]string, []byte, time.Duration, currentTimeProvider, *Config) (*http.Response, error)
-	funcAuthPostType func(context.Context, *http.Client, *url.URL, map[string]string, bodyCreatorType, time.Duration) (*http.Response, error)
+	funcAuthPostType func(context.Context, *http.Client, *url.URL, map[string]string, bodyCreatorType, time.Duration, int) (*http.Response, error)
 	bodyCreatorType  func() ([]byte, error)
 )
 
@@ -58,6 +58,7 @@ type snowflakeRestful struct {
 	Protocol       string
 	LoginTimeout   time.Duration // Login timeout
 	RequestTimeout time.Duration // request timeout
+	MaxRetryCount  int
 
 	Client        *http.Client
 	JWTClient     *http.Client
@@ -165,7 +166,7 @@ func postRestful(
 	currentTimeProvider currentTimeProvider,
 	cfg *Config) (
 	*http.Response, error) {
-	return newRetryHTTP(ctx, sr.Client, http.NewRequest, fullURL, headers, timeout, currentTimeProvider, cfg).
+	return newRetryHTTP(ctx, sr.Client, http.NewRequest, fullURL, headers, timeout, sr.MaxRetryCount, currentTimeProvider, cfg).
 		doPost().
 		setBody(body).
 		execute()
@@ -178,7 +179,7 @@ func getRestful(
 	headers map[string]string,
 	timeout time.Duration) (
 	*http.Response, error) {
-	return newRetryHTTP(ctx, sr.Client, http.NewRequest, fullURL, headers, timeout, defaultTimeProvider, nil).execute()
+	return newRetryHTTP(ctx, sr.Client, http.NewRequest, fullURL, headers, timeout, sr.MaxRetryCount, defaultTimeProvider, nil).execute()
 }
 
 func postAuthRestful(
@@ -187,9 +188,10 @@ func postAuthRestful(
 	fullURL *url.URL,
 	headers map[string]string,
 	bodyCreator bodyCreatorType,
-	timeout time.Duration) (
+	timeout time.Duration,
+	maxRetryCount int) (
 	*http.Response, error) {
-	return newRetryHTTP(ctx, client, http.NewRequest, fullURL, headers, timeout, defaultTimeProvider, nil).
+	return newRetryHTTP(ctx, client, http.NewRequest, fullURL, headers, timeout, maxRetryCount, defaultTimeProvider, nil).
 		doPost().
 		setBodyCreator(bodyCreator).
 		execute()
diff --git a/retry.go b/retry.go
index 725a8fc..8995b9c 100644
--- a/retry.go
+++ b/retry.go
@@ -17,11 +17,6 @@ import (
 	"time"
 )
 
-const (
-	// defaultMaxRetryCount specifies maximum number of subsequent retries
-	defaultMaxRetryCount = 7
-)
-
 type waitAlgo struct {
 	mutex  *sync.Mutex // required for *rand.Rand usage
 	random *rand.Rand
@@ -248,6 +243,7 @@ type retryHTTP struct {
 	headers             map[string]string
 	bodyCreator         bodyCreatorType
 	timeout             time.Duration
+	maxRetryCount       int
 	currentTimeProvider currentTimeProvider
 	cfg                 *Config
 }
@@ -258,6 +254,7 @@ func newRetryHTTP(ctx context.Context,
 	fullURL *url.URL,
 	headers map[string]string,
 	timeout time.Duration,
+	maxRetryCount int,
 	currentTimeProvider currentTimeProvider,
 	cfg *Config) *retryHTTP {
 	instance := retryHTTP{}
@@ -268,6 +265,7 @@ func newRetryHTTP(ctx context.Context,
 	instance.fullURL = fullURL
 	instance.headers = headers
 	instance.timeout = timeout
+	instance.maxRetryCount = maxRetryCount
 	instance.bodyCreator = emptyBodyCreator
 	instance.currentTimeProvider = currentTimeProvider
 	instance.cfg = cfg
@@ -341,7 +339,7 @@ func (r *retryHTTP) execute() (res *http.Response, err error) {
 			logger.WithContext(r.ctx).Infof("to timeout: %v", totalTimeout)
 			// if any timeout is set
 			totalTimeout -= time.Duration(sleepTime * float64(time.Second))
-			if totalTimeout <= 0 || retryCounter >= defaultMaxRetryCount {
+			if totalTimeout <= 0 || retryCounter > r.maxRetryCount {
 				if err != nil {
 					return nil, err
 				}


Test patch:
diff --git a/dsn_test.go b/dsn_test.go
index e49207b..f939f86 100644
--- a/dsn_test.go
+++ b/dsn_test.go
@@ -336,6 +336,23 @@ func TestParseDSN(t *testing.T) {
 				ClientTimeout:             defaultClientTimeout,
 				JWTClientTimeout:          defaultJWTClientTimeout,
 				IncludeRetryReason:        ConfigBoolTrue,
+				MaxRetryCount:             defaultMaxRetryCount,
+			},
+			ocspMode: ocspModeFailOpen,
+		},
+		{
+			dsn: "u:p@a?database=d&maxRetryCount=20",
+			config: &Config{
+				Account: "a", User: "u", Password: "p",
+				Protocol: "https", Host: "a.snowflakecomputing.com", Port: 443,
+				Database: "d", Schema: "",
+				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
+				OCSPFailOpen:              OCSPFailOpenTrue,
+				ValidateDefaultParameters: ConfigBoolTrue,
+				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
+				IncludeRetryReason:        ConfigBoolTrue,
+				MaxRetryCount:             20,
 			},
 			ocspMode: ocspModeFailOpen,
 		},
@@ -1239,6 +1256,16 @@ func TestDSN(t *testing.T) {
 			},
 			dsn: "u:p@a.b.c.snowflakecomputing.com:443?ocspFailOpen=true&region=b.c&tmpDirPath=%2Ftmp&validateDefaultParameters=true",
 		},
+		{
+			cfg: &Config{
+				User:               "u",
+				Password:           "p",
+				Account:            "a.b.c",
+				IncludeRetryReason: ConfigBoolFalse,
+				MaxRetryCount:      30,
+			},
+			dsn: "u:p@a.b.c.snowflakecomputing.com:443?includeRetryReason=false&maxRetryCount=30&ocspFailOpen=true&region=b.c&validateDefaultParameters=true",
+		},
 		{
 			cfg: &Config{
 				User:                     "u",
diff --git a/restful_test.go b/restful_test.go
index 66b3b30..0eb5fc4 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -22,7 +22,7 @@ func postTestError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[str
 	}, errors.New("failed to run post method")
 }
 
-func postAuthTestError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*http.Response, error) {
+func postAuthTestError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration, _ int) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusOK,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
@@ -43,7 +43,7 @@ func postTestAppBadGatewayError(_ context.Context, _ *snowflakeRestful, _ *url.U
 	}, nil
 }
 
-func postAuthTestAppBadGatewayError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*http.Response, error) {
+func postAuthTestAppBadGatewayError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration, _ int) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusBadGateway,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
@@ -57,14 +57,14 @@ func postTestAppForbiddenError(_ context.Context, _ *snowflakeRestful, _ *url.UR
 	}, nil
 }
 
-func postAuthTestAppForbiddenError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*http.Response, error) {
+func postAuthTestAppForbiddenError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration, _ int) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusForbidden,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
 	}, nil
 }
 
-func postAuthTestAppUnexpectedError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*http.Response, error) {
+func postAuthTestAppUnexpectedError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration, _ int) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusInsufficientStorage,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
@@ -110,7 +110,7 @@ func postTestRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[str
 	}, nil
 }
 
-func postAuthTestAfterRenew(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration) (*http.Response, error) {
+func postAuthTestAfterRenew(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ bodyCreatorType, _ time.Duration, _ int) (*http.Response, error) {
 	dd := &execResponseData{}
 	er := &execResponse{
 		Data:    *dd,
diff --git a/retry_test.go b/retry_test.go
index 7a089ec..5815d66 100644
--- a/retry_test.go
+++ b/retry_test.go
@@ -223,20 +223,14 @@ func TestRetryQuerySuccess(t *testing.T) {
 		t: t,
 	}
 	urlPtr, err := url.Parse("https://fakeaccountretrysuccess.snowflakecomputing.com:443/queries/v1/query-request?" + requestIDKey + "=testid")
-	if err != nil {
-		t.Fatal("failed to parse the test URL")
-	}
+	assertNilF(t, err, "failed to parse the test URL")
 	_, err = newRetryHTTP(context.Background(),
 		client,
-		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, constTimeProvider(123456), &Config{IncludeRetryReason: ConfigBoolTrue}).doPost().setBody([]byte{0}).execute()
-	if err != nil {
-		t.Fatal("failed to run retry")
-	}
+		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, 3, constTimeProvider(123456), &Config{IncludeRetryReason: ConfigBoolTrue}).doPost().setBody([]byte{0}).execute()
+	assertNilF(t, err, "failed to run retry")
 	var values url.Values
 	values, err = url.ParseQuery(urlPtr.RawQuery)
-	if err != nil {
-		t.Fatal("failed to parse the URL")
-	}
+	assertNilF(t, err, "failed to parse the test URL")
 	retry, err := strconv.Atoi(values.Get(retryCountKey))
 	if err != nil {
 		t.Fatalf("failed to get retry counter: %v", err)
@@ -272,20 +266,14 @@ func TestRetryQuerySuccessWithRetryReasonDisabled(t *testing.T) {
 		t: t,
 	}
 	urlPtr, err := url.Parse("https://fakeaccountretrysuccess.snowflakecomputing.com:443/queries/v1/query-request?" + requestIDKey + "=testid")
-	if err != nil {
-		t.Fatal("failed to parse the test URL")
-	}
+	assertNilF(t, err, "failed to parse the test URL")
 	_, err = newRetryHTTP(context.Background(),
 		client,
-		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, constTimeProvider(123456), &Config{IncludeRetryReason: ConfigBoolFalse}).doPost().setBody([]byte{0}).execute()
-	if err != nil {
-		t.Fatal("failed to run retry")
-	}
+		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, 3, constTimeProvider(123456), &Config{IncludeRetryReason: ConfigBoolFalse}).doPost().setBody([]byte{0}).execute()
+	assertNilF(t, err, "failed to run retry")
 	var values url.Values
 	values, err = url.ParseQuery(urlPtr.RawQuery)
-	if err != nil {
-		t.Fatal("failed to parse the URL")
-	}
+	assertNilF(t, err, "failed to parse the test URL")
 	retry, err := strconv.Atoi(values.Get(retryCountKey))
 	if err != nil {
 		t.Fatalf("failed to get retry counter: %v", err)
@@ -318,20 +306,14 @@ func TestRetryQuerySuccessWithTimeout(t *testing.T) {
 		t: t,
 	}
 	urlPtr, err := url.Parse("https://fakeaccountretrysuccess.snowflakecomputing.com:443/queries/v1/query-request?" + requestIDKey + "=testid")
-	if err != nil {
-		t.Fatal("failed to parse the test URL")
-	}
+	assertNilF(t, err, "failed to parse the test URL")
 	_, err = newRetryHTTP(context.Background(),
 		client,
-		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, constTimeProvider(123456), nil).doPost().setBody([]byte{0}).execute()
-	if err != nil {
-		t.Fatal("failed to run retry")
-	}
+		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, 3, constTimeProvider(123456), nil).doPost().setBody([]byte{0}).execute()
+	assertNilF(t, err, "failed to run retry")
 	var values url.Values
 	values, err = url.ParseQuery(urlPtr.RawQuery)
-	if err != nil {
-		t.Fatal("failed to parse the URL")
-	}
+	assertNilF(t, err, "failed to parse the test URL")
 	retry, err := strconv.Atoi(values.Get(retryCountKey))
 	if err != nil {
 		t.Fatalf("failed to get retry counter: %v", err)
@@ -341,33 +323,52 @@ func TestRetryQuerySuccessWithTimeout(t *testing.T) {
 	}
 }
 
-func TestRetryQueryFail(t *testing.T) {
+func TestRetryQueryFailWithTimeout(t *testing.T) {
 	logger.Info("Retry N times until there is a timeout and Fail")
 	client := &fakeHTTPClient{
 		statusCode: http.StatusTooManyRequests,
 		success:    false,
 	}
 	urlPtr, err := url.Parse("https://fakeaccountretryfail.snowflakecomputing.com:443/queries/v1/query-request?" + requestIDKey)
-	if err != nil {
-		t.Fatal("failed to parse the test URL")
-	}
+	assertNilF(t, err, "failed to parse the test URL")
 	_, err = newRetryHTTP(context.Background(),
 		client,
-		emptyRequest, urlPtr, make(map[string]string), 15*time.Second, defaultTimeProvider, nil).doPost().setBody([]byte{0}).execute()
-	if err == nil {
-		t.Fatal("should fail to run retry")
+		emptyRequest, urlPtr, make(map[string]string), 15*time.Second, 100, defaultTimeProvider, nil).doPost().setBody([]byte{0}).execute()
+	assertNotNilF(t, err, "should fail to run retry")
+	var values url.Values
+	values, err = url.ParseQuery(urlPtr.RawQuery)
+	assertNilF(t, err, fmt.Sprintf("failed to parse the URL: %v", err))
+	retry, err := strconv.Atoi(values.Get(retryCountKey))
+	assertNilF(t, err, fmt.Sprintf("failed to get retry counter: %v", err))
+	if retry < 2 {
+		t.Fatalf("not enough retries: %v", retry)
 	}
+}
+
+func TestRetryQueryFailWithMaxRetryCount(t *testing.T) {
+	maxRetryCount := 3
+	logger.Info("Retry 3 times until retry reaches MaxRetryCount and Fail")
+	client := &fakeHTTPClient{
+		statusCode: http.StatusTooManyRequests,
+		success:    false,
+	}
+	urlPtr, err := url.Parse("https://fakeaccountretryfail.snowflakecomputing.com:443/queries/v1/query-request?" + requestIDKey)
+	assertNilF(t, err, "failed to parse the test URL")
+	_, err = newRetryHTTP(context.Background(),
+		client,
+		emptyRequest, urlPtr, make(map[string]string), 15*time.Hour, maxRetryCount, defaultTimeProvider, nil).doPost().setBody([]byte{0}).execute()
+	assertNotNilF(t, err, "should fail to run retry")
 	var values url.Values
 	values, err = url.ParseQuery(urlPtr.RawQuery)
 	if err != nil {
 		t.Fatalf("failed to parse the URL: %v", err)
 	}
-	retry, err := strconv.Atoi(values.Get(retryCountKey))
+	retryCount, err := strconv.Atoi(values.Get(retryCountKey))
 	if err != nil {
 		t.Fatalf("failed to get retry counter: %v", err)
 	}
-	if retry < 2 {
-		t.Fatalf("not enough retry counter: %v", retry)
+	if retryCount < 3 {
+		t.Fatalf("not enough retries: %v; expected %v", retryCount, maxRetryCount)
 	}
 }
 
@@ -394,35 +395,26 @@ func TestRetryLoginRequest(t *testing.T) {
 		},
 	}
 	urlPtr, err := url.Parse("https://fakeaccountretrylogin.snowflakecomputing.com:443/login-request?request_id=testid")
-	if err != nil {
-		t.Fatal("failed to parse the test URL")
-	}
+	assertNilF(t, err, "failed to parse the test URL")
 	_, err = newRetryHTTP(context.Background(),
 		client,
-		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, defaultTimeProvider, nil).doPost().setBody([]byte{0}).execute()
-	if err != nil {
-		t.Fatal("failed to run retry")
-	}
+		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, 3, defaultTimeProvider, nil).doPost().setBody([]byte{0}).execute()
+	assertNilF(t, err, "failed to run retry")
 	var values url.Values
 	values, err = url.ParseQuery(urlPtr.RawQuery)
-	if err != nil {
-		t.Fatalf("failed to parse the URL: %v", err)
-	}
+	assertNilF(t, err, "failed to parse the test URL")
 	if values.Get(retryCountKey) != "" {
 		t.Fatalf("no retry counter should be attached: %v", retryCountKey)
 	}
 	logger.Info("Retry N times for timeouts and Fail")
 	client = &fakeHTTPClient{
-		cnt:     10,
 		success: false,
 		timeout: true,
 	}
 	_, err = newRetryHTTP(context.Background(),
 		client,
-		emptyRequest, urlPtr, make(map[string]string), 10*time.Second, defaultTimeProvider, nil).doPost().setBody([]byte{0}).execute()
-	if err == nil {
-		t.Fatal("should fail to run retry")
-	}
+		emptyRequest, urlPtr, make(map[string]string), 5*time.Second, 3, defaultTimeProvider, nil).doPost().setBody([]byte{0}).execute()
+	assertNotNilF(t, err, "should fail to run retry")
 	values, err = url.ParseQuery(urlPtr.RawQuery)
 	if err != nil {
 		t.Fatalf("failed to parse the URL: %v", err)
@@ -440,9 +432,7 @@ func TestRetryAuthLoginRequest(t *testing.T) {
 		timeout: true,
 	}
 	urlPtr, err := url.Parse("https://fakeaccountretrylogin.snowflakecomputing.com:443/login-request?request_id=testid")
-	if err != nil {
-		t.Fatal("failed to parse the test URL")
-	}
+	assertNilF(t, err, "failed to parse the test URL")
 	execID := 0
 	bodyCreator := func() ([]byte, error) {
 		execID++
@@ -450,10 +440,8 @@ func TestRetryAuthLoginRequest(t *testing.T) {
 	}
 	_, err = newRetryHTTP(context.Background(),
 		client,
-		http.NewRequest, urlPtr, make(map[string]string), 60*time.Second, defaultTimeProvider, nil).doPost().setBodyCreator(bodyCreator).execute()
-	if err != nil {
-		t.Fatal("failed to run retry")
-	}
+		http.NewRequest, urlPtr, make(map[string]string), 60*time.Second, 3, defaultTimeProvider, nil).doPost().setBodyCreator(bodyCreator).execute()
+	assertNilF(t, err, "failed to run retry")
 	if lastReqBody := string(client.reqBody); lastReqBody != "execID: 3" {
 		t.Fatalf("body should be updated on each request, expected: execID: 3, last body: %v", lastReqBody)
 	}
@@ -466,20 +454,16 @@ func TestLoginRetry429(t *testing.T) {
 		statusCode: 429,
 	}
 	urlPtr, err := url.Parse("https://fakeaccountretrylogin.snowflakecomputing.com:443/login-request?request_id=testid")
-	if err != nil {
-		t.Fatal("failed to parse the test URL")
-	}
+	assertNilF(t, err, "failed to parse the test URL")
+
 	_, err = newRetryHTTP(context.Background(),
 		client,
-		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, defaultTimeProvider, nil).doPost().setBody([]byte{0}).execute() // enable doRaise4XXX
-	if err != nil {
-		t.Fatal("failed to run retry")
-	}
+		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, 3, defaultTimeProvider, nil).doPost().setBody([]byte{0}).execute() // enable doRaise4XXX
+	assertNilF(t, err, "failed to run retry")
+
 	var values url.Values
 	values, err = url.ParseQuery(urlPtr.RawQuery)
-	if err != nil {
-		t.Fatalf("failed to parse the URL: %v", err)
-	}
+	assertNilF(t, err, fmt.Sprintf("failed to parse the URL: %v", err))
 	if values.Get(retryCountKey) != "" {
 		t.Fatalf("no retry counter should be attached: %v", retryCountKey)
 	}