SNOW-1754426 Fix duplicated requestId and request_guid params (#1224)

snowflakedb · Oct 24, 2024 · 1a90e21 · 1a90e21
1 parent 58c35d4
commit 1a90e21
Show file tree

Hide file tree

Showing 6 changed files with 48 additions and 18 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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)
 

diff --git 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
 }