Repository: snowflakedb/gosnowflake
Pull request #368 (2021-01-15 21:48:23)

SNOW-209082 Introduce TokenAccessor to manage tokens and sessions
### Description
Introduce the concept of a token accessor interface that can be used to
manage the lifecycle of token & master token & session ID.

Currently, token, master token, and session ID are tied directly to `snowflakeRestful`.
Instead, allow a `TokenAccessor` to be defined in the `Config` and let the token accessor
manage setting, and retrieving the token to use when making REST calls.

Added new unit tests and updated existing tests to make sure queries carry
the token from the token accessor, and session renewals set the token back into the accessor.

### 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
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing
- [ ] Extended the README / documentation, if necessary


sfc-gh-dwu: okay i need to rebase, thanks guys

Patch:
diff --git a/auth.go b/auth.go
index ae65b83..33fc060 100644
--- a/auth.go
+++ b/auth.go
@@ -263,7 +263,6 @@ func authenticate(
 	samlResponse []byte,
 	proofKey []byte,
 ) (resp *authResponseMain, err error) {
-
 	headers := getHeaders()
 	clientEnvironment := authRequestClientEnvironment{
 		Application: sc.cfg.Application,
@@ -352,9 +351,7 @@ func authenticate(
 	}
 	if !respd.Success {
 		logger.Errorln("Authentication FAILED")
-		sc.rest.Token = ""
-		sc.rest.MasterToken = ""
-		sc.rest.SessionID = -1
+		sc.rest.TokenAccessor.SetTokens("", "", -1)
 		code, err := strconv.Atoi(respd.Code)
 		if err != nil {
 			code = -1
@@ -367,9 +364,7 @@ func authenticate(
 		}
 	}
 	logger.Info("Authentication SUCCESS")
-	sc.rest.Token = respd.Data.Token
-	sc.rest.MasterToken = respd.Data.MasterToken
-	sc.rest.SessionID = respd.Data.SessionID
+	sc.rest.TokenAccessor.SetTokens(respd.Data.Token, respd.Data.MasterToken, respd.Data.SessionID)
 	return &respd.Data, nil
 }
 
diff --git a/authokta.go b/authokta.go
index 55a7735..5084ee4 100644
--- a/authokta.go
+++ b/authokta.go
@@ -90,9 +90,7 @@ func authenticateBySAML(
 	}
 	if !respd.Success {
 		logger.Errorln("Authentication FAILED")
-		sr.Token = ""
-		sr.MasterToken = ""
-		sr.SessionID = -1
+		sr.TokenAccessor.SetTokens("", "", -1)
 		code, err := strconv.Atoi(respd.Code)
 		if err != nil {
 			code = -1
diff --git a/connection.go b/connection.go
index 6b7964b..51923f7 100644
--- a/connection.go
+++ b/connection.go
@@ -581,8 +581,9 @@ func (sc *snowflakeConn) getQueryResult(ctx context.Context, resultPath string)
 	param.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
 	param.Add("clientStartTime", strconv.FormatInt(time.Now().Unix(), 10))
 	param.Add(requestGUIDKey, uuid.New().String())
-	if sc.rest.Token != "" {
-		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, sc.rest.Token)
+	token, _, _ := sc.rest.TokenAccessor.GetTokens()
+	if token != "" {
+		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
 	}
 	url := sc.rest.getFullURL(resultPath, &param)
 	res, err := sc.rest.FuncGet(ctx, sc.rest, url, headers, sc.rest.RequestTimeout)
@@ -634,7 +635,8 @@ func getAsync(
 		sfError.QueryID = rows.queryID
 	}
 	defer close(errChannel)
-	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, sr.Token)
+	token, _, _ := sr.TokenAccessor.GetTokens()
+	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
 	resp, err := sr.FuncGet(ctx, sr, URL, headers, timeout)
 	if err != nil {
 		logger.WithContext(ctx).Errorf("failed to get response. err: %v", err)
diff --git a/driver.go b/driver.go
index acac351..187ca2d 100644
--- a/driver.go
+++ b/driver.go
@@ -49,6 +49,12 @@ func (d SnowflakeDriver) OpenWithConfig(ctx context.Context, config Config) (dri
 		// use the custom transport
 		st = sc.cfg.Transporter
 	}
+	var tokenAccessor TokenAccessor
+	if sc.cfg.TokenAccessor != nil {
+		tokenAccessor = sc.cfg.TokenAccessor
+	} else {
+		tokenAccessor = getSimpleTokenAccessor()
+	}
 	// authenticate
 	sc.rest = &snowflakeRestful{
 		Host:     sc.cfg.Host,
@@ -59,6 +65,7 @@ func (d SnowflakeDriver) OpenWithConfig(ctx context.Context, config Config) (dri
 			Timeout:   sc.cfg.ClientTimeout,
 			Transport: st,
 		},
+		TokenAccessor:       tokenAccessor,
 		LoginTimeout:        sc.cfg.LoginTimeout,
 		RequestTimeout:      sc.cfg.RequestTimeout,
 		FuncPost:            postRestful,
diff --git a/dsn.go b/dsn.go
index eaff73d..758605d 100644
--- a/dsn.go
+++ b/dsn.go
@@ -71,7 +71,8 @@ type Config struct {
 	InsecureMode bool             // driver doesn't check certificate revocation status
 	OCSPFailOpen OCSPFailOpenMode // OCSP Fail Open
 
-	Token string // Token to use for OAuth other forms of token based auth
+	Token         string        // Token to use for OAuth other forms of token based auth
+	TokenAccessor TokenAccessor // Optional token accessor to use
 
 	PrivateKey *rsa.PrivateKey // Private key used to sign JWT
 
diff --git a/heartbeat.go b/heartbeat.go
index a0c3c81..c852816 100644
--- a/heartbeat.go
+++ b/heartbeat.go
@@ -62,7 +62,8 @@ func (hc *heartbeat) heartbeatMain() error {
 	headers["Content-Type"] = headerContentTypeApplicationJSON
 	headers["accept"] = headerAcceptTypeApplicationSnowflake
 	headers["User-Agent"] = userAgent
-	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, hc.restful.Token)
+	token, _, _ := hc.restful.TokenAccessor.GetTokens()
+	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
 
 	fullURL := hc.restful.getFullURL(heartBeatPath, params)
 	timeout := hc.restful.RequestTimeout
diff --git a/restful.go b/restful.go
index a8efbb1..0be7545 100644
--- a/restful.go
+++ b/restful.go
@@ -49,11 +49,9 @@ type snowflakeRestful struct {
 	LoginTimeout   time.Duration // Login timeout
 	RequestTimeout time.Duration // request timeout
 
-	Client      *http.Client
-	Token       string
-	MasterToken string
-	SessionID   int
-	HeartBeat   *heartbeat
+	Client        *http.Client
+	TokenAccessor TokenAccessor
+	HeartBeat     *heartbeat
 
 	Connection          *snowflakeConn
 	FuncPostQuery       func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, uuid.UUID, *Config) (*execResponse, error)
@@ -174,8 +172,9 @@ func postRestfulQueryHelper(
 	params.Add(requestIDKey, requestID.String())
 	params.Add("clientStartTime", strconv.FormatInt(time.Now().Unix(), 10))
 	params.Add(requestGUIDKey, uuid.New().String())
-	if sr.Token != "" {
-		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, sr.Token)
+	token, _, _ := sr.TokenAccessor.GetTokens()
+	if token != "" {
+		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
 	}
 
 	var resp *http.Response
@@ -244,7 +243,8 @@ func postRestfulQueryHelper(
 			}
 
 			logger.Info("ping pong")
-			headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, sr.Token)
+			token, _, _ := sr.TokenAccessor.GetTokens()
+			headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
 			fullURL := sr.getFullURL(resultURL, nil)
 
 			resp, err = sr.FuncGet(ctx, sr, fullURL, headers, timeout)
@@ -298,7 +298,8 @@ func closeSession(ctx context.Context, sr *snowflakeRestful, timeout time.Durati
 	headers["Content-Type"] = headerContentTypeApplicationJSON
 	headers["accept"] = headerAcceptTypeApplicationSnowflake
 	headers["User-Agent"] = userAgent
-	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, sr.Token)
+	token, _, _ := sr.TokenAccessor.GetTokens()
+	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
 
 	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, nil, 5*time.Second, false)
 	if err != nil {
@@ -346,14 +347,15 @@ func renewRestfulSession(ctx context.Context, sr *snowflakeRestful, timeout time
 	params.Add(requestGUIDKey, uuid.New().String())
 	fullURL := sr.getFullURL(tokenRequestPath, params)
 
+	token, masterToken, _ := sr.TokenAccessor.GetTokens()
 	headers := make(map[string]string)
 	headers["Content-Type"] = headerContentTypeApplicationJSON
 	headers["accept"] = headerAcceptTypeApplicationSnowflake
 	headers["User-Agent"] = userAgent
-	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, sr.MasterToken)
+	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, masterToken)
 
 	body := make(map[string]string)
-	body["oldSessionToken"] = sr.Token
+	body["oldSessionToken"] = token
 	body["requestType"] = "RENEW"
 
 	var reqBody []byte
@@ -384,8 +386,7 @@ func renewRestfulSession(ctx context.Context, sr *snowflakeRestful, timeout time
 				Message: respd.Message,
 			}
 		}
-		sr.Token = respd.Data.SessionToken
-		sr.MasterToken = respd.Data.MasterToken
+		sr.TokenAccessor.SetTokens(respd.Data.SessionToken, respd.Data.MasterToken, respd.Data.SessionID)
 		return nil
 	}
 	b, err := ioutil.ReadAll(resp.Body)
@@ -415,7 +416,8 @@ func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID uuid.UUID,
 	headers["Content-Type"] = headerContentTypeApplicationJSON
 	headers["accept"] = headerAcceptTypeApplicationSnowflake
 	headers["User-Agent"] = userAgent
-	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, sr.Token)
+	token, _, _ := sr.TokenAccessor.GetTokens()
+	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
 
 	req := make(map[string]string)
 	req[requestIDKey] = requestID.String()
diff --git a/util.go b/util.go
index 1ac5a4d..42dfc20 100644
--- a/util.go
+++ b/util.go
@@ -99,3 +99,29 @@ func toNamedValues(values []driver.Value) []driver.NamedValue {
 	}
 	return namedValues
 }
+
+// TokenAccessor manages the session token and master token
+type TokenAccessor interface {
+	GetTokens() (token string, masterToken string, sessionID int)
+	SetTokens(token string, masterToken string, sessionID int)
+}
+
+type simpleTokenAccessor struct {
+	token       string
+	masterToken string
+	sessionID   int
+}
+
+func getSimpleTokenAccessor() TokenAccessor {
+	return &simpleTokenAccessor{sessionID: -1}
+}
+
+func (sta *simpleTokenAccessor) GetTokens() (token string, masterToken string, sessionID int) {
+	return sta.token, sta.masterToken, sta.sessionID
+}
+
+func (sta *simpleTokenAccessor) SetTokens(token string, masterToken string, sessionID int) {
+	sta.token = token
+	sta.masterToken = masterToken
+	sta.sessionID = sessionID
+}


Test patch:
diff --git a/auth_test.go b/auth_test.go
index b2729b5..9e74507 100644
--- a/auth_test.go
+++ b/auth_test.go
@@ -18,8 +18,8 @@ import (
 
 func TestUnitPostAuth(t *testing.T) {
 	sr := &snowflakeRestful{
-		Token:    "token",
-		FuncPost: postTestAfterRenew,
+		TokenAccessor: getSimpleTokenAccessor(),
+		FuncPost:      postTestAfterRenew,
 	}
 	var err error
 	_, err = postAuth(context.TODO(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
@@ -241,7 +241,9 @@ func getDefaultSnowflakeConn() *snowflakeConn {
 		Passcode:           "",
 		Application:        "testapp",
 	}
-	sr := &snowflakeRestful{}
+	sr := &snowflakeRestful{
+		TokenAccessor: getSimpleTokenAccessor(),
+	}
 	sc := &snowflakeConn{
 		rest: sr,
 		cfg:  &cfg,
@@ -254,9 +256,11 @@ func TestUnitAuthenticate(t *testing.T) {
 	var driverErr *SnowflakeError
 	var ok bool
 
+	ta := getSimpleTokenAccessor()
 	sc := getDefaultSnowflakeConn()
 	sr := &snowflakeRestful{
-		FuncPostAuth: postAuthFailServiceIssue,
+		FuncPostAuth:  postAuthFailServiceIssue,
+		TokenAccessor: ta,
 	}
 	sc.rest = sr
 
@@ -286,20 +290,30 @@ func TestUnitAuthenticate(t *testing.T) {
 	if !ok || driverErr.Number != ErrFailedToAuth {
 		t.Fatalf("Snowflake error is expected. err: %v", driverErr)
 	}
+	ta.SetTokens("bad-token", "bad-master-token", 1)
 	sr.FuncPostAuth = postAuthSuccessWithErrorCode
 	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
+	newToken, newMasterToken, newSessionID := ta.GetTokens()
+	if newToken != "" || newMasterToken != "" || newSessionID != -1 {
+		t.Fatalf("failed auth should have reset tokens: %v %v %v", newToken, newMasterToken, newSessionID)
+	}
 	driverErr, ok = err.(*SnowflakeError)
 	if !ok || driverErr.Number != 98765 {
 		t.Fatalf("Snowflake error is expected. err: %v", driverErr)
 	}
+	ta.SetTokens("bad-token", "bad-master-token", 1)
 	sr.FuncPostAuth = postAuthSuccessWithInvalidErrorCode
 	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
+	oldToken, oldMasterToken, oldSessionID := ta.GetTokens()
+	if oldToken != "" || oldMasterToken != "" || oldSessionID != -1 {
+		t.Fatalf("failed auth should have reset tokens: %v %v %v", oldToken, oldMasterToken, oldSessionID)
+	}
 	sr.FuncPostAuth = postAuthSuccess
 	var resp *authResponseMain
 	resp, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
@@ -309,12 +323,23 @@ func TestUnitAuthenticate(t *testing.T) {
 	if resp.SessionInfo.DatabaseName != "dbn" {
 		t.Fatalf("failed to get response from auth")
 	}
+	newToken, newMasterToken, newSessionID = ta.GetTokens()
+	if newToken == oldToken {
+		t.Fatalf("new token was not set: %v", newToken)
+	}
+	if newMasterToken == oldMasterToken {
+		t.Fatalf("new master token was not set: %v", newMasterToken)
+	}
+	if newSessionID == oldSessionID {
+		t.Fatalf("new session id was not set: %v", newSessionID)
+	}
 }
 
 func TestUnitAuthenticateSaml(t *testing.T) {
 	var err error
 	sr := &snowflakeRestful{
-		FuncPostAuth: postAuthCheckSAMLResponse,
+		FuncPostAuth:  postAuthCheckSAMLResponse,
+		TokenAccessor: getSimpleTokenAccessor(),
 	}
 	sc := getDefaultSnowflakeConn()
 	sc.cfg.Authenticator = AuthTypeOkta
@@ -333,7 +358,8 @@ func TestUnitAuthenticateSaml(t *testing.T) {
 func TestUnitAuthenticateOAuth(t *testing.T) {
 	var err error
 	sr := &snowflakeRestful{
-		FuncPostAuth: postAuthCheckOAuth,
+		FuncPostAuth:  postAuthCheckOAuth,
+		TokenAccessor: getSimpleTokenAccessor(),
 	}
 	sc := getDefaultSnowflakeConn()
 	sc.cfg.Token = "oauthToken"
@@ -348,7 +374,8 @@ func TestUnitAuthenticateOAuth(t *testing.T) {
 func TestUnitAuthenticatePasscode(t *testing.T) {
 	var err error
 	sr := &snowflakeRestful{
-		FuncPostAuth: postAuthCheckPasscode,
+		FuncPostAuth:  postAuthCheckPasscode,
+		TokenAccessor: getSimpleTokenAccessor(),
 	}
 	sc := getDefaultSnowflakeConn()
 	sc.cfg.Passcode = "987654321"
@@ -372,7 +399,8 @@ func TestUnitAuthenticateJWT(t *testing.T) {
 	var err error
 
 	sr := &snowflakeRestful{
-		FuncPostAuth: postAuthCheckJWTToken,
+		FuncPostAuth:  postAuthCheckJWTToken,
+		TokenAccessor: getSimpleTokenAccessor(),
 	}
 	sc := getDefaultSnowflakeConn()
 	sc.cfg.Authenticator = AuthTypeJwt
diff --git a/authokta_test.go b/authokta_test.go
index 6586d40..df642dc 100644
--- a/authokta_test.go
+++ b/authokta_test.go
@@ -60,7 +60,8 @@ func getTestHTMLSuccess(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ ma
 
 func TestUnitPostAuthSAML(t *testing.T) {
 	sr := &snowflakeRestful{
-		FuncPost: postTestError,
+		FuncPost:      postTestError,
+		TokenAccessor: getSimpleTokenAccessor(),
 	}
 	var err error
 	_, err = postAuthSAML(context.TODO(), sr, make(map[string]string), []byte{}, 0)
@@ -81,7 +82,8 @@ func TestUnitPostAuthSAML(t *testing.T) {
 
 func TestUnitPostAuthOKTA(t *testing.T) {
 	sr := &snowflakeRestful{
-		FuncPost: postTestError,
+		FuncPost:      postTestError,
+		TokenAccessor: getSimpleTokenAccessor(),
 	}
 	var err error
 	_, err = postAuthOKTA(context.TODO(), sr, make(map[string]string), []byte{}, "hahah", 0)
@@ -102,7 +104,8 @@ func TestUnitPostAuthOKTA(t *testing.T) {
 
 func TestUnitGetSSO(t *testing.T) {
 	sr := &snowflakeRestful{
-		FuncGet: getTestError,
+		FuncGet:       getTestError,
+		TokenAccessor: getSimpleTokenAccessor(),
 	}
 	var err error
 	_, err = getSSO(context.TODO(), sr, &url.Values{}, make(map[string]string), "hahah", 0)
@@ -188,6 +191,7 @@ func TestUnitAuthenticateBySAML(t *testing.T) {
 		Host:             "abc.com",
 		Port:             443,
 		FuncPostAuthSAML: postAuthSAMLError,
+		TokenAccessor:    getSimpleTokenAccessor(),
 	}
 	var err error
 	_, err = authenticateBySAML(context.TODO(), sr, authenticator, application, account, user, password)
diff --git a/connection_test.go b/connection_test.go
index db6b606..922ea5b 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -4,6 +4,9 @@ package gosnowflake
 
 import (
 	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
 	"net/url"
 	"testing"
 	"time"
@@ -66,6 +69,44 @@ func TestExecWithEmptyRequestID(t *testing.T) {
 	}
 }
 
+func TestGetQueryResultUsesTokenFromTokenAccessor(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) {
+		if headers[headerAuthorizationKey] != fmt.Sprintf(headerSnowflakeToken, token) {
+			t.Fatalf("header authorization key is not correct: %v", headers[headerAuthorizationKey])
+		}
+		dd := &execResponseData{}
+		er := &execResponse{
+			Data:    *dd,
+			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
+	}
+	sr := &snowflakeRestful{
+		FuncGet: funcGetMock,
+		TokenAccessor: ta,
+	}
+	sc := &snowflakeConn{
+		cfg:  &Config{Params: map[string]*string{}},
+		rest: sr,
+	}
+	_, err := sc.getQueryResult(context.Background(), "")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+}
+
 func TestExecWithSpecificRequestID(t *testing.T) {
 	origRequestID := uuid.New()
 	ctx := WithRequestID(context.Background(), origRequestID)
diff --git a/restful_test.go b/restful_test.go
index cb8f90c..cbb4ce3 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -6,6 +6,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"net/http"
 	"net/url"
 	"testing"
@@ -91,8 +92,8 @@ func postTestAfterRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ ma
 
 func TestUnitPostQueryHelperError(t *testing.T) {
 	sr := &snowflakeRestful{
-		Token:    "token",
-		FuncPost: postTestError,
+		FuncPost:      postTestError,
+		TokenAccessor: getSimpleTokenAccessor(),
 	}
 	var err error
 	var requestID uuid.UUID
@@ -123,6 +124,36 @@ func renewSessionTestError(_ context.Context, _ *snowflakeRestful, _ time.Durati
 	return errors.New("failed to renew session in tests")
 }
 
+func TestUnitPostQueryHelperUsesToken(t *testing.T) {
+	accessor := getSimpleTokenAccessor()
+	token := "token123"
+	accessor.SetTokens(token, "", 0)
+
+	var err error
+	postQueryTest := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, headers map[string]string, _ []byte, _ time.Duration, _ uuid.UUID, _ *Config) (*execResponse, error) {
+		if headers[headerAuthorizationKey] != fmt.Sprintf(headerSnowflakeToken, token) {
+			t.Fatalf("authorization key doesn't match, %v vs %v", headers[headerAuthorizationKey], fmt.Sprintf(headerSnowflakeToken, token))
+		}
+		dd := &execResponseData{}
+		return &execResponse{
+			Data:    *dd,
+			Message: "",
+			Code:    "0",
+			Success: true,
+		}, nil
+	}
+	sr := &snowflakeRestful{
+		FuncPost:         postTestRenew,
+		FuncPostQuery:    postQueryTest,
+		FuncRenewSession: renewSessionTest,
+		TokenAccessor:    accessor,
+	}
+	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, uuid.New(), &Config{})
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+}
+
 func TestUnitPostQueryHelperRenewSession(t *testing.T) {
 	var err error
 	origRequestID := uuid.New()
@@ -140,10 +171,10 @@ func TestUnitPostQueryHelperRenewSession(t *testing.T) {
 		}, nil
 	}
 	sr := &snowflakeRestful{
-		Token:            "token",
 		FuncPost:         postTestRenew,
 		FuncPostQuery:    postQueryTest,
 		FuncRenewSession: renewSessionTest,
+		TokenAccessor:    getSimpleTokenAccessor(),
 	}
 
 	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, origRequestID, &Config{})
@@ -158,10 +189,35 @@ func TestUnitPostQueryHelperRenewSession(t *testing.T) {
 }
 
 func TestUnitRenewRestfulSession(t *testing.T) {
+	accessor := getSimpleTokenAccessor()
+	oldToken, oldMasterToken, oldSessionID := "oldtoken", "oldmaster", 100
+	newToken, newMasterToken, newSessionID := "newtoken", "newmaster", 200
+	postTestSuccessWithNewTokens := func(_ context.Context, _ *snowflakeRestful, _ *url.URL, headers map[string]string, _ []byte, _ time.Duration, _ bool) (*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))
+		}
+		tr := &renewSessionResponse{
+			Data: renewSessionResponseMain{
+				SessionToken: newToken,
+				MasterToken:  newMasterToken,
+				SessionID:    newSessionID,
+			},
+			Message: "",
+			Success: true,
+		}
+		ba, err := json.Marshal(tr)
+		if err != nil {
+			t.Fatalf("failed to serialize token response %v", err)
+		}
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Body:       &fakeResponseBody{body: ba},
+		}, nil
+	}
+
 	sr := &snowflakeRestful{
-		MasterToken: "mtoken",
-		Token:       "token",
-		FuncPost:    postTestAfterRenew,
+		FuncPost:      postTestAfterRenew,
+		TokenAccessor: accessor,
 	}
 	err := renewRestfulSession(context.Background(), sr, time.Second)
 	if err != nil {
@@ -182,11 +238,28 @@ func TestUnitRenewRestfulSession(t *testing.T) {
 	if err == nil {
 		t.Fatal("should have failed to run post request after the renewal")
 	}
+	accessor.SetTokens(oldToken, oldMasterToken, oldSessionID)
+	sr.FuncPost = postTestSuccessWithNewTokens
+	err = renewRestfulSession(context.Background(), sr, time.Second)
+	if err != nil {
+		t.Fatal("should not have failed to run post request after the renewal")
+	}
+	token, masterToken, sessionID := accessor.GetTokens()
+	if token != newToken {
+		t.Fatalf("unexpected new token %v", token)
+	}
+	if masterToken != newMasterToken {
+		t.Fatalf("unexpected new master token %v", masterToken)
+	}
+	if sessionID != newSessionID {
+		t.Fatalf("unexpected new session id %v", sessionID)
+	}
 }
 
 func TestUnitCloseSession(t *testing.T) {
 	sr := &snowflakeRestful{
-		FuncPost: postTestAfterRenew,
+		FuncPost:      postTestAfterRenew,
+		TokenAccessor: getSimpleTokenAccessor(),
 	}
 	err := closeSession(context.Background(), sr, time.Second)
 	if err != nil {
@@ -211,7 +284,8 @@ func TestUnitCloseSession(t *testing.T) {
 
 func TestUnitCancelQuery(t *testing.T) {
 	sr := &snowflakeRestful{
-		FuncPost: postTestAfterRenew,
+		FuncPost:      postTestAfterRenew,
+		TokenAccessor: getSimpleTokenAccessor(),
 	}
 	ctx := context.Background()
 	err := cancelQuery(ctx, sr, getOrGenerateRequestIDFromContext(ctx), time.Second)
diff --git a/util_test.go b/util_test.go
index 59aa07b..865f9c1 100644
--- a/util_test.go
+++ b/util_test.go
@@ -17,6 +17,33 @@ type tcIntMinMax struct {
 	out int
 }
 
+func TestSimpleTokenAccessor(t *testing.T) {
+	accessor := getSimpleTokenAccessor()
+	token, masterToken, sessionID := accessor.GetTokens()
+	if token != "" {
+		t.Errorf("unexpected token %v", token)
+	}
+	if masterToken != "" {
+		t.Errorf("unexpected master token %v", masterToken)
+	}
+	if sessionID != -1 {
+		t.Errorf("unexpected session id %v", sessionID)
+	}
+
+	expectedToken, expectedMasterToken, expectedSessionID := "token123", "master123", 123
+	accessor.SetTokens(expectedToken, expectedMasterToken, expectedSessionID)
+	token, masterToken, sessionID = accessor.GetTokens()
+	if token != expectedToken {
+		t.Errorf("unexpected token %v", token)
+	}
+	if masterToken != expectedMasterToken {
+		t.Errorf("unexpected master token %v", masterToken)
+	}
+	if sessionID != expectedSessionID {
+		t.Errorf("unexpected session id %v", sessionID)
+	}
+}
+
 func TestGetRequestIDFromContext(t *testing.T) {
 	expectedRequestID := uuid.New()
 	ctx := WithRequestID(context.Background(), expectedRequestID)