Repository: snowflakedb/gosnowflake
Pull request #841 (2023-07-05 07:27:01)

SNOW-833537 Separate JWT client
### Description
Please explain the changes you made here.

### 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
- [x] Extended the README / documentation, if necessary



Patch:
diff --git a/auth.go b/auth.go
index ec1bc1b..d1565cd 100644
--- a/auth.go
+++ b/auth.go
@@ -212,6 +212,7 @@ type authResponse struct {
 func postAuth(
 	ctx context.Context,
 	sr *snowflakeRestful,
+	client *http.Client,
 	params *url.Values,
 	headers map[string]string,
 	body []byte,
@@ -222,7 +223,7 @@ func postAuth(
 
 	fullURL := sr.getFullURL(loginRequestPath, params)
 	logger.Infof("full URL: %v", fullURL)
-	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, true)
+	resp, err := sr.FuncAuthPost(ctx, client, fullURL, headers, body, timeout, true)
 	if err != nil {
 		return nil, err
 	}
@@ -403,7 +404,7 @@ func authenticate(
 	logger.WithContext(sc.ctx).Infof("PARAMS for Auth: %v, %v, %v, %v, %v, %v",
 		params, sc.rest.Protocol, sc.rest.Host, sc.rest.Port, sc.rest.LoginTimeout, sc.cfg.Authenticator.String())
 
-	respd, err := sc.rest.FuncPostAuth(ctx, sc.rest, params, headers, jsonBody, sc.rest.LoginTimeout)
+	respd, err := sc.rest.FuncPostAuth(ctx, sc.rest, sc.rest.getClientFor(sc.cfg.Authenticator), params, headers, jsonBody, sc.rest.LoginTimeout)
 	if err != nil {
 		return nil, err
 	}
diff --git a/connection.go b/connection.go
index 84c2cfe..a4be499 100644
--- a/connection.go
+++ b/connection.go
@@ -73,7 +73,7 @@ type snowflakeConn struct {
 var (
 	queryIDPattern = `[\w\-_]+`
 	queryIDRegexp  = regexp.MustCompile(queryIDPattern)
-	errMutex = &sync.Mutex{}
+	errMutex       = &sync.Mutex{}
 )
 
 func (sc *snowflakeConn) exec(
@@ -718,11 +718,16 @@ func buildSnowflakeConn(ctx context.Context, config Config) (*snowflakeConn, err
 			Timeout:   sc.cfg.ClientTimeout,
 			Transport: st,
 		},
+		JWTClient: &http.Client{
+			Timeout:   sc.cfg.JWTClientTimeout,
+			Transport: st,
+		},
 		TokenAccessor:       tokenAccessor,
 		LoginTimeout:        sc.cfg.LoginTimeout,
 		RequestTimeout:      sc.cfg.RequestTimeout,
 		FuncPost:            postRestful,
 		FuncGet:             getRestful,
+		FuncAuthPost:        postAuthRestful,
 		FuncPostQuery:       postRestfulQuery,
 		FuncPostQueryHelper: postRestfulQueryHelper,
 		FuncRenewSession:    renewRestfulSession,
diff --git a/dsn.go b/dsn.go
index a16dc98..c4df6a1 100644
--- a/dsn.go
+++ b/dsn.go
@@ -20,6 +20,7 @@ import (
 
 const (
 	defaultClientTimeout          = 900 * time.Second // Timeout for network round trip + read out http response
+	defaultJWTClientTimeout       = 10 * time.Second  // Timeout for network round trip + read out http response but used for JWT auth
 	defaultLoginTimeout           = 60 * time.Second  // Timeout for retry for login EXCLUDING clientTimeout
 	defaultRequestTimeout         = 0 * time.Second   // Timeout for retry for request EXCLUDING clientTimeout
 	defaultJWTTimeout             = 60 * time.Second
@@ -71,6 +72,7 @@ type Config struct {
 	RequestTimeout         time.Duration // request retry timeout EXCLUDING network roundtrip and read out http response
 	JWTExpireTimeout       time.Duration // JWT expire after timeout
 	ClientTimeout          time.Duration // Timeout for network round trip + read out http response
+	JWTClientTimeout       time.Duration // Timeout for network round trip + read out http response used when JWT token auth is taking place
 	ExternalBrowserTimeout time.Duration // Timeout for external browser login
 
 	Application  string           // application name.
@@ -169,6 +171,9 @@ func DSN(cfg *Config) (dsn string, err error) {
 	if cfg.ClientTimeout != defaultClientTimeout {
 		params.Add("clientTimeout", strconv.FormatInt(int64(cfg.ClientTimeout/time.Second), 10))
 	}
+	if cfg.JWTClientTimeout != defaultJWTClientTimeout {
+		params.Add("jwtClientTimeout", strconv.FormatInt(int64(cfg.JWTClientTimeout/time.Second), 10))
+	}
 	if cfg.LoginTimeout != defaultLoginTimeout {
 		params.Add("loginTimeout", strconv.FormatInt(int64(cfg.LoginTimeout/time.Second), 10))
 	}
@@ -434,6 +439,9 @@ func fillMissingConfigParameters(cfg *Config) error {
 	if cfg.ClientTimeout == 0 {
 		cfg.ClientTimeout = defaultClientTimeout
 	}
+	if cfg.JWTClientTimeout == 0 {
+		cfg.JWTClientTimeout = defaultJWTClientTimeout
+	}
 	if cfg.ExternalBrowserTimeout == 0 {
 		cfg.ExternalBrowserTimeout = defaultExternalBrowserTimeout
 	}
@@ -566,6 +574,11 @@ func parseDSNParams(cfg *Config, params string) (err error) {
 			if err != nil {
 				return
 			}
+		case "jwtClientTimeout":
+			cfg.JWTClientTimeout, err = parseTimeout(value)
+			if err != nil {
+				return
+			}
 		case "loginTimeout":
 			cfg.LoginTimeout, err = parseTimeout(value)
 			if err != nil {
diff --git a/restful.go b/restful.go
index 04da332..2acbbd7 100644
--- a/restful.go
+++ b/restful.go
@@ -42,8 +42,9 @@ 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)
+	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)
+	funcAuthPostType func(context.Context, *http.Client, *url.URL, map[string]string, []byte, time.Duration, bool) (*http.Response, error)
 )
 
 type snowflakeRestful struct {
@@ -54,6 +55,7 @@ type snowflakeRestful struct {
 	RequestTimeout time.Duration // request timeout
 
 	Client        *http.Client
+	JWTClient     *http.Client
 	TokenAccessor TokenAccessor
 	HeartBeat     *heartbeat
 
@@ -63,11 +65,12 @@ type snowflakeRestful struct {
 	FuncPostQueryHelper func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, UUID, *Config) (*execResponse, error)
 	FuncPost            funcPostType
 	FuncGet             funcGetType
+	FuncAuthPost        funcAuthPostType
 	FuncRenewSession    func(context.Context, *snowflakeRestful, time.Duration) error
-	FuncPostAuth        func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration) (*authResponse, error)
 	FuncCloseSession    func(context.Context, *snowflakeRestful, time.Duration) error
 	FuncCancelQuery     func(context.Context, *snowflakeRestful, UUID, time.Duration) error
 
+	FuncPostAuth     func(context.Context, *snowflakeRestful, *http.Client, *url.Values, map[string]string, []byte, time.Duration) (*authResponse, error)
 	FuncPostAuthSAML func(context.Context, *snowflakeRestful, map[string]string, []byte, time.Duration) (*authResponse, error)
 	FuncPostAuthOKTA func(context.Context, *snowflakeRestful, map[string]string, []byte, string, time.Duration) (*authOKTAResponse, error)
 	FuncGetSSO       func(context.Context, *snowflakeRestful, *url.Values, map[string]string, string, time.Duration) ([]byte, error)
@@ -92,6 +95,16 @@ func (sr *snowflakeRestful) getFullURL(path string, params *url.Values) *url.URL
 	return ret
 }
 
+// We need separate client for JWT, because if token processing takes too long, token may be already expired.
+func (sr *snowflakeRestful) getClientFor(authType AuthType) *http.Client {
+	switch authType {
+	case AuthTypeJwt:
+		return sr.JWTClient
+	default:
+		return sr.Client
+	}
+}
+
 // Renew the snowflake session if the current token is still the stale token specified
 func (sr *snowflakeRestful) renewExpiredSessionToken(ctx context.Context, timeout time.Duration, expiredToken string) error {
 	err := sr.TokenAccessor.Lock()
@@ -161,6 +174,19 @@ func getRestful(
 		ctx, sr.Client, http.NewRequest, fullURL, headers, timeout).execute()
 }
 
+func postAuthRestful(
+	ctx context.Context,
+	client *http.Client,
+	fullURL *url.URL,
+	headers map[string]string,
+	body []byte,
+	timeout time.Duration,
+	raise4XX bool) (
+	*http.Response, error) {
+	return newRetryHTTP(
+		ctx, client, http.NewRequest, fullURL, headers, timeout).doPost().setBody(body).doRaise4XX(raise4XX).execute()
+}
+
 func postRestfulQuery(
 	ctx context.Context,
 	sr *snowflakeRestful,


Test patch:
diff --git a/auth_test.go b/auth_test.go
index f1d1bed..190bd39 100644
--- a/auth_test.go
+++ b/auth_test.go
@@ -10,6 +10,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"net/http"
 	"net/url"
 	"os"
 	"runtime"
@@ -22,54 +23,54 @@ import (
 func TestUnitPostAuth(t *testing.T) {
 	sr := &snowflakeRestful{
 		TokenAccessor: getSimpleTokenAccessor(),
-		FuncPost:      postTestAfterRenew,
+		FuncAuthPost:  postAuthTestAfterRenew,
 	}
 	var err error
-	_, err = postAuth(context.TODO(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
+	_, err = postAuth(context.TODO(), sr, sr.Client, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
-	sr.FuncPost = postTestError
-	_, err = postAuth(context.TODO(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
+	sr.FuncAuthPost = postAuthTestError
+	_, err = postAuth(context.TODO(), sr, sr.Client, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
 	if err == nil {
 		t.Fatal("should have failed to auth for unknown reason")
 	}
-	sr.FuncPost = postTestAppBadGatewayError
-	_, err = postAuth(context.TODO(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
+	sr.FuncAuthPost = postAuthTestAppBadGatewayError
+	_, err = postAuth(context.TODO(), sr, sr.Client, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
 	if err == nil {
 		t.Fatal("should have failed to auth for unknown reason")
 	}
-	sr.FuncPost = postTestAppForbiddenError
-	_, err = postAuth(context.TODO(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
+	sr.FuncAuthPost = postAuthTestAppForbiddenError
+	_, err = postAuth(context.TODO(), sr, sr.Client, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
 	if err == nil {
 		t.Fatal("should have failed to auth for unknown reason")
 	}
-	sr.FuncPost = postTestAppUnexpectedError
-	_, err = postAuth(context.TODO(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
+	sr.FuncAuthPost = postAuthTestAppUnexpectedError
+	_, err = postAuth(context.TODO(), sr, sr.Client, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
 	if err == nil {
 		t.Fatal("should have failed to auth for unknown reason")
 	}
 }
 
-func postAuthFailServiceIssue(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthFailServiceIssue(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return nil, &SnowflakeError{
 		Number: ErrCodeServiceUnavailable,
 	}
 }
 
-func postAuthFailWrongAccount(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthFailWrongAccount(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return nil, &SnowflakeError{
 		Number: ErrCodeFailedToConnect,
 	}
 }
 
-func postAuthFailUnknown(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthFailUnknown(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return nil, &SnowflakeError{
 		Number: ErrFailedToAuth,
 	}
 }
 
-func postAuthSuccessWithErrorCode(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthSuccessWithErrorCode(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return &authResponse{
 		Success: false,
 		Code:    "98765",
@@ -77,7 +78,7 @@ func postAuthSuccessWithErrorCode(_ context.Context, _ *snowflakeRestful, _ *url
 	}, nil
 }
 
-func postAuthSuccessWithInvalidErrorCode(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthSuccessWithInvalidErrorCode(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return &authResponse{
 		Success: false,
 		Code:    "abcdef",
@@ -85,7 +86,7 @@ func postAuthSuccessWithInvalidErrorCode(_ context.Context, _ *snowflakeRestful,
 	}, nil
 }
 
-func postAuthSuccess(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthSuccess(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return &authResponse{
 		Success: true,
 		Data: authResponseMain{
@@ -98,7 +99,7 @@ func postAuthSuccess(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ ma
 	}, nil
 }
 
-func postAuthCheckSAMLResponse(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckSAMLResponse(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
 	var ar authRequest
 	if err := json.Unmarshal(jsonBody, &ar); err != nil {
 		return nil, err
@@ -123,6 +124,7 @@ func postAuthCheckSAMLResponse(_ context.Context, _ *snowflakeRestful, _ *url.Va
 func postAuthCheckOAuth(
 	_ context.Context,
 	_ *snowflakeRestful,
+	_ *http.Client,
 	_ *url.Values, _ map[string]string,
 	jsonBody []byte,
 	_ time.Duration) (*authResponse, error) {
@@ -151,7 +153,7 @@ func postAuthCheckOAuth(
 	}, nil
 }
 
-func postAuthCheckPasscode(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckPasscode(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
 	var ar authRequest
 	if err := json.Unmarshal(jsonBody, &ar); err != nil {
 		return nil, err
@@ -171,7 +173,7 @@ func postAuthCheckPasscode(_ context.Context, _ *snowflakeRestful, _ *url.Values
 	}, nil
 }
 
-func postAuthCheckPasscodeInPassword(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckPasscodeInPassword(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
 	var ar authRequest
 	if err := json.Unmarshal(jsonBody, &ar); err != nil {
 		return nil, err
@@ -193,7 +195,7 @@ func postAuthCheckPasscodeInPassword(_ context.Context, _ *snowflakeRestful, _ *
 
 // JWT token validate callback function to check the JWT token
 // It uses the public key paired with the testPrivKey
-func postAuthCheckJWTToken(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckJWTToken(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
 	var ar authRequest
 	if err := json.Unmarshal(jsonBody, &ar); err != nil {
 		return nil, err
@@ -229,7 +231,7 @@ func postAuthCheckJWTToken(_ context.Context, _ *snowflakeRestful, _ *url.Values
 	}, nil
 }
 
-func postAuthCheckUsernamePasswordMfa(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckUsernamePasswordMfa(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
 	var ar authRequest
 	if err := json.Unmarshal(jsonBody, &ar); err != nil {
 		return nil, err
@@ -251,7 +253,7 @@ func postAuthCheckUsernamePasswordMfa(_ context.Context, _ *snowflakeRestful, _
 	}, nil
 }
 
-func postAuthCheckUsernamePasswordMfaToken(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckUsernamePasswordMfaToken(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
 	var ar authRequest
 	if err := json.Unmarshal(jsonBody, &ar); err != nil {
 		return nil, err
@@ -273,7 +275,7 @@ func postAuthCheckUsernamePasswordMfaToken(_ context.Context, _ *snowflakeRestfu
 	}, nil
 }
 
-func postAuthCheckUsernamePasswordMfaFailed(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckUsernamePasswordMfaFailed(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
 	var ar authRequest
 	if err := json.Unmarshal(jsonBody, &ar); err != nil {
 		return nil, err
@@ -290,7 +292,7 @@ func postAuthCheckUsernamePasswordMfaFailed(_ context.Context, _ *snowflakeRestf
 	}, nil
 }
 
-func postAuthCheckExternalBrowser(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckExternalBrowser(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
 	var ar authRequest
 	if err := json.Unmarshal(jsonBody, &ar); err != nil {
 		return nil, err
@@ -312,7 +314,7 @@ func postAuthCheckExternalBrowser(_ context.Context, _ *snowflakeRestful, _ *url
 	}, nil
 }
 
-func postAuthCheckExternalBrowserToken(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckExternalBrowserToken(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
 	var ar authRequest
 	if err := json.Unmarshal(jsonBody, &ar); err != nil {
 		return nil, err
@@ -334,7 +336,7 @@ func postAuthCheckExternalBrowserToken(_ context.Context, _ *snowflakeRestful, _
 	}, nil
 }
 
-func postAuthCheckExternalBrowserFailed(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckExternalBrowserFailed(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
 	var ar authRequest
 	if err := json.Unmarshal(jsonBody, &ar); err != nil {
 		return nil, err
diff --git a/dsn_test.go b/dsn_test.go
index 314c667..bc5dbd5 100644
--- a/dsn_test.go
+++ b/dsn_test.go
@@ -38,6 +38,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -52,6 +53,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -64,6 +66,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -77,6 +80,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -91,6 +95,7 @@ func TestParseDSN(t *testing.T) {
 				ValidateDefaultParameters: ConfigBoolTrue,
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -105,6 +110,7 @@ func TestParseDSN(t *testing.T) {
 				ValidateDefaultParameters: ConfigBoolTrue,
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -118,6 +124,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -131,6 +138,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -144,6 +152,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -158,6 +167,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -172,6 +182,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -186,6 +197,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -200,6 +212,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -214,6 +227,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -228,6 +242,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -241,6 +256,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -254,6 +270,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -267,6 +284,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -283,6 +301,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -297,6 +316,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
 		},
@@ -310,6 +330,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -322,6 +343,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -341,6 +363,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeInsecure,
@@ -356,6 +379,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -406,6 +430,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -420,6 +445,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -439,6 +465,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -453,6 +480,7 @@ func TestParseDSN(t *testing.T) {
 				OCSPFailOpen:              OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -466,6 +494,7 @@ func TestParseDSN(t *testing.T) {
 				Database: "db", Schema: "s", OCSPFailOpen: OCSPFailOpenTrue,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -479,6 +508,7 @@ func TestParseDSN(t *testing.T) {
 				Database: "db", Schema: "s", OCSPFailOpen: OCSPFailOpenFalse,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailClosed,
@@ -492,6 +522,7 @@ func TestParseDSN(t *testing.T) {
 				Database: "db", Schema: "s", OCSPFailOpen: OCSPFailOpenFalse, InsecureMode: true,
 				ValidateDefaultParameters: ConfigBoolTrue,
 				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
 				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeInsecure,
@@ -504,6 +535,7 @@ func TestParseDSN(t *testing.T) {
 				Protocol: "https", Host: "account.snowflakecomputing.com", Port: 443,
 				Database: "db", Schema: "s", ValidateDefaultParameters: ConfigBoolTrue, OCSPFailOpen: OCSPFailOpenTrue,
 				ClientTimeout:          defaultClientTimeout,
+				JWTClientTimeout:       defaultJWTClientTimeout,
 				ExternalBrowserTimeout: defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -516,6 +548,7 @@ func TestParseDSN(t *testing.T) {
 				Protocol: "https", Host: "account.snowflakecomputing.com", Port: 443,
 				Database: "db", Schema: "s", ValidateDefaultParameters: ConfigBoolFalse, OCSPFailOpen: OCSPFailOpenTrue,
 				ClientTimeout:          defaultClientTimeout,
+				JWTClientTimeout:       defaultJWTClientTimeout,
 				ExternalBrowserTimeout: defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -528,18 +561,20 @@ func TestParseDSN(t *testing.T) {
 				Protocol: "https", Host: "a.r.c.snowflakecomputing.com", Port: 443,
 				Database: "db", Schema: "s", ValidateDefaultParameters: ConfigBoolFalse, OCSPFailOpen: OCSPFailOpenTrue,
 				ClientTimeout:          defaultClientTimeout,
+				JWTClientTimeout:       defaultJWTClientTimeout,
 				ExternalBrowserTimeout: defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
 			err:      nil,
 		},
 		{
-			dsn: "u:p@a.r.c.snowflakecomputing.com/db/s?account=a.r.c&clientTimeout=300",
+			dsn: "u:p@a.r.c.snowflakecomputing.com/db/s?account=a.r.c&clientTimeout=300&jwtClientTimeout=45",
 			config: &Config{
 				Account: "a", User: "u", Password: "p",
 				Protocol: "https", Host: "a.r.c.snowflakecomputing.com", Port: 443,
 				Database: "db", Schema: "s", ValidateDefaultParameters: ConfigBoolTrue, OCSPFailOpen: OCSPFailOpenTrue,
 				ClientTimeout:          300 * time.Second,
+				JWTClientTimeout:       45 * time.Second,
 				ExternalBrowserTimeout: defaultExternalBrowserTimeout,
 			},
 			ocspMode: ocspModeFailOpen,
@@ -627,6 +662,10 @@ func TestParseDSN(t *testing.T) {
 				t.Fatalf("%d: Failed to match ClientTimeout. expected: %v, got: %v",
 					i, test.config.ClientTimeout, cfg.ClientTimeout)
 			}
+			if test.config.JWTClientTimeout != cfg.JWTClientTimeout {
+				t.Fatalf("%d: Failed to match JWTClientTimeout. expected: %v, got: %v",
+					i, test.config.JWTClientTimeout, cfg.JWTClientTimeout)
+			}
 			if test.config.ExternalBrowserTimeout != cfg.ExternalBrowserTimeout {
 				t.Fatalf("%d: Failed to match ExternalBrowserTimeout. expected: %v, got: %v",
 					i, test.config.ExternalBrowserTimeout, cfg.ExternalBrowserTimeout)
@@ -917,12 +956,13 @@ func TestDSN(t *testing.T) {
 		},
 		{
 			cfg: &Config{
-				User:          "u",
-				Password:      "p",
-				Account:       "a.b.c",
-				ClientTimeout: 300 * time.Second,
+				User:             "u",
+				Password:         "p",
+				Account:          "a.b.c",
+				ClientTimeout:    300 * time.Second,
+				JWTClientTimeout: 60 * time.Second,
 			},
-			dsn: "u:p@a.b.c.snowflakecomputing.com:443?clientTimeout=300&ocspFailOpen=true&region=b.c&validateDefaultParameters=true",
+			dsn: "u:p@a.b.c.snowflakecomputing.com:443?clientTimeout=300&jwtClientTimeout=60&ocspFailOpen=true&region=b.c&validateDefaultParameters=true",
 		},
 	}
 	for _, test := range testcases {
diff --git a/restful_test.go b/restful_test.go
index 01b2adb..a49b2a3 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -22,6 +22,13 @@ func postTestError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[str
 	}, errors.New("failed to run post method")
 }
 
+func postAuthTestError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+	return &http.Response{
+		StatusCode: http.StatusOK,
+		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
+	}, errors.New("failed to run post method")
+}
+
 func postTestSuccessButInvalidJSON(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusOK,
@@ -36,6 +43,13 @@ func postTestAppBadGatewayError(_ context.Context, _ *snowflakeRestful, _ *url.U
 	}, nil
 }
 
+func postAuthTestAppBadGatewayError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+	return &http.Response{
+		StatusCode: http.StatusBadGateway,
+		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
+	}, nil
+}
+
 func postTestAppForbiddenError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusForbidden,
@@ -43,7 +57,14 @@ func postTestAppForbiddenError(_ context.Context, _ *snowflakeRestful, _ *url.UR
 	}, nil
 }
 
-func postTestAppUnexpectedError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+func postAuthTestAppForbiddenError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+	return &http.Response{
+		StatusCode: http.StatusForbidden,
+		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
+	}, nil
+}
+
+func postAuthTestAppUnexpectedError(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
 	return &http.Response{
 		StatusCode: http.StatusInsufficientStorage,
 		Body:       &fakeResponseBody{body: []byte{0x12, 0x34}},
@@ -89,6 +110,26 @@ func postTestRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[str
 	}, nil
 }
 
+func postAuthTestAfterRenew(_ context.Context, _ *http.Client, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+	dd := &execResponseData{}
+	er := &execResponse{
+		Data:    *dd,
+		Message: "",
+		Code:    "",
+		Success: true,
+	}
+
+	ba, err := json.Marshal(er)
+	logger.Infof("encoded JSON: %v", ba)
+	if err != nil {
+		panic(err)
+	}
+	return &http.Response{
+		StatusCode: http.StatusOK,
+		Body:       &fakeResponseBody{body: ba},
+	}, nil
+}
+
 func postTestAfterRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
 	dd := &execResponseData{}
 	er := &execResponse{