Repository: snowflakedb/gosnowflake
Pull request #913 (2023-09-13 09:40:35)

SNOW-736353: Implement retry reason
### Description
Introducing retry reason and retry count parameters.

### Checklist
- [ ] Code compiles correctly
- [ ] Run ``make fmt`` to fix inconsistent formats
- [ ] 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



Patch:
diff --git a/async.go b/async.go
index d29b24b..47da1b7 100644
--- a/async.go
+++ b/async.go
@@ -105,7 +105,7 @@ func (sr *snowflakeRestful) getAsync(
 
 	}
 
-	sc := &snowflakeConn{rest: sr, cfg: cfg, queryContextCache: (&queryContextCache{}).init()}
+	sc := &snowflakeConn{rest: sr, cfg: cfg, queryContextCache: (&queryContextCache{}).init(), currentTimeProvider: defaultTimeProvider}
 	if respd.Success {
 		if resType == execResultType {
 			res.insertID = -1
diff --git a/authokta.go b/authokta.go
index 1706f2d..994b51c 100644
--- a/authokta.go
+++ b/authokta.go
@@ -216,7 +216,7 @@ func postAuthSAML(
 	fullURL := sr.getFullURL(authenticatorRequestPath, params)
 
 	logger.Infof("fullURL: %v", fullURL)
-	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, true)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, true, defaultTimeProvider)
 	if err != nil {
 		return nil, err
 	}
@@ -274,7 +274,7 @@ func postAuthOKTA(
 	if err != nil {
 		return nil, err
 	}
-	resp, err := sr.FuncPost(ctx, sr, targetURL, headers, body, timeout, false)
+	resp, err := sr.FuncPost(ctx, sr, targetURL, headers, body, timeout, false, defaultTimeProvider)
 	if err != nil {
 		return nil, err
 	}
diff --git a/chunk_downloader.go b/chunk_downloader.go
index f715992..1a298cf 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).execute()
+	return newRetryHTTP(ctx, sc.rest.Client, http.NewRequest, u, headers, timeout, sc.currentTimeProvider).execute()
 }
 
 func (scd *snowflakeChunkDownloader) startArrowBatches() error {
@@ -636,7 +636,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).execute()
+	res, err := newRetryHTTP(context.Background(), f.client, http.NewRequest, fullURL, f.headers, 0, defaultTimeProvider).execute()
 	if err != nil {
 		return err
 	}
