Repository: snowflakedb/gosnowflake
Pull request #325 (2020-08-17 23:08:15)

SNOW-183510 Update context behavior function signatures
### Description
Added
- Update chunk download helper function signature to return error and handle error
- Parity in function signatures within restful.go / restful object
- Handle request timeout errors for restful object
- Add and expose the query ID for DMLs as a part of Snowflake result

Not added
- New describeContext function (sets parameter describeOnly) at the statement level. Requires exposing this function by making Snowflake Statement object public as well.

### 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-stakeda: @sfc-gh-jbahk can you elaborate how query id should be exposed?
sfc-gh-jbahk: @sfc-gh-stakeda sure, we essentially just add a `queryID` field for the user to be able to call on their `SnowflakeResult` object. It requires simply adding an extra field to the rows (for select queries) and the snowflakeresults (for dml) and creating a function that retrieves it.
sfc-gh-stakeda: @sfc-gh-jbahk , can you find the best way to expose the query id?
sfc-gh-stakeda: @sfc-gh-jbahk https://github.com/sigmacomputing/gosnowflake/commit/9a539cac11808dddcab96afcb49c9a40f55fb222 is missing.
sfc-gh-jbahk: @sfc-gh-stakeda That commit is already implemented here: [https://github.com/snowflakedb/gosnowflake/blob/SNOW-183510-update-context-behavior-function-signatures/retry.go#L282-L291](url)
sfc-gh-stakeda: ok, thanks for the confirmation.
@sfc-gh-jbahk can you please summarize what you merged and not merged here?
sfc-gh-jbahk: @sfc-gh-stakeda I updated the PR description.
sfc-gh-stakeda: The describe functionality can be mapped to https://golang.org/pkg/database/sql/#Conn.PrepareContext, which returns `Stmt`. Let's create another ticket for PrepareContext implementation.
sfc-gh-stakeda: How can the application specify the timeout? Can you add an example as a test case?
sfc-gh-jbahk: @sfc-gh-stakeda The tests are the ones modified in restful_test.go (`time.Second` is the controlled variable) regarding specifying the timeout.

Patch:
diff --git a/auth.go b/auth.go
index 6315236..58a2293 100644
--- a/auth.go
+++ b/auth.go
@@ -188,6 +188,7 @@ type authResponse struct {
 }
 
 func postAuth(
+	ctx context.Context,
 	sr *snowflakeRestful,
 	params *url.Values,
 	headers map[string]string,
@@ -199,7 +200,7 @@ func postAuth(
 
 	fullURL := sr.getFullURL(loginRequestPath, params)
 	glog.V(2).Infof("full URL: %v", fullURL)
-	resp, err := sr.FuncPost(context.TODO(), sr, fullURL, headers, body, timeout, true)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, true)
 	if err != nil {
 		return nil, err
 	}
@@ -261,6 +262,7 @@ func getHeaders() map[string]string {
 
 // Used to authenticate the user with Snowflake.
 func authenticate(
+	ctx context.Context,
 	sc *snowflakeConn,
 	samlResponse []byte,
 	proofKey []byte,
@@ -348,7 +350,7 @@ func authenticate(
 	glog.V(2).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(sc.rest, params, headers, jsonBody, sc.rest.LoginTimeout)
+	respd, err := sc.rest.FuncPostAuth(ctx, sc.rest, params, headers, jsonBody, sc.rest.LoginTimeout)
 	if err != nil {
 		return nil, err
 	}
diff --git a/authexternalbrowser.go b/authexternalbrowser.go
index e8f3911..ff13392 100644
--- a/authexternalbrowser.go
+++ b/authexternalbrowser.go
@@ -6,6 +6,7 @@ package gosnowflake
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -81,6 +82,7 @@ func openBrowser(idpURL string) error {
 // Note: FuncPostAuthSaml will return a fully qualified error if
 // there is something wrong getting data from Snowflake.
 func getIdpURLProofKey(
+	ctx context.Context,
 	sr *snowflakeRestful,
 	authenticator string,
 	application string,
@@ -117,7 +119,7 @@ func getIdpURLProofKey(
 		return "", "", err
 	}
 
-	respd, err := sr.FuncPostAuthSAML(sr, headers, jsonBody, sr.LoginTimeout)
+	respd, err := sr.FuncPostAuthSAML(ctx, sr, headers, jsonBody, sr.LoginTimeout)
 	if err != nil {
 		return "", "", err
 	}
@@ -162,6 +164,7 @@ func getTokenFromResponse(response string) (string, error) {
 // - Snowflake directs the user back to the driver
 // - authenticate is complete!
 func authenticateByExternalBrowser(
+	ctx context.Context,
 	sr *snowflakeRestful,
 	authenticator string,
 	application string,
@@ -177,7 +180,7 @@ func authenticateByExternalBrowser(
 
 	callbackPort := l.Addr().(*net.TCPAddr).Port
 	idpURL, proofKey, err := getIdpURLProofKey(
-		sr, authenticator, application, account, callbackPort)
+		ctx, sr, authenticator, application, account, callbackPort)
 	if err != nil {
 		return nil, nil, err
 	}
diff --git a/authokta.go b/authokta.go
index 2821de9..c937507 100644
--- a/authokta.go
+++ b/authokta.go
@@ -51,6 +51,7 @@ Explanation:
 	another SP.
 */
 func authenticateBySAML(
+	ctx context.Context,
 	sr *snowflakeRestful,
 	oktaURL *url.URL,
 	application string,
@@ -85,7 +86,7 @@ func authenticateBySAML(
 		return nil, err
 	}
 	glog.V(2).Infof("PARAMS for Auth: %v, %v", params, sr)
-	respd, err := sr.FuncPostAuthSAML(sr, headers, jsonBody, sr.LoginTimeout)
+	respd, err := sr.FuncPostAuthSAML(ctx, sr, headers, jsonBody, sr.LoginTimeout)
 	if err != nil {
 		return nil, err
 	}
@@ -131,7 +132,7 @@ func authenticateBySAML(
 	if err != nil {
 		return nil, err
 	}
-	respa, err := sr.FuncPostAuthOKTA(sr, headers, jsonBody, respd.Data.TokenURL, sr.LoginTimeout)
+	respa, err := sr.FuncPostAuthOKTA(ctx, sr, headers, jsonBody, respd.Data.TokenURL, sr.LoginTimeout)
 	if err != nil {
 		return nil, err
 	}
@@ -143,7 +144,7 @@ func authenticateBySAML(
 
 	headers = make(map[string]string)
 	headers["accept"] = "*/*"
-	bd, err := sr.FuncGetSSO(sr, params, headers, respd.Data.SSOURL, sr.LoginTimeout)
+	bd, err := sr.FuncGetSSO(ctx, sr, params, headers, respd.Data.SSOURL, sr.LoginTimeout)
 	if err != nil {
 		return nil, err
 	}
@@ -199,6 +200,7 @@ func isPrefixEqual(u1 *url.URL, u2 *url.URL) bool {
 // Makes a request to /session/authenticator-request to get SAML Information,
 // such as the IDP Url and Proof Key, depending on the authenticator
 func postAuthSAML(
+	ctx context.Context,
 	sr *snowflakeRestful,
 	headers map[string]string,
 	body []byte,
@@ -210,7 +212,7 @@ func postAuthSAML(
 	fullURL := sr.getFullURL(authenticatorRequestPath, params)
 
 	glog.V(2).Infof("fullURL: %v", fullURL)
-	resp, err := sr.FuncPost(context.TODO(), sr, fullURL, headers, body, timeout, true)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, true)
 	if err != nil {
 		return nil, err
 	}
@@ -259,6 +261,7 @@ func postAuthSAML(
 }
 
 func postAuthOKTA(
+	ctx context.Context,
 	sr *snowflakeRestful,
 	headers map[string]string,
 	body []byte,
@@ -270,7 +273,7 @@ func postAuthOKTA(
 	if err != nil {
 		return nil, err
 	}
-	resp, err := sr.FuncPost(context.TODO(), sr, targetURL, headers, body, timeout, false)
+	resp, err := sr.FuncPost(ctx, sr, targetURL, headers, body, timeout, false)
 	if err != nil {
 		return nil, err
 	}
@@ -303,6 +306,7 @@ func postAuthOKTA(
 }
 
 func getSSO(
+	ctx context.Context,
 	sr *snowflakeRestful,
 	params *url.Values,
 	headers map[string]string,
@@ -315,7 +319,7 @@ func getSSO(
 	}
 	fullURL.RawQuery = params.Encode()
 	glog.V(2).Infof("fullURL: %v", fullURL)
-	resp, err := sr.FuncGet(context.TODO(), sr, fullURL, headers, timeout)
+	resp, err := sr.FuncGet(ctx, sr, fullURL, headers, timeout)
 	if err != nil {
 		return nil, err
 	}
diff --git a/benchmark/jsonresultset/Makefile b/benchmark/jsonresultset/Makefile
index fff9ab3..31efd8c 100644
--- a/benchmark/jsonresultset/Makefile
+++ b/benchmark/jsonresultset/Makefile
@@ -3,7 +3,7 @@ SHELL := /bin/bash
 SRC = $(shell find . -type f -name '*.go' -not -path "./vendor/*")
 
 setup:
-	go get github.com/golang/lint/golint
+	go get -u golang.org/lint/golint
 	go get github.com/Songmu/make2help/cmd/make2help
 
 ## Benchmark
diff --git a/benchmark/largesetresult/Makefile b/benchmark/largesetresult/Makefile
index af3ce1c..977ae9c 100644
--- a/benchmark/largesetresult/Makefile
+++ b/benchmark/largesetresult/Makefile
@@ -3,7 +3,7 @@ SHELL := /bin/bash
 SRC = $(shell find . -type f -name '*.go' -not -path "./vendor/*")
 
 setup:
-	go get github.com/golang/lint/golint
+	go get -u golang.org/lint/golint
 	go get github.com/Songmu/make2help/cmd/make2help
 
 ## Benchmark
diff --git a/connection.go b/connection.go
index 286d724..26dc36d 100644
--- a/connection.go
+++ b/connection.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2019 Snowflake Computing Inc. All right reserved.
+// Copyright (c) 2017-2020 Snowflake Computing Inc. All right reserved.
 
 package gosnowflake
 
@@ -60,7 +60,11 @@ func (sc *snowflakeConn) isMultiStmt(data execResponseData) bool {
 
 func (sc *snowflakeConn) exec(
 	ctx context.Context,
-	query string, noResult bool, isInternal bool, bindings []driver.NamedValue) (*execResponse, error) {
+	query string,
+	noResult bool,
+	isInternal bool,
+	bindings []driver.NamedValue) (
+	*execResponse, error) {
 	var err error
 	counter := atomic.AddUint64(&sc.SequenceCounter, 1) // query sequence counter
 
@@ -120,14 +124,14 @@ func (sc *snowflakeConn) exec(
 	requestID := uuid.New()
 	data, err = sc.rest.FuncPostQuery(ctx, sc.rest, &url.Values{}, headers, jsonBody, sc.rest.RequestTimeout, &requestID)
 	if err != nil {
-		return nil, err
+		return data, err
 	}
 	var code int
 	if data.Code != "" {
 		code, err = strconv.Atoi(data.Code)
 		if err != nil {
 			code = -1
-			return nil, err
+			return data, err
 		}
 	} else {
 		code = -1
@@ -192,13 +196,14 @@ func (sc *snowflakeConn) Close() (err error) {
 	glog.V(2).Infoln("Close")
 	sc.stopHeartBeat()
 
-	err = sc.rest.FuncCloseSession(sc.rest)
+	err = sc.rest.FuncCloseSession(context.TODO(), sc.rest, sc.rest.RequestTimeout)
 	if err != nil {
 		glog.V(2).Info(err)
 	}
 	sc.cleanup()
 	return nil
 }
+
 func (sc *snowflakeConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
 	glog.V(2).Infoln("Prepare")
 	if sc.rest == nil {
@@ -223,8 +228,21 @@ func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []d
 	// TODO: handle noResult and isInternal
 	data, err := sc.exec(ctx, query, false, false, args)
 	if err != nil {
+		glog.V(2).Infof("error: %v", err)
+		if data != nil {
+			code, err := strconv.Atoi(data.Code)
+			if err != nil {
+				return nil, err
+			}
+			return nil, &SnowflakeError{
+				Number:   code,
+				SQLState: data.Data.SQLState,
+				Message:  err.Error(),
+				QueryID:  data.Data.QueryID}
+		}
 		return nil, err
 	}
+
 	var updatedRows int64 = 0
 	if sc.isDml(data.Data.StatementTypeID) {
 		// collects all values from the returned row sets
@@ -235,18 +253,44 @@ func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []d
 		glog.V(2).Infof("number of updated rows: %#v", updatedRows)
 		return &snowflakeResult{
 			affectedRows: updatedRows,
-			insertID:     -1}, nil // last insert id is not supported by Snowflake
+			insertID:     -1,
+			queryID:      sc.QueryID,
+		}, nil // last insert id is not supported by Snowflake
 	} else if sc.isMultiStmt(data.Data) {
 		childResults := getChildResults(data.Data.ResultIDs, data.Data.ResultTypes)
 		for _, child := range childResults {
 			resultPath := fmt.Sprintf("/queries/%s/result", child.id)
 			childData, err := sc.getQueryResult(ctx, resultPath)
 			if err != nil {
+				glog.V(2).Infof("error: %v", err)
+				code, err := strconv.Atoi(childData.Code)
+				if err != nil {
+					return nil, err
+				}
+				if childData != nil {
+					return nil, &SnowflakeError{
+						Number:   code,
+						SQLState: childData.Data.SQLState,
+						Message:  err.Error(),
+						QueryID:  childData.Data.QueryID}
+				}
 				return nil, err
 			}
 			if sc.isDml(childData.Data.StatementTypeID) {
 				count, err := updateRows(childData.Data)
 				if err != nil {
+					glog.V(2).Infof("error: %v", err)
+					if childData != nil {
+						code, err := strconv.Atoi(childData.Code)
+						if err != nil {
+							return nil, err
+						}
+						return nil, &SnowflakeError{
+							Number:   code,
+							SQLState: childData.Data.SQLState,
+							Message:  err.Error(),
+							QueryID:  childData.Data.QueryID}
+					}
 					return nil, err
 				}
 				updatedRows += count
@@ -255,7 +299,9 @@ func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []d
 		glog.V(2).Infof("number of updated rows: %#v", updatedRows)
 		return &snowflakeResult{
 			affectedRows: updatedRows,
-			insertID:     -1}, nil
+			insertID:     -1,
+			queryID:      sc.QueryID,
+		}, nil
 	}
 	glog.V(2).Info("DDL")
 	return driver.ResultNoRows, nil
@@ -270,6 +316,17 @@ func (sc *snowflakeConn) QueryContext(ctx context.Context, query string, args []
 	data, err := sc.exec(ctx, query, false, false, args)
 	if err != nil {
 		glog.V(2).Infof("error: %v", err)
+		if data != nil {
+			code, err := strconv.Atoi(data.Code)
+			if err != nil {
+				return nil, err
+			}
+			return nil, &SnowflakeError{
+				Number:   code,
+				SQLState: data.Data.SQLState,
+				Message:  err.Error(),
+				QueryID:  data.Data.QueryID}
+		}
 		return nil, err
 	}
 
@@ -295,6 +352,7 @@ func (sc *snowflakeConn) QueryContext(ctx context.Context, query string, args []
 			RowSetBase64: data.Data.RowSetBase64,
 		},
 	}
+	rows.queryID = sc.QueryID
 
 	if sc.isMultiStmt(data.Data) {
 		childResults := getChildResults(data.Data.ResultIDs, data.Data.ResultTypes)
@@ -305,6 +363,18 @@ func (sc *snowflakeConn) QueryContext(ctx context.Context, query string, args []
 			resultPath := fmt.Sprintf("/queries/%s/result", child.id)
 			childData, err := sc.getQueryResult(ctx, resultPath)
 			if err != nil {
+				glog.V(2).Infof("error: %v", err)
+				if childData != nil {
+					code, err := strconv.Atoi(childData.Code)
+					if err != nil {
+						return nil, err
+					}
+					return nil, &SnowflakeError{
+						Number:   code,
+						SQLState: childData.Data.SQLState,
+						Message:  err.Error(),
+						QueryID:  childData.Data.QueryID}
+				}
 				return nil, err
 			}
 			if !firstResultSet {
@@ -430,7 +500,7 @@ func getChildResults(IDs string, types string) []childResult {
 	return res
 }
 
-func (sc *snowflakeConn) getQueryResult(ctx context.Context, resultPath string) (execResponse, error) {
+func (sc *snowflakeConn) getQueryResult(ctx context.Context, resultPath string) (*execResponse, error) {
 	headers := make(map[string]string)
 	headers["Content-Type"] = headerContentTypeApplicationJSON
 	headers["accept"] = headerAcceptTypeApplicationSnowflake
@@ -450,14 +520,14 @@ func (sc *snowflakeConn) getQueryResult(ctx context.Context, resultPath string)
 	if err != nil {
 		glog.V(1).Infof("failed to get response. err: %v", err)
 		glog.Flush()
-		return execResponse{}, err
+		return nil, err
 	}
-	var respd execResponse
+	var respd *execResponse
 	err = json.NewDecoder(res.Body).Decode(&respd)
 	if err != nil {
 		glog.V(1).Infof("failed to decode JSON. err: %v", err)
 		glog.Flush()
-		return execResponse{}, err
+		return nil, err
 	}
 	return respd, nil
 }
diff --git a/driver.go b/driver.go
index 0f38626..a9982fa 100644
--- a/driver.go
+++ b/driver.go
@@ -3,6 +3,7 @@
 package gosnowflake
 
 import (
+	"context"
 	"database/sql"
 	"database/sql/driver"
 	"net/http"
@@ -18,6 +19,7 @@ func (d SnowflakeDriver) Open(dsn string) (driver.Conn, error) {
 	sc := &snowflakeConn{
 		SequenceCounter: 0,
 	}
+	ctx := context.TODO()
 	sc.cfg, err = ParseDSN(dsn)
 	if err != nil {
 		sc.cleanup()
@@ -65,6 +67,7 @@ func (d SnowflakeDriver) Open(dsn string) (driver.Conn, error) {
 	switch sc.cfg.Authenticator {
 	case AuthTypeExternalBrowser:
 		samlResponse, proofKey, err = authenticateByExternalBrowser(
+			ctx,
 			sc.rest,
 			sc.cfg.Authenticator.String(),
 			sc.cfg.Application,
@@ -77,6 +80,7 @@ func (d SnowflakeDriver) Open(dsn string) (driver.Conn, error) {
 		}
 	case AuthTypeOkta:
 		samlResponse, err = authenticateBySAML(
+			ctx,
 			sc.rest,
 			sc.cfg.OktaURL,
 			sc.cfg.Application,
@@ -89,6 +93,7 @@ func (d SnowflakeDriver) Open(dsn string) (driver.Conn, error) {
 		}
 	}
 	authData, err = authenticate(
+		ctx,
 		sc,
 		samlResponse,
 		proofKey)
diff --git a/dsn.go b/dsn.go
index ee72154..becd6e8 100644
--- a/dsn.go
+++ b/dsn.go
@@ -174,7 +174,6 @@ func DSN(cfg *Config) (dsn string, err error) {
 		keyBase64 := base64.URLEncoding.EncodeToString(privateKeyInBytes)
 		params.Add("privateKey", keyBase64)
 	}
-
 	if cfg.InsecureMode {
 		params.Add("insecureMode", strconv.FormatBool(cfg.InsecureMode))
 	}
@@ -521,6 +520,11 @@ func parseDSNParams(cfg *Config, params string) (err error) {
 			if err != nil {
 				return
 			}
+		case "requestTimeout":
+			cfg.RequestTimeout, err = parseTimeout(value)
+			if err != nil {
+				return
+			}
 		case "jwtTimeout":
 			cfg.JWTExpireTimeout, err = parseTimeout(value)
 			if err != nil {
diff --git a/gosnowflake.mak b/gosnowflake.mak
index bc48cb3..1029344 100644
--- a/gosnowflake.mak
+++ b/gosnowflake.mak
@@ -3,7 +3,7 @@ SHELL := /bin/bash
 SRC = $(shell find . -type f -name '*.go' -not -path "./vendor/*")
 
 setup:
-	@which golint &> /dev/null  || go get golang.org/x/lint/golint
+	@which golint &> /dev/null  || go get -u golang.org/x/lint/golint
 	@which make2help &> /dev/null || go get github.com/Songmu/make2help/cmd/make2help
 	@which staticcheck &> /dev/null || go get honnef.co/go/tools/cmd/staticcheck
 
diff --git a/heartbeat.go b/heartbeat.go
index 8e3a9b6..b019a10 100644
--- a/heartbeat.go
+++ b/heartbeat.go
@@ -65,7 +65,8 @@ func (hc *heartbeat) heartbeatMain() error {
 	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, hc.restful.Token)
 
 	fullURL := hc.restful.getFullURL(heartBeatPath, params)
-	resp, err := hc.restful.FuncPost(context.TODO(), hc.restful, fullURL, headers, nil, hc.restful.RequestTimeout, false)
+	timeout := hc.restful.RequestTimeout
+	resp, err := hc.restful.FuncPost(context.Background(), hc.restful, fullURL, headers, nil, timeout, false)
 	if err != nil {
 		return err
 	}
@@ -80,7 +81,7 @@ func (hc *heartbeat) heartbeatMain() error {
 			return err
 		}
 		if respd.Code == sessionExpiredCode {
-			err = hc.restful.FuncRenewSession(context.TODO(), hc.restful)
+			err = hc.restful.FuncRenewSession(context.TODO(), hc.restful, timeout)
 			if err != nil {
 				return err
 			}
diff --git a/ocsp.go b/ocsp.go
index 2bc07e1..e90ec99 100644
--- a/ocsp.go
+++ b/ocsp.go
@@ -352,6 +352,7 @@ func isTestUnknownStatus() bool {
 }
 
 func checkOCSPCacheServer(
+	ctx context.Context,
 	client clientInterface,
 	req requestFunc,
 	ocspServerHost *url.URL,
@@ -360,7 +361,7 @@ func checkOCSPCacheServer(
 	ocspS *ocspStatus) {
 	var respd map[string][]interface{}
 	headers := make(map[string]string)
-	res, err := newRetryHTTP(context.TODO(), client, req, ocspServerHost, headers, totalTimeout).execute()
+	res, err := newRetryHTTP(ctx, client, req, ocspServerHost, headers, totalTimeout).execute()
 	if err != nil {
 		glog.V(2).Infof("failed to get OCSP cache from OCSP Cache Server. %v\n", err)
 		return nil, &ocspStatus{
@@ -398,6 +399,7 @@ func checkOCSPCacheServer(
 // retryOCSP is the second level of retry method if the returned contents are corrupted. It often happens with OCSP
 // serer and retry helps.
 func retryOCSP(
+	ctx context.Context,
 	client clientInterface,
 	req requestFunc,
 	ocspHost *url.URL,
@@ -413,7 +415,7 @@ func retryOCSP(
 		multiplier = 3 // up to 3 times for Fail Close mode
 	}
 	res, err := newRetryHTTP(
-		context.TODO(), client, req, ocspHost, headers,
+		ctx, client, req, ocspHost, headers,
 		totalTimeout*time.Duration(multiplier)).doPost().setBody(reqBody).execute()
 	if err != nil {
 		return ocspRes, ocspResBytes, &ocspStatus{
@@ -452,7 +454,7 @@ func retryOCSP(
 }
 
 // getRevocationStatus checks the certificate revocation status for subject using issuer certificate.
-func getRevocationStatus(subject, issuer *x509.Certificate) *ocspStatus {
+func getRevocationStatus(ctx context.Context, subject, issuer *x509.Certificate) *ocspStatus {
 	glog.V(2).Infof("Subject: %v, Issuer: %v\n", subject.Subject, issuer.Subject)
 
 	status, ocspReq, encodedCertID := validateWithCache(subject, issuer)
@@ -510,7 +512,7 @@ func getRevocationStatus(subject, issuer *x509.Certificate) *ocspStatus {
 		Transport: snowflakeInsecureTransport,
 	}
 	ocspRes, ocspResBytes, ocspS := retryOCSP(
-		ocspClient, http.NewRequest, u, headers, ocspReq, issuer, timeout)
+		ctx, ocspClient, http.NewRequest, u, headers, ocspReq, issuer, timeout)
 	if ocspS.code != ocspSuccess {
 		return ocspS
 	}
@@ -536,7 +538,7 @@ func isValidOCSPStatus(status ocspStatusCode) bool {
 }
 
 // verifyPeerCertificate verifies all of certificate revocation status
-func verifyPeerCertificate(verifiedChains [][]*x509.Certificate) (err error) {
+func verifyPeerCertificate(ctx context.Context, verifiedChains [][]*x509.Certificate) (err error) {
 	for i := 0; i < len(verifiedChains); i++ {
 		// Certificate signed by Root CA. This should be one before the last in the Certificate Chain
 		numberOfNoneRootCerts := len(verifiedChains[i]) - 1
@@ -550,7 +552,7 @@ func verifyPeerCertificate(verifiedChains [][]*x509.Certificate) (err error) {
 			verifiedChains[i] = append(verifiedChains[i], rca)
 			numberOfNoneRootCerts++
 		}
-		results := getAllRevocationStatus(verifiedChains[i])
+		results := getAllRevocationStatus(ctx, verifiedChains[i])
 		if r := canEarlyExitForOCSP(results, numberOfNoneRootCerts); r != nil {
 			return r.err
 		}
@@ -663,7 +665,7 @@ func downloadOCSPCacheServer() {
 		Timeout:   timeout,
 		Transport: snowflakeInsecureTransport,
 	}
-	ret, ocspStatus := checkOCSPCacheServer(ocspClient, http.NewRequest, u, timeout)
+	ret, ocspStatus := checkOCSPCacheServer(context.TODO(), ocspClient, http.NewRequest, u, timeout)
 	if ocspStatus.code != ocspSuccess {
 		return
 	}
@@ -681,7 +683,7 @@ func downloadOCSPCacheServer() {
 	ocspResponseCacheLock.Unlock()
 }
 
-func getAllRevocationStatus(verifiedChains []*x509.Certificate) []*ocspStatus {
+func getAllRevocationStatus(ctx context.Context, verifiedChains []*x509.Certificate) []*ocspStatus {
 	cached := validateWithCacheForAllCertificates(verifiedChains)
 	if !cached {
 		downloadOCSPCacheServer()
@@ -689,7 +691,7 @@ func getAllRevocationStatus(verifiedChains []*x509.Certificate) []*ocspStatus {
 	n := len(verifiedChains) - 1
 	results := make([]*ocspStatus, n)
 	for j := 0; j < n; j++ {
-		results[j] = getRevocationStatus(verifiedChains[j], verifiedChains[j+1])
+		results[j] = getRevocationStatus(ctx, verifiedChains[j], verifiedChains[j+1])
 		if !isValidOCSPStatus(results[j].code) {
 			return results
 		}
@@ -700,7 +702,7 @@ func getAllRevocationStatus(verifiedChains []*x509.Certificate) []*ocspStatus {
 // verifyPeerCertificateSerial verifies the certificate revocation status in serial.
 func verifyPeerCertificateSerial(_ [][]byte, verifiedChains [][]*x509.Certificate) (err error) {
 	overrideCacheDir()
-	return verifyPeerCertificate(verifiedChains)
+	return verifyPeerCertificate(context.TODO(), verifiedChains)
 }
 
 func overrideCacheDir() {
diff --git a/restful.go b/restful.go
index ad0d06f..1f9d4e4 100644
--- a/restful.go
+++ b/restful.go
@@ -60,14 +60,14 @@ type snowflakeRestful struct {
 	FuncPostQueryHelper func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, *uuid.UUID) (*execResponse, error)
 	FuncPost            func(context.Context, *snowflakeRestful, *url.URL, map[string]string, []byte, time.Duration, bool) (*http.Response, error)
 	FuncGet             func(context.Context, *snowflakeRestful, *url.URL, map[string]string, time.Duration) (*http.Response, error)
-	FuncRenewSession    func(context.Context, *snowflakeRestful) error
-	FuncPostAuth        func(*snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration) (*authResponse, error)
-	FuncCloseSession    func(*snowflakeRestful) error
-	FuncCancelQuery     func(*snowflakeRestful, *uuid.UUID) error
-
-	FuncPostAuthSAML func(*snowflakeRestful, map[string]string, []byte, time.Duration) (*authResponse, error)
-	FuncPostAuthOKTA func(*snowflakeRestful, map[string]string, []byte, string, time.Duration) (*authOKTAResponse, error)
-	FuncGetSSO       func(*snowflakeRestful, *url.Values, map[string]string, string, time.Duration) ([]byte, error)
+	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.UUID, time.Duration) 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)
 }
 
 func (sr *snowflakeRestful) getURL() *url.URL {
@@ -152,8 +152,7 @@ func postRestfulQuery(
 		return data, err
 	}
 
-	// For context cancel/timeout cases, special cancel request need to be sent
-	err = sr.FuncCancelQuery(sr, requestID)
+	err = sr.FuncCancelQuery(context.TODO(), sr, requestID, timeout)
 	if err != nil {
 		return nil, err
 	}
@@ -192,7 +191,7 @@ func postRestfulQueryHelper(
 			return nil, err
 		}
 		if respd.Code == sessionExpiredCode {
-			err = sr.FuncRenewSession(ctx, sr)
+			err = sr.FuncRenewSession(ctx, sr, timeout)
 			if err != nil {
 				return nil, err
 			}
@@ -228,7 +227,7 @@ func postRestfulQueryHelper(
 				return nil, err
 			}
 			if respd.Code == sessionExpiredCode {
-				err = sr.FuncRenewSession(ctx, sr)
+				err = sr.FuncRenewSession(ctx, sr, timeout)
 				if err != nil {
 					return nil, err
 				}
@@ -255,7 +254,7 @@ func postRestfulQueryHelper(
 	}
 }
 
-func closeSession(sr *snowflakeRestful) error {
+func closeSession(ctx context.Context, sr *snowflakeRestful, timeout time.Duration) error {
 	glog.V(2).Info("close session")
 	params := &url.Values{}
 	params.Add("delete", "true")
@@ -269,7 +268,7 @@ func closeSession(sr *snowflakeRestful) error {
 	headers["User-Agent"] = userAgent
 	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, sr.Token)
 
-	resp, err := sr.FuncPost(context.TODO(), sr, fullURL, headers, nil, 5*time.Second, false)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, nil, 5*time.Second, false)
 	if err != nil {
 		return err
 	}
@@ -311,7 +310,7 @@ func closeSession(sr *snowflakeRestful) error {
 	}
 }
 
-func renewRestfulSession(ctx context.Context, sr *snowflakeRestful) error {
+func renewRestfulSession(ctx context.Context, sr *snowflakeRestful, timeout time.Duration) error {
 	glog.V(2).Info("start renew session")
 	params := &url.Values{}
 	params.Add(requestIDKey, uuid.New().String())
@@ -334,7 +333,7 @@ func renewRestfulSession(ctx context.Context, sr *snowflakeRestful) error {
 		return err
 	}
 
-	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, reqBody, sr.RequestTimeout, false)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, reqBody, timeout, false)
 	if err != nil {
 		return err
 	}
@@ -378,7 +377,7 @@ func renewRestfulSession(ctx context.Context, sr *snowflakeRestful) error {
 	}
 }
 
-func cancelQuery(sr *snowflakeRestful, requestID *uuid.UUID) error {
+func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID *uuid.UUID, timeout time.Duration) error {
 	glog.V(2).Info("cancel query")
 	params := &url.Values{}
 	params.Add(requestIDKey, uuid.New().String())
@@ -400,7 +399,7 @@ func cancelQuery(sr *snowflakeRestful, requestID *uuid.UUID) error {
 		return err
 	}
 
-	resp, err := sr.FuncPost(context.TODO(), sr, fullURL, headers, reqByte, sr.RequestTimeout, false)
+	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, reqByte, timeout, false)
 	if err != nil {
 		return err
 	}
@@ -414,11 +413,11 @@ func cancelQuery(sr *snowflakeRestful, requestID *uuid.UUID) error {
 			return err
 		}
 		if !respd.Success && respd.Code == sessionExpiredCode {
-			err := sr.FuncRenewSession(context.TODO(), sr)
+			err := sr.FuncRenewSession(ctx, sr, timeout)
 			if err != nil {
 				return err
 			}
-			return sr.FuncCancelQuery(sr, requestID)
+			return sr.FuncCancelQuery(ctx, sr, requestID, timeout)
 		} else if respd.Success {
 			return nil
 		} else {
diff --git a/result.go b/result.go
index e15475f..6e271fc 100644
--- a/result.go
+++ b/result.go
@@ -2,9 +2,15 @@
 
 package gosnowflake
 
+// SnowflakeResult provides the associated query ID
+type SnowflakeResult interface {
+	QueryID() string
+}
+
 type snowflakeResult struct {
 	affectedRows int64
 	insertID     int64 // Snowflake doesn't support last insert id
+	queryID      string
 }
 
 func (res *snowflakeResult) LastInsertId() (int64, error) {
@@ -14,3 +20,7 @@ func (res *snowflakeResult) LastInsertId() (int64, error) {
 func (res *snowflakeResult) RowsAffected() (int64, error) {
 	return res.affectedRows, nil
 }
+
+func (res *snowflakeResult) QueryID() string {
+	return res.queryID
+}
diff --git a/retry.go b/retry.go
index c603a45..d8ca92b 100644
--- a/retry.go
+++ b/retry.go
@@ -5,7 +5,6 @@ package gosnowflake
 import (
 	"bytes"
 	"crypto/x509"
-	"errors"
 	"fmt"
 	"github.com/google/uuid"
 	"io"
@@ -264,9 +263,9 @@ func (r *retryHTTP) execute() (res *http.Response, err error) {
 					return nil, err
 				}
 				if res != nil {
-					return nil, fmt.Errorf("timeout. HTTP Status: %v. Hanging?", res.StatusCode)
+					return nil, fmt.Errorf("timeout after %s. HTTP Status: %v. Hanging?", r.timeout, res.StatusCode)
 				}
-				return nil, errors.New("timeout. Hanging?")
+				return nil, fmt.Errorf("timeout after %s. Hanging?", r.timeout)
 			}
 		}
 		retryCounter++
diff --git a/rows.go b/rows.go
index 588d738..b95bfa3 100644
--- a/rows.go
+++ b/rows.go
@@ -42,6 +42,7 @@ type snowflakeRows struct {
 	sc              *snowflakeConn
 	RowType         []execResponseRowType
 	ChunkDownloader *snowflakeChunkDownloader
+	queryID         string
 }
 
 func (rows *snowflakeRows) Close() (err error) {
@@ -88,8 +89,8 @@ type snowflakeChunkDownloader struct {
 	RowSet             rowSetType
 	ChunkHeader        map[string]string
 	CurrentIndex       int
-	FuncDownload       func(*snowflakeChunkDownloader, int)
-	FuncDownloadHelper func(context.Context, *snowflakeChunkDownloader, int)
+	FuncDownload       func(context.Context, *snowflakeChunkDownloader, int)
+	FuncDownloadHelper func(context.Context, *snowflakeChunkDownloader, int) error
 	FuncGet            func(context.Context, *snowflakeChunkDownloader, string, map[string]string, time.Duration) (*http.Response, error)
 	DoneDownloadCond   *sync.Cond
 	NextDownloader     *snowflakeChunkDownloader
@@ -147,6 +148,10 @@ func (rows *snowflakeRows) ColumnTypeScanType(index int) reflect.Type {
 	return snowflakeTypeToGo(rows.RowType[index].Type, rows.RowType[index].Scale)
 }
 
+func (rows *snowflakeRows) QueryID() string {
+	return rows.queryID
+}
+
 func (rows *snowflakeRows) Next(dest []driver.Value) (err error) {
 	row, err := rows.ChunkDownloader.Next()
 	if err != nil {
@@ -256,7 +261,7 @@ func (scd *snowflakeChunkDownloader) schedule() {
 	select {
 	case nextIdx := <-scd.ChunksChan:
 		glog.V(2).Infof("schedule chunk: %v", nextIdx+1)
-		go scd.FuncDownload(scd, nextIdx)
+		go scd.FuncDownload(scd.ctx, scd, nextIdx)
 	default:
 		// no more download
 		glog.V(2).Info("no more download")
@@ -268,7 +273,7 @@ func (scd *snowflakeChunkDownloader) checkErrorRetry() (err error) {
 	case errc := <-scd.ChunksError:
 		if scd.ChunksErrorCounter < maxChunkDownloaderErrorCounter && errc.Error != context.Canceled {
 			// add the index to the chunks channel so that the download will be retried.
-			go scd.FuncDownload(scd, errc.Index)
+			go scd.FuncDownload(scd.ctx, scd, errc.Index)
 			scd.ChunksErrorCounter++
 			glog.V(2).Infof("chunk idx: %v, err: %v. retrying (%v/%v)...",
 				errc.Index, errc.Error, scd.ChunksErrorCounter, maxChunkDownloaderErrorCounter)
@@ -377,16 +382,21 @@ func (r *largeResultSetReader) Read(p []byte) (n int, err error) {
 	return 0, io.EOF
 }
 
-func downloadChunk(scd *snowflakeChunkDownloader, idx int) {
+func downloadChunk(ctx context.Context, scd *snowflakeChunkDownloader, idx int) {
 	glog.V(2).Infof("download start chunk: %v", idx+1)
+	defer scd.DoneDownloadCond.Broadcast()
 
-	scd.FuncDownloadHelper(scd.ctx, scd, idx)
-	if scd.ctx.Err() == context.Canceled || scd.ctx.Err() == context.DeadlineExceeded {
+	if err := scd.FuncDownloadHelper(ctx, scd, idx); err != nil {
+		glog.V(1).Infof(
+			"failed to extract HTTP response body. URL: %v, err: %v", scd.ChunkMetas[idx].URL, err)
+		glog.Flush()
+		scd.ChunksError <- &chunkError{Index: idx, Error: err}
+	} else if scd.ctx.Err() == context.Canceled || scd.ctx.Err() == context.DeadlineExceeded {
 		scd.ChunksError <- &chunkError{Index: idx, Error: scd.ctx.Err()}
 	}
 }
 
-func downloadChunkHelper(ctx context.Context, scd *snowflakeChunkDownloader, idx int) {
+func downloadChunkHelper(ctx context.Context, scd *snowflakeChunkDownloader, idx int) error {
 	headers := make(map[string]string)
 	if len(scd.ChunkHeader) > 0 {
 		glog.V(2).Info("chunk header is provided.")
@@ -400,116 +410,98 @@ func downloadChunkHelper(ctx context.Context, scd *snowflakeChunkDownloader, idx
 
 	resp, err := scd.FuncGet(ctx, scd, scd.ChunkMetas[idx].URL, headers, scd.sc.rest.RequestTimeout)
 	if err != nil {
-		raiseDownloadError(scd, idx, err)
-		return
+		return err
 	}
 	bufStream := bufio.NewReader(resp.Body)
 	defer resp.Body.Close()
 	glog.V(2).Infof("response returned chunk: %v, resp: %v", idx+1, resp)
-	if resp.StatusCode == http.StatusOK {
-		gzipMagic, err := bufStream.Peek(2)
+	if resp.StatusCode != http.StatusOK {
+		b, err := ioutil.ReadAll(bufStream)
 		if err != nil {
-			raiseDownloadError(scd, idx, err)
-			return
+			return err
 		}
-		start := time.Now()
-		var source io.Reader
-		if gzipMagic[0] == 0x1f && gzipMagic[1] == 0x8b {
-			// detects and uncompresses Gzip format data
-			bufStream0, err := gzip.NewReader(bufStream)
-			if err != nil {
-				raiseDownloadError(scd, idx, err)
-				return
-			}
-			defer bufStream0.Close()
-			source = bufStream0
-		} else {
-			source = bufStream
+		glog.V(1).Infof("HTTP: %v, URL: %v, Body: %v", resp.StatusCode, scd.ChunkMetas[idx].URL, b)
+		glog.V(1).Infof("Header: %v", resp.Header)
+		glog.Flush()
+		return &SnowflakeError{
+			Number:      ErrFailedToGetChunk,
+			SQLState:    SQLStateConnectionFailure,
+			Message:     errMsgFailedToGetChunk,
+			MessageArgs: []interface{}{idx},
 		}
-		st := &largeResultSetReader{
-			status: 0,
-			body:   source,
+	}
+	return decodeChunk(scd, idx, bufStream)
+}
+
+func decodeChunk(scd *snowflakeChunkDownloader, idx int, bufStream *bufio.Reader) (err error) {
+	gzipMagic, err := bufStream.Peek(2)
+	if err != nil {
+		return err
+	}
+	start := time.Now()
+	var source io.Reader
+	if gzipMagic[0] == 0x1f && gzipMagic[1] == 0x8b {
+		// detects and uncompresses Gzip format data
+		bufStream0, err := gzip.NewReader(bufStream)
+		if err != nil {
+			return err
 		}
-		var respd []chunkRowType
-		if scd.QueryResultFormat != arrowFormat {
-			var decRespd [][]*string
-			if !CustomJSONDecoderEnabled {
-				dec := json.NewDecoder(st)
-				for {
-					if err := dec.Decode(&decRespd); err == io.EOF {
-						break
-					} else if err != nil {
-						raiseDownloadError(scd, idx, err)
-						return
-					}
-				}
-			} else {
-				decRespd, err = decodeLargeChunk(st, scd.ChunkMetas[idx].RowCount, scd.CellCount)
-				if err != nil {
-					raiseDownloadError(scd, idx, err)
-					return
+		defer bufStream0.Close()
+		source = bufStream0
+	} else {
+		source = bufStream
+	}
+	st := &largeResultSetReader{
+		status: 0,
+		body:   source,
+	}
+	var respd []chunkRowType
+	if scd.QueryResultFormat != arrowFormat {
+		var decRespd [][]*string
+		if !CustomJSONDecoderEnabled {
+			dec := json.NewDecoder(st)
+			for {
+				if err := dec.Decode(&decRespd); err == io.EOF {
+					break
+				} else if err != nil {
+					return err
 				}
 			}
-			respd = make([]chunkRowType, len(decRespd))
-			populateJSONRowSet(respd, decRespd)
 		} else {
-			ipcReader, err := ipc.NewReader(source)
-			if err != nil {
-				return
-			}
-			arc := arrowResultChunk{
-				*ipcReader,
-				0,
-				int(scd.totalUncompressedSize()),
-				memory.NewGoAllocator(),
-			}
-			respd, err = arc.decodeArrowChunk(scd.RowSet.RowType)
+			decRespd, err = decodeLargeChunk(st, scd.ChunkMetas[idx].RowCount, scd.CellCount)
 			if err != nil {
-				return
+				return err
 			}
 		}
-		glog.V(2).Infof(
-			"decoded %d rows w/ %d bytes in %s (chunk %v)",
-			scd.ChunkMetas[idx].RowCount,
-			scd.ChunkMetas[idx].UncompressedSize,
-			time.Since(start), idx+1,
-		)
-
-		scd.ChunksMutex.Lock()
-		defer scd.ChunksMutex.Unlock()
-		scd.Chunks[idx] = respd
-		scd.DoneDownloadCond.Broadcast()
+		respd = make([]chunkRowType, len(decRespd))
+		populateJSONRowSet(respd, decRespd)
 	} else {
-		b, err := ioutil.ReadAll(bufStream)
+		ipcReader, err := ipc.NewReader(source)
 		if err != nil {
-			raiseDownloadError(scd, idx, err)
-			return
+			return err
+		}
+		arc := arrowResultChunk{
+			*ipcReader,
+			0,
+			int(scd.totalUncompressedSize()),
+			memory.NewGoAllocator(),
+		}
+		respd, err = arc.decodeArrowChunk(scd.RowSet.RowType)
+		if err != nil {
+			return err
 		}
-		glog.V(1).Infof("HTTP: %v, URL: %v, Body: %v", resp.StatusCode, scd.ChunkMetas[idx].URL, b)
-		glog.V(1).Infof("Header: %v", resp.Header)
-		glog.Flush()
-		scd.ChunksMutex.Lock()
-		defer scd.ChunksMutex.Unlock()
-		scd.ChunksError <- &chunkError{
-			Index: idx,
-			Error: &SnowflakeError{
-				Number:      ErrFailedToGetChunk,
-				SQLState:    SQLStateConnectionFailure,
-				Message:     errMsgFailedToGetChunk,
-				MessageArgs: []interface{}{idx},
-			}}
-		scd.DoneDownloadCond.Broadcast()
 	}
-}
+	glog.V(2).Infof(
+		"decoded %d rows w/ %d bytes in %s (chunk %v)",
+		scd.ChunkMetas[idx].RowCount,
+		scd.ChunkMetas[idx].UncompressedSize,
+		time.Since(start), idx+1,
+	)
 
-func raiseDownloadError(scd *snowflakeChunkDownloader, idx int, err error) {
-	glog.V(1).Infof(
-		"failed to extract HTTP response body. URL: %v, err: %v", scd.ChunkMetas[idx].URL, err)
-	glog.Flush()
 	scd.ChunksMutex.Lock()
 	defer scd.ChunksMutex.Unlock()
-	scd.ChunksError <- &chunkError{Index: idx, Error: err}
-	scd.DoneDownloadCond.Broadcast()
+	scd.Chunks[idx] = respd
+	return nil
 }
 
 func populateJSONRowSet(dst []chunkRowType, src [][]*string) {


Test patch:
diff --git a/auth_test.go b/auth_test.go
index 75589ed..372985e 100644
--- a/auth_test.go
+++ b/auth_test.go
@@ -3,6 +3,7 @@
 package gosnowflake
 
 import (
+	"context"
 	"crypto/rand"
 	"crypto/rsa"
 	"encoding/json"
@@ -20,51 +21,51 @@ func TestUnitPostAuth(t *testing.T) {
 		FuncPost: postTestAfterRenew,
 	}
 	var err error
-	_, err = postAuth(sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
+	_, err = postAuth(context.TODO(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
 	sr.FuncPost = postTestError
-	_, err = postAuth(sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
+	_, err = postAuth(context.TODO(), sr, &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(sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
+	_, err = postAuth(context.TODO(), sr, &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(sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
+	_, err = postAuth(context.TODO(), sr, &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(sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0)
+	_, err = postAuth(context.TODO(), sr, &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(_ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthFailServiceIssue(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return nil, &SnowflakeError{
 		Number: ErrCodeServiceUnavailable,
 	}
 }
 
-func postAuthFailWrongAccount(_ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthFailWrongAccount(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return nil, &SnowflakeError{
 		Number: ErrCodeFailedToConnect,
 	}
 }
 
-func postAuthFailUnknown(_ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthFailUnknown(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return nil, &SnowflakeError{
 		Number: ErrFailedToAuth,
 	}
 }
 
-func postAuthSuccessWithErrorCode(_ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthSuccessWithErrorCode(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return &authResponse{
 		Success: false,
 		Code:    "98765",
@@ -72,7 +73,7 @@ func postAuthSuccessWithErrorCode(_ *snowflakeRestful, _ *url.Values, _ map[stri
 	}, nil
 }
 
-func postAuthSuccessWithInvalidErrorCode(_ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthSuccessWithInvalidErrorCode(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return &authResponse{
 		Success: false,
 		Code:    "abcdef",
@@ -80,7 +81,7 @@ func postAuthSuccessWithInvalidErrorCode(_ *snowflakeRestful, _ *url.Values, _ m
 	}, nil
 }
 
-func postAuthSuccess(_ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthSuccess(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return &authResponse{
 		Success: true,
 		Data: authResponseMain{
@@ -93,7 +94,7 @@ func postAuthSuccess(_ *snowflakeRestful, _ *url.Values, _ map[string]string, _
 	}, nil
 }
 
-func postAuthCheckSAMLResponse(_ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckSAMLResponse(_ context.Context, _ *snowflakeRestful, _ *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
@@ -116,6 +117,7 @@ func postAuthCheckSAMLResponse(_ *snowflakeRestful, _ *url.Values, _ map[string]
 // Checks that the request body generated when authenticating with OAuth
 // contains all the necessary values.
 func postAuthCheckOAuth(
+	_ context.Context,
 	_ *snowflakeRestful,
 	_ *url.Values, _ map[string]string,
 	jsonBody []byte,
@@ -145,7 +147,7 @@ func postAuthCheckOAuth(
 	}, nil
 }
 
-func postAuthCheckPasscode(_ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckPasscode(_ context.Context, _ *snowflakeRestful, _ *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
@@ -165,7 +167,7 @@ func postAuthCheckPasscode(_ *snowflakeRestful, _ *url.Values, _ map[string]stri
 	}, nil
 }
 
-func postAuthCheckPasscodeInPassword(_ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckPasscodeInPassword(_ context.Context, _ *snowflakeRestful, _ *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
@@ -187,7 +189,7 @@ func postAuthCheckPasscodeInPassword(_ *snowflakeRestful, _ *url.Values, _ map[s
 
 // JWT token validate callback function to check the JWT token
 // It uses the public key paired with the testPrivKey
-func postAuthCheckJWTToken(_ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+func postAuthCheckJWTToken(_ context.Context, _ *snowflakeRestful, _ *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
@@ -257,7 +259,7 @@ func TestUnitAuthenticate(t *testing.T) {
 	}
 	sc.rest = sr
 
-	_, err = authenticate(sc, []byte{}, []byte{})
+	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
@@ -266,7 +268,7 @@ func TestUnitAuthenticate(t *testing.T) {
 		t.Fatalf("Snowflake error is expected. err: %v", driverErr)
 	}
 	sr.FuncPostAuth = postAuthFailWrongAccount
-	_, err = authenticate(sc, []byte{}, []byte{})
+	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
@@ -275,7 +277,7 @@ func TestUnitAuthenticate(t *testing.T) {
 		t.Fatalf("Snowflake error is expected. err: %v", driverErr)
 	}
 	sr.FuncPostAuth = postAuthFailUnknown
-	_, err = authenticate(sc, []byte{}, []byte{})
+	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
@@ -284,7 +286,7 @@ func TestUnitAuthenticate(t *testing.T) {
 		t.Fatalf("Snowflake error is expected. err: %v", driverErr)
 	}
 	sr.FuncPostAuth = postAuthSuccessWithErrorCode
-	_, err = authenticate(sc, []byte{}, []byte{})
+	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
@@ -293,13 +295,13 @@ func TestUnitAuthenticate(t *testing.T) {
 		t.Fatalf("Snowflake error is expected. err: %v", driverErr)
 	}
 	sr.FuncPostAuth = postAuthSuccessWithInvalidErrorCode
-	_, err = authenticate(sc, []byte{}, []byte{})
+	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
 	sr.FuncPostAuth = postAuthSuccess
 	var resp *authResponseMain
-	resp, err = authenticate(sc, []byte{}, []byte{})
+	resp, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err != nil {
 		t.Fatalf("failed to auth. err: %v", err)
 	}
@@ -320,7 +322,7 @@ func TestUnitAuthenticateSaml(t *testing.T) {
 		Host:   "blah.okta.com",
 	}
 	sc.rest = sr
-	_, err = authenticate(sc, []byte("HTML data in bytes from"), []byte{})
+	_, err = authenticate(context.TODO(), sc, []byte("HTML data in bytes from"), []byte{})
 	if err != nil {
 		t.Fatalf("failed to run. err: %v", err)
 	}
@@ -336,7 +338,7 @@ func TestUnitAuthenticateOAuth(t *testing.T) {
 	sc.cfg.Token = "oauthToken"
 	sc.cfg.Authenticator = AuthTypeOAuth
 	sc.rest = sr
-	_, err = authenticate(sc, []byte{}, []byte{})
+	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err != nil {
 		t.Fatalf("failed to run. err: %v", err)
 	}
@@ -351,14 +353,14 @@ func TestUnitAuthenticatePasscode(t *testing.T) {
 	sc.cfg.Passcode = "987654321"
 	sc.rest = sr
 
-	_, err = authenticate(sc, []byte{}, []byte{})
+	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err != nil {
 		t.Fatalf("failed to run. err: %v", err)
 	}
 	sr.FuncPostAuth = postAuthCheckPasscodeInPassword
 	sc.rest = sr
 	sc.cfg.PasscodeInPassword = true
-	_, err = authenticate(sc, []byte{}, []byte{})
+	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err != nil {
 		t.Fatalf("failed to run. err: %v", err)
 	}
@@ -378,7 +380,7 @@ func TestUnitAuthenticateJWT(t *testing.T) {
 	sc.rest = sr
 
 	// A valid JWT token should pass
-	_, err = authenticate(sc, []byte{}, []byte{})
+	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err != nil {
 		t.Fatalf("failed to run. err: %v", err)
 	}
@@ -386,7 +388,7 @@ func TestUnitAuthenticateJWT(t *testing.T) {
 	// An invalid JWT token should not pass
 	invalidPrivateKey, _ := rsa.GenerateKey(rand.Reader, 2048)
 	sc.cfg.PrivateKey = invalidPrivateKey
-	_, err = authenticate(sc, []byte{}, []byte{})
+	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err == nil {
 		t.Fatalf("invalid token passed")
 	}
diff --git a/authokta_test.go b/authokta_test.go
index 8a1688c..6586d40 100644
--- a/authokta_test.go
+++ b/authokta_test.go
@@ -63,17 +63,17 @@ func TestUnitPostAuthSAML(t *testing.T) {
 		FuncPost: postTestError,
 	}
 	var err error
-	_, err = postAuthSAML(sr, make(map[string]string), []byte{}, 0)
+	_, err = postAuthSAML(context.TODO(), sr, make(map[string]string), []byte{}, 0)
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
 	sr.FuncPost = postTestAppBadGatewayError
-	_, err = postAuthSAML(sr, make(map[string]string), []byte{}, 0)
+	_, err = postAuthSAML(context.TODO(), sr, make(map[string]string), []byte{}, 0)
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
 	sr.FuncPost = postTestSuccessButInvalidJSON
-	_, err = postAuthSAML(sr, make(map[string]string), []byte{0x12, 0x34}, 0)
+	_, err = postAuthSAML(context.TODO(), sr, make(map[string]string), []byte{0x12, 0x34}, 0)
 	if err == nil {
 		t.Fatalf("should have failed to post")
 	}
@@ -84,17 +84,17 @@ func TestUnitPostAuthOKTA(t *testing.T) {
 		FuncPost: postTestError,
 	}
 	var err error
-	_, err = postAuthOKTA(sr, make(map[string]string), []byte{}, "hahah", 0)
+	_, err = postAuthOKTA(context.TODO(), sr, make(map[string]string), []byte{}, "hahah", 0)
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
 	sr.FuncPost = postTestAppBadGatewayError
-	_, err = postAuthOKTA(sr, make(map[string]string), []byte{}, "hahah", 0)
+	_, err = postAuthOKTA(context.TODO(), sr, make(map[string]string), []byte{}, "hahah", 0)
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
 	sr.FuncPost = postTestSuccessButInvalidJSON
-	_, err = postAuthOKTA(sr, make(map[string]string), []byte{0x12, 0x34}, "haha", 0)
+	_, err = postAuthOKTA(context.TODO(), sr, make(map[string]string), []byte{0x12, 0x34}, "haha", 0)
 	if err == nil {
 		t.Fatal("should have failed to run post request after the renewal")
 	}
@@ -105,34 +105,34 @@ func TestUnitGetSSO(t *testing.T) {
 		FuncGet: getTestError,
 	}
 	var err error
-	_, err = getSSO(sr, &url.Values{}, make(map[string]string), "hahah", 0)
+	_, err = getSSO(context.TODO(), sr, &url.Values{}, make(map[string]string), "hahah", 0)
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
 	sr.FuncGet = getTestAppBadGatewayError
-	_, err = getSSO(sr, &url.Values{}, make(map[string]string), "hahah", 0)
+	_, err = getSSO(context.TODO(), sr, &url.Values{}, make(map[string]string), "hahah", 0)
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
 	sr.FuncGet = getTestHTMLSuccess
-	_, err = getSSO(sr, &url.Values{}, make(map[string]string), "hahah", 0)
+	_, err = getSSO(context.TODO(), sr, &url.Values{}, make(map[string]string), "hahah", 0)
 	if err != nil {
 		t.Fatalf("failed to get HTML content. err: %v", err)
 	}
 }
 
-func postAuthSAMLError(_ *snowflakeRestful, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthSAMLError(_ context.Context, _ *snowflakeRestful, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return &authResponse{}, errors.New("failed to get SAML response")
 }
 
-func postAuthSAMLAuthFail(_ *snowflakeRestful, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthSAMLAuthFail(_ context.Context, _ *snowflakeRestful, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return &authResponse{
 		Success: false,
 		Message: "SAML auth failed",
 	}, nil
 }
 
-func postAuthSAMLAuthSuccessButInvalidURL(_ *snowflakeRestful, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthSAMLAuthSuccessButInvalidURL(_ context.Context, _ *snowflakeRestful, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return &authResponse{
 		Success: true,
 		Message: "",
@@ -143,7 +143,7 @@ func postAuthSAMLAuthSuccessButInvalidURL(_ *snowflakeRestful, _ map[string]stri
 	}, nil
 }
 
-func postAuthSAMLAuthSuccess(_ *snowflakeRestful, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
+func postAuthSAMLAuthSuccess(_ context.Context, _ *snowflakeRestful, _ map[string]string, _ []byte, _ time.Duration) (*authResponse, error) {
 	return &authResponse{
 		Success: true,
 		Message: "",
@@ -154,23 +154,23 @@ func postAuthSAMLAuthSuccess(_ *snowflakeRestful, _ map[string]string, _ []byte,
 	}, nil
 }
 
-func postAuthOKTAError(_ *snowflakeRestful, _ map[string]string, _ []byte, _ string, _ time.Duration) (*authOKTAResponse, error) {
+func postAuthOKTAError(_ context.Context, _ *snowflakeRestful, _ map[string]string, _ []byte, _ string, _ time.Duration) (*authOKTAResponse, error) {
 	return &authOKTAResponse{}, errors.New("failed to get SAML response")
 }
 
-func postAuthOKTASuccess(_ *snowflakeRestful, _ map[string]string, _ []byte, _ string, _ time.Duration) (*authOKTAResponse, error) {
+func postAuthOKTASuccess(_ context.Context, _ *snowflakeRestful, _ map[string]string, _ []byte, _ string, _ time.Duration) (*authOKTAResponse, error) {
 	return &authOKTAResponse{}, nil
 }
 
-func getSSOError(_ *snowflakeRestful, _ *url.Values, _ map[string]string, _ string, _ time.Duration) ([]byte, error) {
+func getSSOError(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ string, _ time.Duration) ([]byte, error) {
 	return []byte{}, errors.New("failed to get SSO html")
 }
 
-func getSSOSuccessButInvalidURL(_ *snowflakeRestful, _ *url.Values, _ map[string]string, _ string, _ time.Duration) ([]byte, error) {
+func getSSOSuccessButInvalidURL(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ string, _ time.Duration) ([]byte, error) {
 	return []byte(`<html><form id="1"/></html>`), nil
 }
 
-func getSSOSuccess(_ *snowflakeRestful, _ *url.Values, _ map[string]string, _ string, _ time.Duration) ([]byte, error) {
+func getSSOSuccess(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ string, _ time.Duration) ([]byte, error) {
 	return []byte(`<html><form id="1" action="https&#x3a;&#x2f;&#x2f;abc.com&#x2f;"></form></html>`), nil
 }
 
@@ -190,17 +190,17 @@ func TestUnitAuthenticateBySAML(t *testing.T) {
 		FuncPostAuthSAML: postAuthSAMLError,
 	}
 	var err error
-	_, err = authenticateBySAML(sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.TODO(), sr, authenticator, application, account, user, password)
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
 	sr.FuncPostAuthSAML = postAuthSAMLAuthFail
-	_, err = authenticateBySAML(sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.TODO(), sr, authenticator, application, account, user, password)
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
 	sr.FuncPostAuthSAML = postAuthSAMLAuthSuccessButInvalidURL
-	_, err = authenticateBySAML(sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.TODO(), sr, authenticator, application, account, user, password)
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
@@ -213,23 +213,23 @@ func TestUnitAuthenticateBySAML(t *testing.T) {
 	}
 	sr.FuncPostAuthSAML = postAuthSAMLAuthSuccess
 	sr.FuncPostAuthOKTA = postAuthOKTAError
-	_, err = authenticateBySAML(sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.TODO(), sr, authenticator, application, account, user, password)
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
 	sr.FuncPostAuthOKTA = postAuthOKTASuccess
 	sr.FuncGetSSO = getSSOError
-	_, err = authenticateBySAML(sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.TODO(), sr, authenticator, application, account, user, password)
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
 	sr.FuncGetSSO = getSSOSuccessButInvalidURL
-	_, err = authenticateBySAML(sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.TODO(), sr, authenticator, application, account, user, password)
 	if err == nil {
 		t.Fatal("should have failed.")
 	}
 	sr.FuncGetSSO = getSSOSuccess
-	_, err = authenticateBySAML(sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.TODO(), sr, authenticator, application, account, user, password)
 	if err != nil {
 		t.Fatalf("failed. err: %v", err)
 	}
diff --git a/connection_test.go b/connection_test.go
index c5e760b..1fb3014 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -62,7 +62,7 @@ func TestServiceName(t *testing.T) {
 	}
 }
 
-func closeSessionMock(_ *snowflakeRestful) error {
+func closeSessionMock(_ context.Context, _ *snowflakeRestful, _ time.Duration) error {
 	return &SnowflakeError{
 		Number: ErrSessionGone,
 	}
diff --git a/ocsp_test.go b/ocsp_test.go
index 8c547b4..975a86f 100644
--- a/ocsp_test.go
+++ b/ocsp_test.go
@@ -4,6 +4,7 @@ package gosnowflake
 
 import (
 	"bytes"
+	"context"
 	"crypto"
 	"crypto/tls"
 	"crypto/x509"
@@ -312,6 +313,7 @@ func TestOCSPRetry(t *testing.T) {
 		body:    []byte{1, 2, 3},
 	}
 	res, b, st := retryOCSP(
+		context.TODO(),
 		client, fakeRequestFunc,
 		dummyOCSPHost,
 		make(map[string]string), []byte{0}, certs[len(certs)-1], 10*time.Second)
@@ -324,6 +326,7 @@ func TestOCSPRetry(t *testing.T) {
 		body:    []byte{1, 2, 3},
 	}
 	res, b, st = retryOCSP(
+		context.TODO(),
 		client, fakeRequestFunc,
 		dummyOCSPHost,
 		make(map[string]string), []byte{0}, certs[len(certs)-1], 5*time.Second)
@@ -343,7 +346,7 @@ func TestOCSPCacheServerRetry(t *testing.T) {
 		body:    []byte{1, 2, 3},
 	}
 	res, st := checkOCSPCacheServer(
-		client, fakeRequestFunc, dummyOCSPHost, 20*time.Second)
+		context.TODO(), client, fakeRequestFunc, dummyOCSPHost, 20*time.Second)
 	if st.err == nil {
 		t.Errorf("should fail: %v", res)
 	}
@@ -353,7 +356,7 @@ func TestOCSPCacheServerRetry(t *testing.T) {
 		body:    []byte{1, 2, 3},
 	}
 	res, st = checkOCSPCacheServer(
-		client, fakeRequestFunc, dummyOCSPHost, 10*time.Second)
+		context.TODO(), client, fakeRequestFunc, dummyOCSPHost, 10*time.Second)
 	if st.err == nil {
 		t.Errorf("should fail: %v", res)
 	}
diff --git a/restful_test.go b/restful_test.go
index c8e77bf..0cb8b81 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -113,11 +113,11 @@ func TestUnitPostQueryHelperError(t *testing.T) {
 	}
 }
 
-func renewSessionTest(_ context.Context, _ *snowflakeRestful) error {
+func renewSessionTest(_ context.Context, _ *snowflakeRestful, _ time.Duration) error {
 	return nil
 }
 
-func renewSessionTestError(_ context.Context, _ *snowflakeRestful) error {
+func renewSessionTestError(_ context.Context, _ *snowflakeRestful, _ time.Duration) error {
 	return errors.New("failed to renew session in tests")
 }
 
@@ -161,22 +161,22 @@ func TestUnitRenewRestfulSession(t *testing.T) {
 		Token:       "token",
 		FuncPost:    postTestAfterRenew,
 	}
-	err := renewRestfulSession(context.Background(), sr)
+	err := renewRestfulSession(context.Background(), sr, time.Second)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
 	sr.FuncPost = postTestError
-	err = renewRestfulSession(context.Background(), sr)
+	err = renewRestfulSession(context.Background(), sr, time.Second)
 	if err == nil {
 		t.Fatal("should have failed to run post request after the renewal")
 	}
 	sr.FuncPost = postTestAppBadGatewayError
-	err = renewRestfulSession(context.Background(), sr)
+	err = renewRestfulSession(context.Background(), sr, time.Second)
 	if err == nil {
 		t.Fatal("should have failed to run post request after the renewal")
 	}
 	sr.FuncPost = postTestSuccessButInvalidJSON
-	err = renewRestfulSession(context.Background(), sr)
+	err = renewRestfulSession(context.Background(), sr, time.Second)
 	if err == nil {
 		t.Fatal("should have failed to run post request after the renewal")
 	}
@@ -186,22 +186,22 @@ func TestUnitCloseSession(t *testing.T) {
 	sr := &snowflakeRestful{
 		FuncPost: postTestAfterRenew,
 	}
-	err := closeSession(sr)
+	err := closeSession(context.Background(), sr, time.Second)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
 	sr.FuncPost = postTestError
-	err = closeSession(sr)
+	err = closeSession(context.Background(), sr, time.Second)
 	if err == nil {
 		t.Fatal("should have failed to close session")
 	}
 	sr.FuncPost = postTestAppBadGatewayError
-	err = closeSession(sr)
+	err = closeSession(context.Background(), sr, time.Second)
 	if err == nil {
 		t.Fatal("should have failed to close session")
 	}
 	sr.FuncPost = postTestSuccessButInvalidJSON
-	err = closeSession(sr)
+	err = closeSession(context.Background(), sr, time.Second)
 	if err == nil {
 		t.Fatal("should have failed to close session")
 	}
@@ -213,25 +213,25 @@ func TestUnitCancelQuery(t *testing.T) {
 	}
 	var requestID uuid.UUID
 	requestID = uuid.New()
-	err := cancelQuery(sr, &requestID)
+	err := cancelQuery(context.Background(), sr, &requestID, time.Second)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
 	sr.FuncPost = postTestError
 	requestID = uuid.New()
-	err = cancelQuery(sr, &requestID)
+	err = cancelQuery(context.Background(), sr, &requestID, time.Second)
 	if err == nil {
 		t.Fatal("should have failed to close session")
 	}
 	sr.FuncPost = postTestAppBadGatewayError
 	requestID = uuid.New()
-	err = cancelQuery(sr, &requestID)
+	err = cancelQuery(context.Background(), sr, &requestID, time.Second)
 	if err == nil {
 		t.Fatal("should have failed to close session")
 	}
 	sr.FuncPost = postTestSuccessButInvalidJSON
 	requestID = uuid.New()
-	err = cancelQuery(sr, &requestID)
+	err = cancelQuery(context.Background(), sr, &requestID, time.Second)
 	if err == nil {
 		t.Fatal("should have failed to close session")
 	}
diff --git a/rows_test.go b/rows_test.go
index 52d81ea..403507c 100644
--- a/rows_test.go
+++ b/rows_test.go
@@ -68,7 +68,7 @@ func TestRowsWithoutChunkDownloader(t *testing.T) {
 
 }
 
-func downloadChunkTest(scd *snowflakeChunkDownloader, idx int) {
+func downloadChunkTest(ctx context.Context, scd *snowflakeChunkDownloader, idx int) {
 	d := make([][]*string, 0)
 	for i := 0; i < rowsInChunk; i++ {
 		v1 := fmt.Sprintf("%v", idx*1000+i)
@@ -139,7 +139,7 @@ func TestRowsWithChunkDownloader(t *testing.T) {
 	glog.V(2).Info("END TESTS")
 }
 
-func downloadChunkTestError(scd *snowflakeChunkDownloader, idx int) {
+func downloadChunkTestError(ctx context.Context, scd *snowflakeChunkDownloader, idx int) {
 	// fail to download 6th and 10th chunk, and retry up to N times and success
 	// NOTE: zero based index
 	scd.ChunksMutex.Lock()
@@ -220,7 +220,7 @@ func TestRowsWithChunkDownloaderError(t *testing.T) {
 	glog.V(2).Info("END TESTS")
 }
 
-func downloadChunkTestErrorFail(scd *snowflakeChunkDownloader, idx int) {
+func downloadChunkTestErrorFail(ctx context.Context, scd *snowflakeChunkDownloader, idx int) {
 	// fail to download 6th and 10th chunk, and retry up to N times and fail
 	// NOTE: zero based index
 	scd.ChunksMutex.Lock()
@@ -326,7 +326,7 @@ func TestDownloadChunkInvalidResponseBody(t *testing.T) {
 	scd.DoneDownloadCond = sync.NewCond(scd.ChunksMutex)
 	scd.Chunks = make(map[int][]chunkRowType)
 	scd.ChunksError = make(chan *chunkError, 1)
-	scd.FuncDownload(scd, 1)
+	scd.FuncDownload(scd.ctx, scd, 1)
 	select {
 	case errc := <-scd.ChunksError:
 		if errc.Index != 1 {
@@ -368,7 +368,7 @@ func TestDownloadChunkErrorStatus(t *testing.T) {
 	scd.DoneDownloadCond = sync.NewCond(scd.ChunksMutex)
 	scd.Chunks = make(map[int][]chunkRowType)
 	scd.ChunksError = make(chan *chunkError, 1)
-	scd.FuncDownload(scd, 1)
+	scd.FuncDownload(scd.ctx, scd, 1)
 	select {
 	case errc := <-scd.ChunksError:
 		if errc.Index != 1 {
diff --git a/statement_test.go b/statement_test.go
index a79b565..1f18f75 100644
--- a/statement_test.go
+++ b/statement_test.go
@@ -1,8 +1,11 @@
+// Copyright (c) 2020 Snowflake Computing Inc. All right reserved.
+
 package gosnowflake
 
 import (
 	"context"
 	"database/sql"
+	"database/sql/driver"
 	"testing"
 )
 
@@ -439,3 +442,58 @@ func TestMultiStatementCountMismatch(t *testing.T) {
 		t.Fatal("should have failed to query multiple statements")
 	}
 }
+
+func TestGetQueryID(t *testing.T) {
+	var db *sql.DB
+	var err error
+
+	if db, err = sql.Open("snowflake", dsn); err != nil {
+		t.Fatalf("failed to open db. %v, err: %v", dsn, err)
+	}
+	defer db.Close()
+
+	ctx := context.TODO()
+	conn, _ := db.Conn(ctx)
+
+	err = conn.Raw(func(x interface{}) error {
+		stmt, err := x.(driver.ConnPrepareContext).PrepareContext(ctx, "select 1")
+		if err != nil {
+			return err
+		}
+		rows, err := stmt.(driver.StmtQueryContext).QueryContext(ctx, nil)
+		if err != nil {
+			return err
+		}
+		defer rows.Close()
+		qid := rows.(SnowflakeResult).QueryID()
+		if qid == "" {
+			t.Fatal("should have returned a query ID string")
+		}
+
+		stmt, err = x.(driver.ConnPrepareContext).PrepareContext(ctx, "selectt 1")
+		if err != nil {
+			return err
+		}
+		rows, err = stmt.(driver.StmtQueryContext).QueryContext(ctx, nil)
+		if err == nil {
+			t.Fatal("should have failed to execute query")
+		}
+		if driverErr, ok := err.(*SnowflakeError); ok {
+			if driverErr.Number != 1003 {
+				t.Fatalf("incorrect error code. expected: 1003, got: %v", driverErr.Number)
+			}
+			if driverErr.QueryID == "" {
+				t.Fatal("should have an associated query ID")
+			}
+		} else {
+			t.Fatal("should have been able to cast to Snowflake Error")
+		}
+		if rows != nil {
+			t.Fatal("rows should be empty")
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("failed to prepare statement. err: %v", err)
+	}
+}