Repository: snowflakedb/gosnowflake
Pull request #1160 (2024-06-18 06:54:46)

SNOW-1346547: add check for session expiry when getting query result
### Description

drivers issue 957: https://github.com/snowflakedb/snowflake-sdks-drivers-issues-teamwork/issues/957

The root cause is that the heartbeat process has renewed the session token [here](https://github.com/snowflakedb/gosnowflake/blob/v1.10.1/heartbeat.go#L79) but the new token in the get query result request's header was not updated in [getQueryResultWithRetriesForAsyncMode()](https://github.com/snowflakedb/gosnowflake/blob/822-support-put-cancel/async.go#L136).

This PR adds a check for session expiry in `getQueryResultWithRetriesForAsyncMode()` and update the token in the header if needed. Similar to ODBC's implementation [here](https://github.com/snowflakedb/snowflake-odbc/blob/master/Source/Platform/Statement.cpp#L385), where there should be a check for the session expired code 390112 during the getting query result request and allows a retry after updating the session token.

### Checklist
- [X] Created tests which fail without the change (if possible)
- [ ] Extended the README / documentation, if necessary


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1160?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
Attention: Patch coverage is `88.46154%` with `3 lines` in your changes missing coverage. Please review.
> Project coverage is 72.38%. Comparing base [(`5d28db8`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/5d28db80c2ffac67e9e8991eb2d78c591faedb47?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) to head [(`80bc5ee`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/80bc5ee238eeca4c28e00d6740b17e9e6117ad27?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).

| [Files](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1160?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | Patch % | Lines |
|---|---|---|
| [async.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1160?src=pr&el=tree&filepath=async.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-YXN5bmMuZ28=) | 88.46% | [3 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1160?src=pr&el=tree&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    #1160      +/-   ##
==========================================
+ Coverage   72.34%   72.38%   +0.03%     
==========================================
  Files          53       53              
  Lines       11909    11934      +25     
==========================================
+ Hits         8616     8638      +22     
- Misses       3293     3296       +3     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1160?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).

sfc-gh-ext-simba-jl: Added a check for session expiry in getQueryResultWithRetriesForAsyncMode() and update the token in the header if needed.

Patch:
diff --git a/async.go b/async.go
index ed870d4..d8b88c9 100644
--- a/async.go
+++ b/async.go
@@ -148,6 +148,7 @@ func getQueryResultWithRetriesForAsyncMode(
 	retry := 0
 	retryPattern := []int32{1, 1, 2, 3, 4, 8, 10}
 	retryPatternIndex := 0
+	retryCountForSessionRenewal := 0
 
 	for {
 		logger.WithContext(ctx).Debugf("Retry count for get query result request in async mode: %v", retry)
@@ -165,7 +166,34 @@ func getQueryResultWithRetriesForAsyncMode(
 			logger.WithContext(ctx).Errorf("failed to decode JSON. err: %v", err)
 			return respd, err
 		}
-		if respd.Code != queryInProgressAsyncCode {
+		if respd.Code == sessionExpiredCode {
+			// Update the session token in the header and retry
+			token, _, _ := sr.TokenAccessor.GetTokens()
+			if token != "" && headers[headerAuthorizationKey] != fmt.Sprintf(headerSnowflakeToken, token) {
+				headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
+				logger.WithContext(ctx).Info("Session token has been updated.")
+				retry++
+				continue
+			}
+
+			// Renew the session token
+			if err = sr.renewExpiredSessionToken(ctx, timeout, token); err != nil {
+				logger.WithContext(ctx).Errorf("failed to renew session token. err: %v", err)
+				return respd, err
+			}
+			retryCountForSessionRenewal++
+
+			// If this is the first response, go back to retry the query
+			// since it failed due to session expiration
+			logger.WithContext(ctx).Infof("retry count for session renewal: %v", retryCountForSessionRenewal)
+			if retryCountForSessionRenewal < 2 {
+				retry++
+				continue
+			} else {
+				logger.WithContext(ctx).Errorf("failed to get query result with the renewed session token. err: %v", err)
+				return respd, err
+			}
+		} else if respd.Code != queryInProgressAsyncCode {
 			// If the query takes longer than 45 seconds to complete the results are not returned.
 			// If the query is still in progress after 45 seconds, retry the request to the /results endpoint.
 			// For all other scenarios continue processing results response


Test patch:
diff --git a/connection_test.go b/connection_test.go
index 658cd5a..46a9c34 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -113,7 +113,7 @@ func TestGetQueryResultUsesTokenFromTokenAccessor(t *testing.T) {
 		er := &execResponse{
 			Data:    *dd,
 			Message: "",
-			Code:    sessionExpiredCode,
+			Code:    "0",
 			Success: true,
 		}
 		ba, err := json.Marshal(er)
@@ -139,6 +139,106 @@ func TestGetQueryResultUsesTokenFromTokenAccessor(t *testing.T) {
 	}
 }
 
+func TestGetQueryResultTokenExpiry(t *testing.T) {
+	ta := getSimpleTokenAccessor()
+	token := "snowflake-test-token"
+	ta.SetTokens(token, "", 1)
+	funcGetMock := func(_ context.Context, _ *snowflakeRestful, _ *url.URL,
+		headers map[string]string, _ time.Duration) (*http.Response, error) {
+		respData := execResponseData{}
+		er := &execResponse{
+			Data:    respData,
+			Message: "",
+			Code:    sessionExpiredCode,
+			Success: true,
+		}
+		ba, err := json.Marshal(er)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Body:       &fakeResponseBody{body: ba},
+		}, nil
+	}
+
+	expectedToken := "new token"
+	expectedMaster := "new master"
+	expectedSession := int64(321)
+
+	renewSessionDummy := func(_ context.Context, sr *snowflakeRestful, _ time.Duration) error {
+		ta.SetTokens(expectedToken, expectedMaster, expectedSession)
+		return nil
+	}
+
+	sr := &snowflakeRestful{
+		FuncGet:          funcGetMock,
+		FuncRenewSession: renewSessionDummy,
+		TokenAccessor:    ta,
+	}
+	sc := &snowflakeConn{
+		cfg:                 &Config{Params: map[string]*string{}},
+		rest:                sr,
+		currentTimeProvider: defaultTimeProvider,
+	}
+	_, err := sc.getQueryResultResp(context.Background(), "")
+	assertNilF(t, err, fmt.Sprintf("err: %v", err))
+
+	updatedToken, updatedMaster, updatedSession := ta.GetTokens()
+	assertEqualF(t, updatedToken, expectedToken)
+	assertEqualF(t, updatedMaster, expectedMaster)
+	assertEqualF(t, updatedSession, expectedSession)
+}
+
+func TestGetQueryResultTokenNotSet(t *testing.T) {
+	ta := getSimpleTokenAccessor()
+	funcGetMock := func(_ context.Context, _ *snowflakeRestful, _ *url.URL,
+		headers map[string]string, _ time.Duration) (*http.Response, error) {
+		respData := execResponseData{}
+		er := &execResponse{
+			Data:    respData,
+			Message: "",
+			Code:    sessionExpiredCode,
+			Success: true,
+		}
+		ba, err := json.Marshal(er)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Body:       &fakeResponseBody{body: ba},
+		}, nil
+	}
+
+	expectedToken := "new token"
+	expectedMaster := "new master"
+	expectedSession := int64(321)
+
+	renewSessionDummy := func(_ context.Context, sr *snowflakeRestful, _ time.Duration) error {
+		ta.SetTokens(expectedToken, expectedMaster, expectedSession)
+		return nil
+	}
+
+	sr := &snowflakeRestful{
+		FuncGet:          funcGetMock,
+		FuncRenewSession: renewSessionDummy,
+		TokenAccessor:    ta,
+	}
+	sc := &snowflakeConn{
+		cfg:                 &Config{Params: map[string]*string{}},
+		rest:                sr,
+		currentTimeProvider: defaultTimeProvider,
+	}
+	_, err := sc.getQueryResultResp(context.Background(), "")
+	assertNilF(t, err, fmt.Sprintf("err: %v", err))
+
+	updatedToken, updatedMaster, updatedSession := ta.GetTokens()
+	assertEqualF(t, updatedToken, expectedToken)
+	assertEqualF(t, updatedMaster, expectedMaster)
+	assertEqualF(t, updatedSession, expectedSession)
+}
+
 func TestExecWithSpecificRequestID(t *testing.T) {
 	origRequestID := NewUUID()
 	ctx := WithRequestID(context.Background(), origRequestID)