Repository: snowflakedb/gosnowflake
Pull request #1224 (2024-10-22 10:25:36)

SNOW-1754426 Fix duplicated requestId and request_guid params
### Description

SNOW-1754426 While migrating from Nginx to Envoy it was observed that when session expired and request was retried after renewal, `requestId` and `request_guid` params where duplicated. The reason was that headers were added to existing value list instead of being replaced.

### Checklist
- [x] Created tests which fail without the change (if possible)
- [x] Extended the README / documentation, if necessary


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1224?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 82.79%. Comparing base [(`251a6cd`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/251a6cdcb10eccd696b0d1b3439f63cee48041d2?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) to head [(`6c5907f`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/6c5907f8b3fc6beaeafa3f3914c1d14f7050cc10?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #1224   +/-   ##
=======================================
  Coverage   82.79%   82.79%           
=======================================
  Files          55       55           
  Lines       13194    13194           
=======================================
  Hits        10924    10924           
  Misses       2270     2270           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1224?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?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 3839e78..c59913d 100644
--- a/auth.go
+++ b/auth.go
@@ -223,8 +223,8 @@ func postAuth(
 	bodyCreator bodyCreatorType,
 	timeout time.Duration) (
 	data *authResponse, err error) {
-	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
-	params.Add(requestGUIDKey, NewUUID().String())
+	params.Set(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
+	params.Set(requestGUIDKey, NewUUID().String())
 
 	fullURL := sr.getFullURL(loginRequestPath, params)
 	logger.WithContext(ctx).Infof("full URL: %v", fullURL)
diff --git a/authokta.go b/authokta.go
index a2dfed8..681a09c 100644
--- a/authokta.go
+++ b/authokta.go
@@ -215,7 +215,7 @@ func postAuthSAML(
 	data *authResponse, err error) {
 
 	params := &url.Values{}
-	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
+	params.Set(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
 	fullURL := sr.getFullURL(authenticatorRequestPath, params)
 
 	logger.WithContext(ctx).Infof("fullURL: %v", fullURL)
diff --git a/heartbeat.go b/heartbeat.go
index 4b40ef6..81ba9f2 100644
--- a/heartbeat.go
+++ b/heartbeat.go
@@ -54,8 +54,8 @@ func (hc *heartbeat) stop() {
 func (hc *heartbeat) heartbeatMain() error {
 	logger.Info("Heartbeating!")
 	params := &url.Values{}
-	params.Add(requestIDKey, NewUUID().String())
-	params.Add(requestGUIDKey, NewUUID().String())
+	params.Set(requestIDKey, NewUUID().String())
+	params.Set(requestGUIDKey, NewUUID().String())
 	headers := getHeaders()
 	token, _, _ := hc.restful.TokenAccessor.GetTokens()
 	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
diff --git a/monitoring.go b/monitoring.go
index 19af715..0831385 100644
--- a/monitoring.go
+++ b/monitoring.go
@@ -132,7 +132,7 @@ func (sc *snowflakeConn) checkQueryStatus(
 	*retStatus, error) {
 	headers := make(map[string]string)
 	param := make(url.Values)
-	param.Add(requestGUIDKey, NewUUID().String())
+	param.Set(requestGUIDKey, NewUUID().String())
 	if tok, _, _ := sc.rest.TokenAccessor.GetTokens(); tok != "" {
 		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, tok)
 	}
@@ -206,9 +206,9 @@ func (sc *snowflakeConn) getQueryResultResp(
 	}
 	paramsMutex.Unlock()
 	param := make(url.Values)
-	param.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
-	param.Add("clientStartTime", strconv.FormatInt(sc.currentTimeProvider.currentTime(), 10))
-	param.Add(requestGUIDKey, NewUUID().String())
+	param.Set(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
+	param.Set("clientStartTime", strconv.FormatInt(sc.currentTimeProvider.currentTime(), 10))
+	param.Set(requestGUIDKey, NewUUID().String())
 	token, _, _ := sc.rest.TokenAccessor.GetTokens()
 	if token != "" {
 		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
diff --git a/restful.go b/restful.go
index efc3fb8..69105d0 100644
--- a/restful.go
+++ b/restful.go
@@ -234,8 +234,8 @@ func postRestfulQueryHelper(
 	cfg *Config) (
 	data *execResponse, err error) {
 	logger.WithContext(ctx).Infof("params: %v", params)
-	params.Add(requestIDKey, requestID.String())
-	params.Add(requestGUIDKey, NewUUID().String())
+	params.Set(requestIDKey, requestID.String())
+	params.Set(requestGUIDKey, NewUUID().String())
 	token, _, _ := sr.TokenAccessor.GetTokens()
 	if token != "" {
 		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
@@ -326,9 +326,9 @@ func postRestfulQueryHelper(
 func closeSession(ctx context.Context, sr *snowflakeRestful, timeout time.Duration) error {
 	logger.WithContext(ctx).Info("close session")
 	params := &url.Values{}
-	params.Add("delete", "true")
-	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
-	params.Add(requestGUIDKey, NewUUID().String())
+	params.Set("delete", "true")
+	params.Set(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
+	params.Set(requestGUIDKey, NewUUID().String())
 	fullURL := sr.getFullURL(sessionRequestPath, params)
 
 	headers := getHeaders()
@@ -376,8 +376,8 @@ func closeSession(ctx context.Context, sr *snowflakeRestful, timeout time.Durati
 func renewRestfulSession(ctx context.Context, sr *snowflakeRestful, timeout time.Duration) error {
 	logger.WithContext(ctx).Info("start renew session")
 	params := &url.Values{}
-	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
-	params.Add(requestGUIDKey, NewUUID().String())
+	params.Set(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
+	params.Set(requestGUIDKey, NewUUID().String())
 	fullURL := sr.getFullURL(tokenRequestPath, params)
 
 	token, masterToken, _ := sr.TokenAccessor.GetTokens()
@@ -449,8 +449,8 @@ func getCancelRetry(ctx context.Context) int {
 func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID UUID, timeout time.Duration) error {
 	logger.WithContext(ctx).Info("cancel query")
 	params := &url.Values{}
-	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
-	params.Add(requestGUIDKey, NewUUID().String())
+	params.Set(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
+	params.Set(requestGUIDKey, NewUUID().String())
 
 	fullURL := sr.getFullURL(abortRequestPath, params)
 


Test patch:
diff --git a/restful_test.go b/restful_test.go
index 0eb5fc4..867ac4b 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -202,6 +202,36 @@ func TestUnitPostQueryHelperError(t *testing.T) {
 	}
 }
 
+func TestUnitPostQueryHelperOnRenewSessionKeepsRequestIdButGeneratesNewRequestGuid(t *testing.T) {
+	postCount := 0
+	requestID := NewUUID()
+
+	sr := &snowflakeRestful{
+		FuncPost: func(ctx context.Context, restful *snowflakeRestful, url *url.URL, headers map[string]string, bytes []byte, duration time.Duration, provider currentTimeProvider, config *Config) (*http.Response, error) {
+			assertEqualF(t, len((url.Query())[requestIDKey]), 1)
+			assertEqualF(t, len((url.Query())[requestGUIDKey]), 1)
+			return &http.Response{
+				StatusCode: 200,
+				Body:       &fakeResponseBody{body: []byte(`{"data":null,"code":"390112","message":"token expired for testing","success":false,"headers":null}`)},
+			}, nil
+		},
+		FuncPostQuery: func(ctx context.Context, restful *snowflakeRestful, values *url.Values, headers map[string]string, bytes []byte, timeout time.Duration, uuid UUID, config *Config) (*execResponse, error) {
+			assertEqualF(t, requestID.String(), uuid.String())
+			assertEqualF(t, len((*values)[requestIDKey]), 1)
+			assertEqualF(t, len((*values)[requestGUIDKey]), 1)
+			if postCount == 0 {
+				postCount++
+				return postRestfulQueryHelper(ctx, restful, values, headers, bytes, timeout, uuid, config)
+			}
+			return nil, nil
+		},
+		FuncRenewSession: renewSessionTest,
+		TokenAccessor:    getSimpleTokenAccessor(),
+	}
+	_, err := postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), make([]byte, 0), time.Second, requestID, nil)
+	assertNilE(t, err)
+}
+
 func renewSessionTest(_ context.Context, _ *snowflakeRestful, _ time.Duration) error {
 	return nil
 }