diff --git a/client.go b/client.go
index 358be47..a5d7f74 100644
--- a/client.go
+++ b/client.go
@@ -12,7 +12,7 @@ import (
 // InternalClient is implemented by HTTPClient
 type InternalClient interface {
 	Get(context.Context, *url.URL, map[string]string, time.Duration) (*http.Response, error)
-	Post(context.Context, *url.URL, map[string]string, []byte, time.Duration, bool) (*http.Response, error)
+	Post(context.Context, *url.URL, map[string]string, []byte, time.Duration, bool, currentTimeProvider) (*http.Response, error)
 }
 
 type httpClient struct {
@@ -33,6 +33,7 @@ func (cli *httpClient) Post(
 	headers map[string]string,
 	body []byte,
 	timeout time.Duration,
-	raise4xx bool) (*http.Response, error) {
-	return cli.sr.FuncPost(ctx, cli.sr, url, headers, body, timeout, raise4xx)
+	raise4xx bool,
+	currentTimeProvider currentTimeProvider) (*http.Response, error) {
+	return cli.sr.FuncPost(ctx, cli.sr, url, headers, body, timeout, raise4xx, currentTimeProvider)
 }
diff --git a/connection.go b/connection.go
index 7f716df..769253d 100644
--- a/connection.go
+++ b/connection.go
@@ -61,13 +61,14 @@ const (
 const privateLinkSuffix = "privatelink.snowflakecomputing.com"
 
 type snowflakeConn struct {
-	ctx               context.Context
-	cfg               *Config
-	rest              *snowflakeRestful
-	SequenceCounter   uint64
-	telemetry         *snowflakeTelemetry
-	internal          InternalClient
-	queryContextCache *queryContextCache
+	ctx                 context.Context
+	cfg                 *Config
+	rest                *snowflakeRestful
+	SequenceCounter     uint64
+	telemetry           *snowflakeTelemetry
+	internal            InternalClient
+	queryContextCache   *queryContextCache
+	currentTimeProvider currentTimeProvider
 }
 
 var (
@@ -727,10 +728,11 @@ func (scd *snowflakeArrowStreamChunkDownloader) GetBatches() (out []ArrowStreamB
 
 func buildSnowflakeConn(ctx context.Context, config Config) (*snowflakeConn, error) {
 	sc := &snowflakeConn{
-		SequenceCounter:   0,
-		ctx:               ctx,
-		cfg:               &config,
-		queryContextCache: (&queryContextCache{}).init(),
+		SequenceCounter:     0,
+		ctx:                 ctx,
+		cfg:                 &config,
+		queryContextCache:   (&queryContextCache{}).init(),
+		currentTimeProvider: defaultTimeProvider,
 	}
 	var st http.RoundTripper = SnowflakeTransport
 	if sc.cfg.Transporter == nil {
diff --git a/heartbeat.go b/heartbeat.go
index 8f9b020..60c4bf1 100644
--- a/heartbeat.go
+++ b/heartbeat.go
@@ -62,7 +62,7 @@ func (hc *heartbeat) heartbeatMain() error {
 
 	fullURL := hc.restful.getFullURL(heartBeatPath, params)
 	timeout := hc.restful.RequestTimeout
-	resp, err := hc.restful.FuncPost(context.Background(), hc.restful, fullURL, headers, nil, timeout, false)
+	resp, err := hc.restful.FuncPost(context.Background(), hc.restful, fullURL, headers, nil, timeout, false, defaultTimeProvider)
 	if err != nil {
 		return err
 	}
diff --git a/monitoring.go b/monitoring.go
index 8ac4c1d..07b17c0 100644
--- a/monitoring.go
+++ b/monitoring.go
@@ -9,7 +9,6 @@ import (
 	"fmt"
 	"net/url"
 	"strconv"
-	"time"
 )
 
 const urlQueriesResultFmt = "/queries/%s/result"
@@ -208,7 +207,7 @@ func (sc *snowflakeConn) getQueryResultResp(
 	paramsMutex.Unlock()
 	param := make(url.Values)
 	param.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
-	param.Add("clientStartTime", strconv.FormatInt(time.Now().Unix(), 10))
+	param.Add("clientStartTime", strconv.FormatInt(sc.currentTimeProvider.currentTime(), 10))
 	param.Add(requestGUIDKey, NewUUID().String())
 	token, _, _ := sc.rest.TokenAccessor.GetTokens()
 	if token != "" {
diff --git a/ocsp.go b/ocsp.go
index f0a30b9..037f103 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).execute()
+	res, err := newRetryHTTP(ctx, client, req, ocspServerHost, headers, totalTimeout, defaultTimeProvider).execute()
 	if err != nil {
 		logger.Errorf("failed to get OCSP cache from OCSP Cache Server. %v\n", err)
 		return nil, &ocspStatus{
@@ -413,7 +413,7 @@ func retryOCSP(
 	}
 	res, err := newRetryHTTP(
 		ctx, client, req, ocspHost, headers,
-		totalTimeout*time.Duration(multiplier)).doPost().setBody(reqBody).execute()
+		totalTimeout*time.Duration(multiplier), defaultTimeProvider).doPost().setBody(reqBody).execute()
 	if err != nil {
 		return ocspRes, ocspResBytes, &ocspStatus{
 			code: ocspFailedSubmit,
diff --git a/restful.go b/restful.go
index 3278717..34297c1 100644
--- a/restful.go
+++ b/restful.go
@@ -43,7 +43,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, bool) (*http.Response, error)
+	funcPostType     func(context.Context, *snowflakeRestful, *url.URL, map[string]string, []byte, time.Duration, bool, currentTimeProvider) (*http.Response, error)
 	funcAuthPostType func(context.Context, *http.Client, *url.URL, map[string]string, bodyCreatorType, time.Duration, bool) (*http.Response, error)
 	bodyCreatorType  func() ([]byte, error)
 )
@@ -162,9 +162,10 @@ func postRestful(
 	headers map[string]string,
 	body []byte,
 	timeout time.Duration,
-	raise4XX bool) (
+	raise4XX bool,
+	currentTimeProvider currentTimeProvider) (
 	*http.Response, error) {
-	return newRetryHTTP(ctx, sr.Client, http.NewRequest, fullURL, headers, timeout).
+	return newRetryHTTP(ctx, sr.Client, http.NewRequest, fullURL, headers, timeout, currentTimeProvider).
 		doPost().
 		setBody(body).
 		doRaise4XX(raise4XX).
@@ -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).execute()
+	return newRetryHTTP(ctx, sr.Client, http.NewRequest, fullURL, headers, timeout, defaultTimeProvider).execute()
 }
 
 func postAuthRestful(
@@ -190,7 +191,7 @@ func postAuthRestful(
 	timeout time.Duration,
 	raise4XX bool) (
 	*http.Response, error) {
-	return newRetryHTTP(ctx, client, http.NewRequest, fullURL, headers, timeout).
+	return newRetryHTTP(ctx, client, http.NewRequest, fullURL, headers, timeout, defaultTimeProvider).
 		doPost().
 		setBodyCreator(bodyCreator).
 		doRaise4XX(raise4XX).
@@ -233,7 +234,6 @@ func postRestfulQueryHelper(
 	data *execResponse, err error) {
 	logger.Infof("params: %v", params)
 	params.Add(requestIDKey, requestID.String())
-	params.Add("clientStartTime", strconv.FormatInt(time.Now().Unix(), 10))
 	params.Add(requestGUIDKey, NewUUID().String())
 	token, _, _ := sr.TokenAccessor.GetTokens()
 	if token != "" {
@@ -242,7 +242,7 @@ func postRestfulQueryHelper(
 
 	var resp *http.Response
 	fullURL := sr.getFullURL(queryRequestPath, params)
-	resp, err = sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, true)
+	resp, err = sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, true, defaultTimeProvider)
 	if err != nil {
 		return nil, err
 	}
@@ -334,7 +334,7 @@ func closeSession(ctx context.Context, sr *snowflakeRestful, timeout time.Durati
 	token, _, _ := sr.TokenAccessor.GetTokens()
 	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
 
-	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, nil, 5*time.Second, false)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, nil, 5*time.Second, false, defaultTimeProvider)
 	if err != nil {
 		return err
 	}
@@ -393,7 +393,7 @@ func renewRestfulSession(ctx context.Context, sr *snowflakeRestful, timeout time
 		return err
 	}
 
-	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, reqBody, timeout, false)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, reqBody, timeout, false, defaultTimeProvider)
 	if err != nil {
 		return err
 	}
@@ -465,7 +465,7 @@ func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID UUID, time
 		return err
 	}
 
-	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, reqByte, timeout, false)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, reqByte, timeout, false, defaultTimeProvider)
 	if err != nil {
 		return err
 	}
diff --git a/retry.go b/retry.go
index 95efbc7..f465b66 100644
--- a/retry.go
+++ b/retry.go
@@ -27,8 +27,12 @@ func init() {
 const (
 	// requestGUIDKey is attached to every request against Snowflake
 	requestGUIDKey string = "request_guid"
-	// retryCounterKey is attached to query-request from the second time
-	retryCounterKey string = "retryCounter"
+	// retryCountKey is attached to query-request from the second time
+	retryCountKey string = "retryCount"
+	// retryReasonKey contains last HTTP status or 0 if timeout
+	retryReasonKey string = "retryReason"
+	// clientStartTime contains a time when client started request (first request, not retries)
+	clientStartTimeKey string = "clientStartTime"
 	// requestIDKey is attached to all requests to Snowflake
 	requestIDKey string = "requestId"
 )
@@ -86,42 +90,91 @@ func (replacer *requestGUIDReplace) replace() *url.URL {
 	return replacer.urlPtr
 }
 
-type retryCounterUpdater interface {
+type retryCountUpdater interface {
 	replaceOrAdd(retry int) *url.URL
 }
 
-type retryCounterUpdate struct {
+type retryCountUpdate struct {
 	urlPtr    *url.URL
 	urlValues url.Values
 }
 
 // this replacer does nothing but replace the url
-type transientReplaceOrAdd struct {
+type transientRetryCountUpdater struct {
 	urlPtr *url.URL
 }
 
-func (replaceOrAdder *transientReplaceOrAdd) replaceOrAdd(retry int) *url.URL {
+func (replaceOrAdder *transientRetryCountUpdater) replaceOrAdd(retry int) *url.URL {
 	return replaceOrAdder.urlPtr
 }
 
-func (replacer *retryCounterUpdate) replaceOrAdd(retry int) *url.URL {
-	replacer.urlValues.Del(retryCounterKey)
-	replacer.urlValues.Add(retryCounterKey, strconv.Itoa(retry))
+func (replacer *retryCountUpdate) replaceOrAdd(retry int) *url.URL {
+	replacer.urlValues.Del(retryCountKey)
+	replacer.urlValues.Add(retryCountKey, strconv.Itoa(retry))
 	replacer.urlPtr.RawQuery = replacer.urlValues.Encode()
 	return replacer.urlPtr
 }
 
-func newRetryUpdate(urlPtr *url.URL) retryCounterUpdater {
-	if !strings.HasPrefix(urlPtr.Path, queryRequestPath) {
+func newRetryCountUpdater(urlPtr *url.URL) retryCountUpdater {
+	if !isQueryRequest(urlPtr) {
 		// nop if not query-request
-		return &transientReplaceOrAdd{urlPtr}
+		return &transientRetryCountUpdater{urlPtr}
 	}
 	values, err := url.ParseQuery(urlPtr.RawQuery)
 	if err != nil {
 		// nop if the URL is not valid
-		return &transientReplaceOrAdd{urlPtr}
+		return &transientRetryCountUpdater{urlPtr}
 	}
-	return &retryCounterUpdate{urlPtr, values}
+	return &retryCountUpdate{urlPtr, values}
+}
+
+type retryReasonUpdater interface {
+	replaceOrAdd(reason int) *url.URL
+}
+
+type retryReasonUpdate struct {
+	url *url.URL
+}
+
+func (retryReasonUpdater *retryReasonUpdate) replaceOrAdd(reason int) *url.URL {
+	query := retryReasonUpdater.url.Query()
+	query.Del(retryReasonKey)
+	query.Add(retryReasonKey, strconv.Itoa(reason))
+	retryReasonUpdater.url.RawQuery = query.Encode()
+	return retryReasonUpdater.url
+}
+
+type transientRetryReasonUpdater struct {
+	url *url.URL
+}
+
+func (retryReasonUpdater *transientRetryReasonUpdater) replaceOrAdd(_ int) *url.URL {
+	return retryReasonUpdater.url
+}
+
+func newRetryReasonUpdater(url *url.URL) retryReasonUpdater {
+	if !isQueryRequest(url) {
+		return &transientRetryReasonUpdater{url}
+	}
+	return &retryReasonUpdate{url}
+}
+
+func ensureClientStartTimeIsSet(url *url.URL, clientStartTime string) *url.URL {
+	if !isQueryRequest(url) {
+		// nop if not query-request
+		return url
+	}
+	query := url.Query()
+	if query.Has(clientStartTimeKey) {
+		return url
+	}
+	query.Add(clientStartTimeKey, clientStartTime)
+	url.RawQuery = query.Encode()
+	return url
+}
+
+func isQueryRequest(url *url.URL) bool {
+	return strings.HasPrefix(url.Path, queryRequestPath)
 }
 
 type waitAlgo struct {
@@ -161,15 +214,16 @@ type clientInterface interface {
 }
 
 type retryHTTP struct {
-	ctx         context.Context
-	client      clientInterface
-	req         requestFunc
-	method      string
-	fullURL     *url.URL
-	headers     map[string]string
-	bodyCreator bodyCreatorType
-	timeout     time.Duration
-	raise4XX    bool
+	ctx                 context.Context
+	client              clientInterface
+	req                 requestFunc
+	method              string
+	fullURL             *url.URL
+	headers             map[string]string
+	bodyCreator         bodyCreatorType
+	timeout             time.Duration
+	raise4XX            bool
+	currentTimeProvider currentTimeProvider
 }
 
 func newRetryHTTP(ctx context.Context,
@@ -177,7 +231,8 @@ func newRetryHTTP(ctx context.Context,
 	req requestFunc,
 	fullURL *url.URL,
 	headers map[string]string,
-	timeout time.Duration) *retryHTTP {
+	timeout time.Duration,
+	currentTimeProvider currentTimeProvider) *retryHTTP {
 	instance := retryHTTP{}
 	instance.ctx = ctx
 	instance.client = client
@@ -188,6 +243,7 @@ func newRetryHTTP(ctx context.Context,
 	instance.timeout = timeout
 	instance.bodyCreator = emptyBodyCreator
 	instance.raise4XX = false
+	instance.currentTimeProvider = currentTimeProvider
 	return &instance
 }
 
@@ -218,9 +274,11 @@ func (r *retryHTTP) execute() (res *http.Response, err error) {
 	logger.WithContext(r.ctx).Infof("retryHTTP.totalTimeout: %v", totalTimeout)
 	retryCounter := 0
 	sleepTime := time.Duration(0)
+	clientStartTime := strconv.FormatInt(r.currentTimeProvider.currentTime(), 10)
 
-	var rIDReplacer requestGUIDReplacer
-	var rUpdater retryCounterUpdater
+	var requestGUIDReplacer requestGUIDReplacer
+	var retryCountUpdater retryCountUpdater
+	var retryReasonUpdater retryReasonUpdater
 
 	for {
 		logger.Debugf("retry count: %v", retryCounter)
@@ -279,15 +337,25 @@ func (r *retryHTTP) execute() (res *http.Response, err error) {
 			}
 		}
 		retryCounter++
-		if rIDReplacer == nil {
-			rIDReplacer = newRequestGUIDReplace(r.fullURL)
+		if requestGUIDReplacer == nil {
+			requestGUIDReplacer = newRequestGUIDReplace(r.fullURL)
+		}
+		r.fullURL = requestGUIDReplacer.replace()
+		if retryCountUpdater == nil {
+			retryCountUpdater = newRetryCountUpdater(r.fullURL)
+		}
+		r.fullURL = retryCountUpdater.replaceOrAdd(retryCounter)
+		if retryReasonUpdater == nil {
+			retryReasonUpdater = newRetryReasonUpdater(r.fullURL)
 		}
-		r.fullURL = rIDReplacer.replace()
-		if rUpdater == nil {
-			rUpdater = newRetryUpdate(r.fullURL)
+		retryReason := 0
+		if res != nil {
+			retryReason = res.StatusCode
 		}
-		r.fullURL = rUpdater.replaceOrAdd(retryCounter)
+		r.fullURL = retryReasonUpdater.replaceOrAdd(retryReason)
+		r.fullURL = ensureClientStartTimeIsSet(r.fullURL, clientStartTime)
 		logger.WithContext(r.ctx).Infof("sleeping %v. to timeout: %v. retrying", sleepTime, totalTimeout)
+		logger.WithContext(r.ctx).Infof("retry count: %v, retry reason: %v", retryCounter, retryReason)
 
 		await := time.NewTimer(sleepTime)
 		select {
diff --git a/telemetry.go b/telemetry.go
index 542b8a0..1adb085 100644
--- a/telemetry.go
+++ b/telemetry.go
@@ -97,7 +97,7 @@ func (st *snowflakeTelemetry) sendBatch() error {
 	}
 	resp, err := st.sr.FuncPost(context.Background(), st.sr,
 		st.sr.getFullURL(telemetryPath, nil), headers, body,
-		defaultTelemetryTimeout, true)
+		defaultTelemetryTimeout, true, defaultTimeProvider)
 	if err != nil {
 		logger.Info("failed to upload metrics to telemetry. err: %v", err)
 		return err
diff --git a/util.go b/util.go
index d112189..ade1093 100644
--- a/util.go
+++ b/util.go
@@ -36,6 +36,10 @@ const (
 	streamChunkDownload contextKey = "STREAM_CHUNK_DOWNLOAD"
 )
 
+var (
+	defaultTimeProvider = &unixTimeProvider{}
+)
+
 // WithMultiStatement returns a context that allows the user to execute the desired number of sql queries in one query
 func WithMultiStatement(ctx context.Context, num int) (context.Context, error) {
 	return context.WithValue(ctx, multiStatementCount, num), nil
@@ -235,3 +239,14 @@ func GetFromEnv(name string, failOnMissing bool) (string, error) {
 	}
 	return "", nil
 }
+
+type currentTimeProvider interface {
+	currentTime() int64
+}
+
+type unixTimeProvider struct {
+}
+
+func (utp *unixTimeProvider) currentTime() int64 {
+	return time.Now().UnixMilli()
+}


Test patch:
diff --git a/client_test.go b/client_test.go
index b8e582a..6ef66ef 100644
--- a/client_test.go
+++ b/client_test.go
@@ -48,7 +48,7 @@ func TestInternalClient(t *testing.T) {
 		t.Fatalf("Expected exactly one GET request, got %v", transport.getRequests)
 	}
 
-	resp, err = internalClient.Post(context.Background(), &url.URL{}, make(map[string]string), make([]byte, 0), 0, false)
+	resp, err = internalClient.Post(context.Background(), &url.URL{}, make(map[string]string), make([]byte, 0), 0, false, defaultTimeProvider)
 	if err != nil || resp.StatusCode != 200 {
 		t.Fail()
 	}
diff --git a/connection_test.go b/connection_test.go
index 0aa2e90..0a23684 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -130,8 +130,9 @@ func TestGetQueryResultUsesTokenFromTokenAccessor(t *testing.T) {
 		TokenAccessor: ta,
 	}
 	sc := &snowflakeConn{
-		cfg:  &Config{Params: map[string]*string{}},
-		rest: sr,
+		cfg:                 &Config{Params: map[string]*string{}},
+		rest:                sr,
+		currentTimeProvider: defaultTimeProvider,
 	}
 	if _, err := sc.getQueryResultResp(context.Background(), ""); err != nil {
 		t.Fatalf("err: %v", err)
diff --git a/ocsp_test.go b/ocsp_test.go
index 94c2dc8..257e879 100644
--- a/ocsp_test.go
+++ b/ocsp_test.go
@@ -319,7 +319,7 @@ func TestOCSPRetry(t *testing.T) {
 	}
 	res, b, st := retryOCSP(
 		context.TODO(),
-		client, fakeRequestFunc,
+		client, emptyRequest,
 		dummyOCSPHost,
 		make(map[string]string), []byte{0}, certs[len(certs)-1], 10*time.Second)
 	if st.err == nil {
diff --git a/restful_test.go b/restful_test.go
index ec27213..f23cacd 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -15,7 +15,7 @@ import (
 	"time"
 )
 
-func postTestError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+func postTestError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusOK,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
@@ -29,14 +29,14 @@ func postAuthTestError(_ context.Context, _ *http.Client, _ *url.URL, _ map[stri
 	}, errors.New("failed to run post method")
 }
 
-func postTestSuccessButInvalidJSON(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+func postTestSuccessButInvalidJSON(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusOK,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
 	}, nil
 }
 
-func postTestAppBadGatewayError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+func postTestAppBadGatewayError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusBadGateway,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
@@ -50,7 +50,7 @@ func postAuthTestAppBadGatewayError(_ context.Context, _ *http.Client, _ *url.UR
 	}, nil
 }
 
-func postTestAppForbiddenError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+func postTestAppForbiddenError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusForbidden,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
@@ -71,7 +71,7 @@ func postAuthTestAppUnexpectedError(_ context.Context, _ *http.Client, _ *url.UR
 	}, nil
 }
 
-func postTestQueryNotExecuting(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+func postTestQueryNotExecuting(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider) (*http.Response, error) {
 	dd := &execResponseData{}
 	er := &execResponse{
 		Data:    *dd,
@@ -90,7 +90,7 @@ func postTestQueryNotExecuting(_ context.Context, _ *snowflakeRestful, _ *url.UR
 	}, nil
 }
 
-func postTestRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+func postTestRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider) (*http.Response, error) {
 	dd := &execResponseData{}
 	er := &execResponse{
 		Data:    *dd,
@@ -130,7 +130,7 @@ func postAuthTestAfterRenew(_ context.Context, _ *http.Client, _ *url.URL, _ map
 	}, nil
 }
 
-func postTestAfterRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+func postTestAfterRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider) (*http.Response, error) {
 	dd := &execResponseData{}
 	er := &execResponse{
 		Data:    *dd,
@@ -157,7 +157,7 @@ func cancelTestRetry(ctx context.Context, sr *snowflakeRestful, requestID UUID,
 	if err != nil {
 		return err
 	}
-	resp, err := sr.FuncPost(ctx, sr, &u, getHeaders(), reqByte, timeout, false)
+	resp, err := sr.FuncPost(ctx, sr, &u, getHeaders(), reqByte, timeout, false, defaultTimeProvider)
 	if err != nil {
 		return err
 	}
@@ -462,7 +462,7 @@ func TestUnitRenewRestfulSession(t *testing.T) {
 	accessor := getSimpleTokenAccessor()
 	oldToken, oldMasterToken, oldSessionID := "oldtoken", "oldmaster", int64(100)
 	newToken, newMasterToken, newSessionID := "newtoken", "newmaster", int64(200)
-	postTestSuccessWithNewTokens := func(_ context.Context, _ *snowflakeRestful, _ *url.URL, headers map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+	postTestSuccessWithNewTokens := func(_ context.Context, _ *snowflakeRestful, _ *url.URL, headers map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider) (*http.Response, error) {
 		if headers[headerAuthorizationKey] != fmt.Sprintf(headerSnowflakeToken, oldMasterToken) {
 			t.Fatalf("authorization key doesn't match, %v vs %v", headers[headerAuthorizationKey], fmt.Sprintf(headerSnowflakeToken, oldMasterToken))
 		}
diff --git a/retry_test.go b/retry_test.go
index 5acb380..e8590a2 100644
--- a/retry_test.go
+++ b/retry_test.go
@@ -19,6 +19,10 @@ func fakeRequestFunc(_, _ string, _ io.Reader) (*http.Request, error) {
 	return nil, nil
 }
 
+func emptyRequest(method string, urlStr string, body io.Reader) (*http.Request, error) {
+	return http.NewRequest(method, urlStr, body)
+}
+
 type fakeHTTPError struct {
 	err     string
 	timeout bool
@@ -48,21 +52,39 @@ func (b *fakeResponseBody) Close() error {
 }
 
 type fakeHTTPClient struct {
-	cnt        int    // number of retry
-	success    bool   // return success after retry in cnt times
-	timeout    bool   // timeout
-	body       []byte // return body
-	reqBody    []byte // last request body
-	statusCode int    // status code
+	t                   *testing.T                // for assertions
+	cnt                 int                       // number of retry
+	success             bool                      // return success after retry in cnt times
+	timeout             bool                      // timeout
+	body                []byte                    // return body
+	reqBody             []byte                    // last request body
+	statusCode          int                       // status code
+	retryNumber         int                       // consecutive number of  retries
+	expectedQueryParams map[int]map[string]string // expected query params per each retry (0-based)
 }
 
 func (c *fakeHTTPClient) Do(req *http.Request) (*http.Response, error) {
+	defer func() {
+		c.retryNumber++
+	}()
 	if req != nil {
 		buf := new(bytes.Buffer)
 		buf.ReadFrom(req.Body)
 		c.reqBody = buf.Bytes()
 	}
 
+	if len(c.expectedQueryParams) > 0 {
+		expectedQueryParams, ok := c.expectedQueryParams[c.retryNumber]
+		if ok {
+			for queryParamName, expectedValue := range expectedQueryParams {
+				actualValue := req.URL.Query().Get(queryParamName)
+				if actualValue != expectedValue {
+					c.t.Fatalf("expected query param %v to be %v, got %v", queryParamName, expectedValue, actualValue)
+				}
+			}
+		}
+	}
+
 	c.cnt--
 	if c.cnt < 0 {
 		c.cnt = 0
@@ -175,18 +197,83 @@ func TestRequestGUID(t *testing.T) {
 }
 
 func TestRetryQuerySuccess(t *testing.T) {
+	logger.Info("Retry N times and Success")
+	client := &fakeHTTPClient{
+		cnt:        3,
+		success:    true,
+		statusCode: 429,
+		expectedQueryParams: map[int]map[string]string{
+			0: {
+				"retryCount":      "",
+				"retryReason":     "",
+				"clientStartTime": "",
+			},
+			1: {
+				"retryCount":      "1",
+				"retryReason":     "429",
+				"clientStartTime": "123456",
+			},
+			2: {
+				"retryCount":      "2",
+				"retryReason":     "429",
+				"clientStartTime": "123456",
+			},
+		},
+		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")
+	}
+	_, err = newRetryHTTP(context.TODO(),
+		client,
+		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, constTimeProvider(123456)).doPost().setBody([]byte{0}).execute()
+	if err != nil {
+		t.Fatal("failed to run retry")
+	}
+	var values url.Values
+	values, err = url.ParseQuery(urlPtr.RawQuery)
+	if err != nil {
+		t.Fatal("failed to fail to parse the URL")
+	}
+	retry, 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)
+	}
+}
+
+func TestRetryQuerySuccessWithTimeout(t *testing.T) {
 	logger.Info("Retry N times and Success")
 	client := &fakeHTTPClient{
 		cnt:     3,
 		success: true,
-	}
-	urlPtr, err := url.Parse("https://fakeaccountretrysuccess.snowflakecomputing.com:443/queries/v1/query-request?" + requestIDKey + "=testid&clientStartTime=123456")
+		timeout: true,
+		expectedQueryParams: map[int]map[string]string{
+			0: {
+				"retryCount":  "",
+				"retryReason": "",
+			},
+			1: {
+				"retryCount":  "1",
+				"retryReason": "0",
+			},
+			2: {
+				"retryCount":  "2",
+				"retryReason": "0",
+			},
+		},
+		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")
 	}
 	_, err = newRetryHTTP(context.TODO(),
 		client,
-		fakeRequestFunc, urlPtr, make(map[string]string), 60*time.Second).doPost().setBody([]byte{0}).execute()
+		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, constTimeProvider(123456)).doPost().setBody([]byte{0}).execute()
 	if err != nil {
 		t.Fatal("failed to run retry")
 	}
@@ -195,7 +282,7 @@ func TestRetryQuerySuccess(t *testing.T) {
 	if err != nil {
 		t.Fatal("failed to fail to parse the URL")
 	}
-	retry, err := strconv.Atoi(values.Get(retryCounterKey))
+	retry, err := strconv.Atoi(values.Get(retryCountKey))
 	if err != nil {
 		t.Fatalf("failed to get retry counter: %v", err)
 	}
@@ -203,19 +290,20 @@ func TestRetryQuerySuccess(t *testing.T) {
 		t.Fatalf("not enough retry counter: %v", retry)
 	}
 }
+
 func TestRetryQueryFail(t *testing.T) {
 	logger.Info("Retry N times and Fail")
 	client := &fakeHTTPClient{
 		cnt:     4,
 		success: false,
 	}
-	urlPtr, err := url.Parse("https://fakeaccountretryfail.snowflakecomputing.com:443/queries/v1/query-request?" + requestIDKey + "=testid&clientStartTime=123456")
+	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")
 	}
 	_, err = newRetryHTTP(context.TODO(),
 		client,
-		fakeRequestFunc, urlPtr, make(map[string]string), 60*time.Second).doPost().setBody([]byte{0}).execute()
+		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, defaultTimeProvider).doPost().setBody([]byte{0}).execute()
 	if err == nil {
 		t.Fatal("should fail to run retry")
 	}
@@ -224,7 +312,7 @@ func TestRetryQueryFail(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to fail to parse the URL: %v", err)
 	}
-	retry, err := strconv.Atoi(values.Get(retryCounterKey))
+	retry, err := strconv.Atoi(values.Get(retryCountKey))
 	if err != nil {
 		t.Fatalf("failed to get retry counter: %v", err)
 	}
@@ -232,12 +320,28 @@ func TestRetryQueryFail(t *testing.T) {
 		t.Fatalf("not enough retry counter: %v", retry)
 	}
 }
+
 func TestRetryLoginRequest(t *testing.T) {
 	logger.Info("Retry N times for timeouts and Success")
 	client := &fakeHTTPClient{
 		cnt:     3,
 		success: true,
 		timeout: true,
+		t:       t,
+		expectedQueryParams: map[int]map[string]string{
+			0: {
+				"retryCount":  "",
+				"retryReason": "",
+			},
+			1: {
+				"retryCount":  "",
+				"retryReason": "",
+			},
+			2: {
+				"retryCount":  "",
+				"retryReason": "",
+			},
+		},
 	}
 	urlPtr, err := url.Parse("https://fakeaccountretrylogin.snowflakecomputing.com:443/login-request?request_id=testid")
 	if err != nil {
@@ -245,7 +349,7 @@ func TestRetryLoginRequest(t *testing.T) {
 	}
 	_, err = newRetryHTTP(context.TODO(),
 		client,
-		fakeRequestFunc, urlPtr, make(map[string]string), 60*time.Second).doPost().setBody([]byte{0}).execute()
+		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, defaultTimeProvider).doPost().setBody([]byte{0}).execute()
 	if err != nil {
 		t.Fatal("failed to run retry")
 	}
@@ -254,8 +358,8 @@ func TestRetryLoginRequest(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to fail to parse the URL: %v", err)
 	}
-	if values.Get(retryCounterKey) != "" {
-		t.Fatalf("no retry counter should be attached: %v", retryCounterKey)
+	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{
@@ -265,7 +369,7 @@ func TestRetryLoginRequest(t *testing.T) {
 	}
 	_, err = newRetryHTTP(context.TODO(),
 		client,
-		fakeRequestFunc, urlPtr, make(map[string]string), 10*time.Second).doPost().setBody([]byte{0}).execute()
+		emptyRequest, urlPtr, make(map[string]string), 10*time.Second, defaultTimeProvider).doPost().setBody([]byte{0}).execute()
 	if err == nil {
 		t.Fatal("should fail to run retry")
 	}
@@ -273,8 +377,8 @@ func TestRetryLoginRequest(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to fail to parse the URL: %v", err)
 	}
-	if values.Get(retryCounterKey) != "" {
-		t.Fatalf("no retry counter should be attached: %v", retryCounterKey)
+	if values.Get(retryCountKey) != "" {
+		t.Fatalf("no retry counter should be attached: %v", retryCountKey)
 	}
 }
 
@@ -296,7 +400,7 @@ func TestRetryAuthLoginRequest(t *testing.T) {
 	}
 	_, err = newRetryHTTP(context.TODO(),
 		client,
-		http.NewRequest, urlPtr, make(map[string]string), 60*time.Second).doPost().setBodyCreator(bodyCreator).execute()
+		http.NewRequest, urlPtr, make(map[string]string), 60*time.Second, defaultTimeProvider).doPost().setBodyCreator(bodyCreator).execute()
 	if err != nil {
 		t.Fatal("failed to run retry")
 	}
@@ -317,7 +421,7 @@ func TestLoginRetry429(t *testing.T) {
 	}
 	_, err = newRetryHTTP(context.TODO(),
 		client,
-		fakeRequestFunc, urlPtr, make(map[string]string), 60*time.Second).doRaise4XX(true).doPost().setBody([]byte{0}).execute() // enable doRaise4XXX
+		emptyRequest, urlPtr, make(map[string]string), 60*time.Second, defaultTimeProvider).doRaise4XX(true).doPost().setBody([]byte{0}).execute() // enable doRaise4XXX
 	if err != nil {
 		t.Fatal("failed to run retry")
 	}
@@ -326,7 +430,7 @@ func TestLoginRetry429(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to fail to parse the URL: %v", err)
 	}
-	if values.Get(retryCounterKey) != "" {
-		t.Fatalf("no retry counter should be attached: %v", retryCounterKey)
+	if values.Get(retryCountKey) != "" {
+		t.Fatalf("no retry counter should be attached: %v", retryCountKey)
 	}
 }
diff --git a/telemetry_test.go b/telemetry_test.go
index c738de7..c6a7436 100644
--- a/telemetry_test.go
+++ b/telemetry_test.go
@@ -107,7 +107,7 @@ func TestEnableTelemetry(t *testing.T) {
 	})
 }
 
-func funcPostTelemetryRespFail(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+func funcPostTelemetryRespFail(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool, _ currentTimeProvider) (*http.Response, error) {
 	return nil, errors.New("failed to upload metrics to telemetry")
 }
 
diff --git a/util_test.go b/util_test.go
index 89a86f1..2d8576f 100644
--- a/util_test.go
+++ b/util_test.go
@@ -24,6 +24,18 @@ type tcUUID struct {
 	uuid string
 }
 
+type constTypeProvider struct {
+	constTime int64
+}
+
+func (ctp *constTypeProvider) currentTime() int64 {
+	return ctp.constTime
+}
+
+func constTimeProvider(constTime int64) *constTypeProvider {
+	return &constTypeProvider{constTime: constTime}
+}
+
 func TestSimpleTokenAccessor(t *testing.T) {
 	accessor := getSimpleTokenAccessor()
 	token, masterToken, sessionID := accessor.GetTokens()