Repository: snowflakedb/gosnowflake
Pull request #467 (2021-08-10 22:28:09)

SNOW-421272 Query Status Polling
### Description
Make query status polling a publicly available API
- refactor connection functions to monitoring and utils
- make CheckQueryStatus a public function via SnowflakeConnection interface

### 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


rleighton: Can you have the status returned when a query is complete have:
```
        StartTime    int64
	EndTime      int64
	ScanBytes    int64
	ProducedRows int64
```
(which are in the api struct returned see https://github.com/snowflakedb/gosnowflake/pull/458 ).

If you do not then a query is required to get this info from the client.


rleighton: Thank you!

Patch:
diff --git a/connection.go b/connection.go
index faa07d8..d06104c 100644
--- a/connection.go
+++ b/connection.go
@@ -3,13 +3,11 @@
 package gosnowflake
 
 import (
-	"bytes"
 	"context"
 	"database/sql"
 	"database/sql/driver"
 	"encoding/json"
 	"fmt"
-	"io"
 	"net/http"
 	"net/url"
 	"os"
@@ -19,8 +17,6 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
-
-	"github.com/google/uuid"
 )
 
 const (
@@ -35,13 +31,8 @@ const (
 )
 
 const (
-	statementTypeIDMulti = int64(0x1000)
-
+	statementTypeIDMulti            = int64(0x1000)
 	statementTypeIDDml              = int64(0x3000)
-	statementTypeIDInsert           = statementTypeIDDml + int64(0x100)
-	statementTypeIDUpdate           = statementTypeIDDml + int64(0x200)
-	statementTypeIDDelete           = statementTypeIDDml + int64(0x300)
-	statementTypeIDMerge            = statementTypeIDDml + int64(0x400)
 	statementTypeIDMultiTableInsert = statementTypeIDDml + int64(0x500)
 )
 
@@ -71,24 +62,11 @@ type snowflakeConn struct {
 	internal        InternalClient
 }
 
-var queryIDPattern = `[\w\-_]+`
-var queryIDRegexp = regexp.MustCompile(queryIDPattern)
-
-const (
-	urlQueriesResultFmt string = "/queries/%s/result"
+var (
+	queryIDPattern = `[\w\-_]+`
+	queryIDRegexp  = regexp.MustCompile(queryIDPattern)
 )
 
-// isDml returns true if the statement type code is in the range of DML.
-func (sc *snowflakeConn) isDml(v int64) bool {
-	return statementTypeIDDml <= v && v <= statementTypeIDMultiTableInsert
-}
-
-// isMultiStmt returns true if the statement type code is of type multistatement
-// Note that the statement type code is also equivalent to type INSERT, so an additional check of the name is required
-func (sc *snowflakeConn) isMultiStmt(data *execResponseData) bool {
-	return data.StatementTypeID == statementTypeIDMulti && data.RowType[0].Name == "multiple statement execution"
-}
-
 func (sc *snowflakeConn) exec(
 	ctx context.Context,
 	query string,
@@ -117,7 +95,8 @@ func (sc *snowflakeConn) exec(
 	if len(bindings) > 0 {
 		arrayBindThreshold := sc.getArrayBindStageThreshold()
 		numBinds := arrayBindValueCount(bindings)
-		if 0 < arrayBindThreshold && arrayBindThreshold <= numBinds && !describeOnly && isArrayBind(bindings) {
+		if 0 < arrayBindThreshold && arrayBindThreshold <= numBinds &&
+			!describeOnly && isArrayBind(bindings) {
 			// bulk array insert binding
 			uploader := bindUploader{
 				sc:        sc,
@@ -152,7 +131,8 @@ func (sc *snowflakeConn) exec(
 	}
 
 	var data *execResponse
-	data, err = sc.rest.FuncPostQuery(ctx, sc.rest, &url.Values{}, headers, jsonBody, sc.rest.RequestTimeout, requestID, sc.cfg)
+	data, err = sc.rest.FuncPostQuery(ctx, sc.rest, &url.Values{}, headers,
+		jsonBody, sc.rest.RequestTimeout, requestID, sc.cfg)
 	if err != nil {
 		return data, err
 	}
@@ -176,28 +156,7 @@ func (sc *snowflakeConn) exec(
 		}
 	}
 	if isFileTransfer(query) {
-		sfa := snowflakeFileTransferAgent{
-			sc:      sc,
-			data:    &data.Data,
-			command: query,
-			options: new(SnowflakeFileTransferOptions),
-		}
-		if fs := getFileStream(ctx); fs != nil {
-			sfa.sourceStream = fs
-			if isInternal {
-				sfa.data.AutoCompress = false
-			}
-		}
-		if op := getFileTransferOptions(ctx); op != nil {
-			sfa.options = op
-		}
-		if sfa.options.multiPartThreshold == 0 {
-			sfa.options.multiPartThreshold = dataSizeThreshold
-		}
-		if err = sfa.execute(); err != nil {
-			return nil, err
-		}
-		data, err = sfa.result()
+		data, err = sc.processFileTransfer(ctx, data, query, isInternal)
 		if err != nil {
 			return nil, err
 		}
@@ -218,7 +177,10 @@ func (sc *snowflakeConn) Begin() (driver.Tx, error) {
 	return sc.BeginTx(sc.ctx, driver.TxOptions{})
 }
 
-func (sc *snowflakeConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+func (sc *snowflakeConn) BeginTx(
+	ctx context.Context,
+	opts driver.TxOptions) (
+	driver.Tx, error) {
 	logger.WithContext(ctx).Info("BeginTx")
 	if opts.ReadOnly {
 		return nil, &SnowflakeError{
@@ -238,11 +200,11 @@ func (sc *snowflakeConn) BeginTx(ctx context.Context, opts driver.TxOptions) (dr
 		return nil, driver.ErrBadConn
 	}
 	isDesc := isDescribeOnly(ctx)
-	_, err := sc.exec(ctx, "BEGIN", false /* noResult */, false /* isInternal */, isDesc, nil)
-	if err != nil {
+	if _, err := sc.exec(ctx, "BEGIN", false, /* noResult */
+		false /* isInternal */, isDesc, nil); err != nil {
 		return nil, err
 	}
-	return &snowflakeTx{sc}, err
+	return &snowflakeTx{sc}, nil
 }
 
 func (sc *snowflakeConn) cleanup() {
@@ -257,8 +219,7 @@ func (sc *snowflakeConn) Close() (err error) {
 	sc.stopHeartBeat()
 
 	if !sc.cfg.KeepSessionAlive {
-		err = sc.rest.FuncCloseSession(sc.ctx, sc.rest, sc.rest.RequestTimeout)
-		if err != nil {
+		if err = sc.rest.FuncCloseSession(sc.ctx, sc.rest, sc.rest.RequestTimeout); err != nil {
 			logger.Error(err)
 		}
 	}
@@ -266,13 +227,17 @@ func (sc *snowflakeConn) Close() (err error) {
 	return nil
 }
 
-func (sc *snowflakeConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
+func (sc *snowflakeConn) PrepareContext(
+	ctx context.Context,
+	query string) (
+	driver.Stmt, error) {
 	logger.WithContext(sc.ctx).Infoln("Prepare")
 	if sc.rest == nil {
 		return nil, driver.ErrBadConn
 	}
 	noResult := isAsyncMode(ctx)
-	data, err := sc.exec(ctx, query, noResult, false /* isInternal */, true /* describeOnly */, []driver.NamedValue{})
+	data, err := sc.exec(ctx, query, noResult, false, /* isInternal */
+		true /* describeOnly */, []driver.NamedValue{})
 	if err != nil {
 		if data != nil {
 			code, err := strconv.Atoi(data.Code)
@@ -295,7 +260,11 @@ func (sc *snowflakeConn) PrepareContext(ctx context.Context, query string) (driv
 	return stmt, nil
 }
 
-func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
+func (sc *snowflakeConn) ExecContext(
+	ctx context.Context,
+	query string,
+	args []driver.NamedValue) (
+	driver.Result, error) {
 	logger.WithContext(ctx).Infof("Exec: %#v, %v", query, args)
 	if sc.rest == nil {
 		return nil, driver.ErrBadConn
@@ -326,7 +295,7 @@ func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []d
 		return data.Data.AsyncResult, nil
 	}
 
-	if sc.isDml(data.Data.StatementTypeID) {
+	if isDml(data.Data.StatementTypeID) {
 		// collects all values from the returned row sets
 		updatedRows, err := updateRows(data.Data)
 		if err != nil {
@@ -338,14 +307,18 @@ func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []d
 			insertID:     -1,
 			queryID:      sc.QueryID,
 		}, nil // last insert id is not supported by Snowflake
-	} else if sc.isMultiStmt(&data.Data) {
+	} else if isMultiStmt(&data.Data) {
 		return sc.handleMultiExec(ctx, data.Data)
 	}
 	logger.Debug("DDL")
 	return driver.ResultNoRows, nil
 }
 
-func (sc *snowflakeConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+func (sc *snowflakeConn) QueryContext(
+	ctx context.Context,
+	query string,
+	args []driver.NamedValue) (
+	driver.Rows, error) {
 	qid, err := getResumeQueryID(ctx)
 	if err != nil {
 		return nil, err
@@ -354,17 +327,20 @@ func (sc *snowflakeConn) QueryContext(ctx context.Context, query string, args []
 		return sc.queryContextInternal(ctx, query, args)
 	}
 
-	// first we will check the status of this particular query to find out if there is result to fetch
-	err = sc.checkQueryStatus(ctx, qid)
+	// check the query status to find out if there is a result to fetch
+	_, err = sc.checkQueryStatus(ctx, qid)
 	if err == nil || (err != nil && err.(*SnowflakeError).Number == ErrQueryIsRunning) {
 		// the query is running. Rows object will be returned from here.
 		return sc.buildRowsForRunningQuery(ctx, qid)
 	}
-
 	return nil, err
 }
 
-func (sc *snowflakeConn) queryContextInternal(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+func (sc *snowflakeConn) queryContextInternal(
+	ctx context.Context,
+	query string,
+	args []driver.NamedValue) (
+	driver.Rows, error) {
 	logger.WithContext(ctx).Infof("Query: %#v, %v", query, args)
 	if sc.rest == nil {
 		return nil, driver.ErrBadConn
@@ -400,7 +376,7 @@ func (sc *snowflakeConn) queryContextInternal(ctx context.Context, query string,
 	rows.sc = sc
 	rows.queryID = sc.QueryID
 
-	if sc.isMultiStmt(&data.Data) {
+	if isMultiStmt(&data.Data) {
 		// handleMultiQuery is responsible to fill rows with childResults
 		if err = sc.handleMultiQuery(ctx, data.Data, rows); err != nil {
 			return nil, err
@@ -439,7 +415,8 @@ func (sc *snowflakeConn) Ping(ctx context.Context) error {
 	noResult := isAsyncMode(ctx)
 	isDesc := isDescribeOnly(ctx)
 	// TODO: handle isInternal
-	_, err := sc.exec(ctx, "SELECT 1", noResult, false /* isInternal */, isDesc, []driver.NamedValue{})
+	_, err := sc.exec(ctx, "SELECT 1", noResult, false, /* isInternal */
+		isDesc, []driver.NamedValue{})
 	return err
 }
 
@@ -452,72 +429,26 @@ func (sc *snowflakeConn) CheckNamedValue(nv *driver.NamedValue) error {
 	return nil
 }
 
-func (sc *snowflakeConn) populateSessionParameters(parameters []nameValueParameter) {
-	// other session parameters (not all)
-	logger.WithContext(sc.ctx).Infof("params: %#v", parameters)
-	for _, param := range parameters {
-		v := ""
-		switch param.Value.(type) {
-		case int64:
-			if vv, ok := param.Value.(int64); ok {
-				v = strconv.FormatInt(vv, 10)
-			}
-		case float64:
-			if vv, ok := param.Value.(float64); ok {
-				v = strconv.FormatFloat(vv, 'g', -1, 64)
-			}
-		case bool:
-			if vv, ok := param.Value.(bool); ok {
-				v = strconv.FormatBool(vv)
-			}
-		default:
-			if vv, ok := param.Value.(string); ok {
-				v = vv
-			}
-		}
-		logger.Debugf("parameter. name: %v, value: %v", param.Name, v)
-		sc.cfg.Params[strings.ToLower(param.Name)] = &v
-	}
-}
-
-func (sc *snowflakeConn) isClientSessionKeepAliveEnabled() bool {
-	v, ok := sc.cfg.Params[sessionClientSessionKeepAlive]
-	if !ok {
-		return false
-	}
-	return strings.Compare(*v, "true") == 0
-}
-
-func (sc *snowflakeConn) getArrayBindStageThreshold() int {
-	v, ok := sc.cfg.Params[sessionArrayBindStageThreshold]
-	if !ok {
-		return 0
-	}
-	num, err := strconv.Atoi(*v)
-	if err != nil {
-		return 0
-	}
-	return num
-}
-
-func (sc *snowflakeConn) startHeartBeat() {
-	if !sc.isClientSessionKeepAliveEnabled() {
-		return
-	}
-	sc.rest.HeartBeat = &heartbeat{
-		restful: sc.rest,
-	}
-	sc.rest.HeartBeat.start()
-}
-
-func (sc *snowflakeConn) stopHeartBeat() {
-	if !sc.isClientSessionKeepAliveEnabled() {
-		return
-	}
-	sc.rest.HeartBeat.stop()
-}
-
-func (sc *snowflakeConn) handleMultiExec(ctx context.Context, data execResponseData) (driver.Result, error) {
+func (sc *snowflakeConn) GetQueryStatus(
+	ctx context.Context,
+	queryID string) (
+	*SnowflakeQueryStatus, error) {
+	queryRet, err := sc.checkQueryStatus(ctx, queryID)
+	return &SnowflakeQueryStatus{
+		queryRet.SQLText,
+		queryRet.StartTime,
+		queryRet.EndTime,
+		queryRet.ErrorCode,
+		queryRet.ErrorMessage,
+		queryRet.Stats.ScanBytes,
+		queryRet.Stats.ProducedRows,
+	}, err
+}
+
+func (sc *snowflakeConn) handleMultiExec(
+	ctx context.Context,
+	data execResponseData) (
+	driver.Result, error) {
 	var updatedRows int64
 	childResults := getChildResults(data.ResultIDs, data.ResultTypes)
 	for _, child := range childResults {
@@ -538,7 +469,7 @@ func (sc *snowflakeConn) handleMultiExec(ctx context.Context, data execResponseD
 			}
 			return nil, err
 		}
-		if sc.isDml(childData.Data.StatementTypeID) {
+		if isDml(childData.Data.StatementTypeID) {
 			count, err := updateRows(childData.Data)
 			if err != nil {
 				logger.WithContext(ctx).Errorf("error: %v", err)
@@ -566,10 +497,13 @@ func (sc *snowflakeConn) handleMultiExec(ctx context.Context, data execResponseD
 	}, nil
 }
 
-// Fill the correspondent rows and add chunk downloader into the rows when iterate the childResults
-func (sc *snowflakeConn) handleMultiQuery(ctx context.Context, data execResponseData, rows *snowflakeRows) error {
+// Fill the corresponding rows and add chunk downloader into the rows when
+// iterating across the childResults
+func (sc *snowflakeConn) handleMultiQuery(
+	ctx context.Context,
+	data execResponseData,
+	rows *snowflakeRows) error {
 	childResults := getChildResults(data.ResultIDs, data.ResultTypes)
-
 	for _, child := range childResults {
 		if err := sc.rowsForRunningQuery(ctx, child.id, rows); err != nil {
 			return err
@@ -578,224 +512,6 @@ func (sc *snowflakeConn) handleMultiQuery(ctx context.Context, data execResponse
 	return nil
 }
 
-func setResultType(ctx context.Context, resType resultType) context.Context {
-	return context.WithValue(ctx, snowflakeResultType, resType)
-}
-
-func getResultType(ctx context.Context) resultType {
-	return ctx.Value(snowflakeResultType).(resultType)
-}
-
-func updateRows(data execResponseData) (int64, error) {
-	var count int64
-	for i, n := 0, len(data.RowType); i < n; i++ {
-		v, err := strconv.ParseInt(*data.RowSet[0][i], 10, 64)
-		if err != nil {
-			return -1, err
-		}
-		count += v
-	}
-	return count, nil
-}
-
-type childResult struct {
-	id  string
-	typ string
-}
-
-func getChildResults(IDs string, types string) []childResult {
-	if IDs == "" {
-		return nil
-	}
-	queryIDs := strings.Split(IDs, ",")
-	resultTypes := strings.Split(types, ",")
-	res := make([]childResult, len(queryIDs))
-	for i, id := range queryIDs {
-		res[i] = childResult{id, resultTypes[i]}
-	}
-	return res
-}
-
-func (sc *snowflakeConn) getQueryResultResp(ctx context.Context, resultPath string) (*execResponse, error) {
-	headers := getHeaders()
-	if serviceName, ok := sc.cfg.Params[serviceName]; ok {
-		headers[httpHeaderServiceName] = *serviceName
-	}
-	param := make(url.Values)
-	param.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
-	param.Add("clientStartTime", strconv.FormatInt(time.Now().Unix(), 10))
-	param.Add(requestGUIDKey, uuid.New().String())
-	token, _, _ := sc.rest.TokenAccessor.GetTokens()
-	if token != "" {
-		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
-	}
-	url := sc.rest.getFullURL(resultPath, &param)
-	res, err := sc.rest.FuncGet(ctx, sc.rest, url, headers, sc.rest.RequestTimeout)
-	if err != nil {
-		logger.WithContext(ctx).Errorf("failed to get response. err: %v", err)
-		return nil, err
-	}
-	var respd *execResponse
-	err = json.NewDecoder(res.Body).Decode(&respd)
-	if err != nil {
-		logger.WithContext(ctx).Errorf("failed to decode JSON. err: %v", err)
-		return nil, err
-	}
-	return respd, nil
-}
-
-// checkQueryStatus returns the status given the query ID. If successful,
-// the error will be nil, indicating there is a complete query result to fetch.
-// Other than nil, there are three error types that can be returned:
-// 1. ErrQueryStatus, if GS cannot return any kind of status due to any reason,
-// i.e. connection, permission, if a query was just submitted, etc.
-// 2, ErrQueryReportedError, if the requested query was terminated or aborted
-// and GS returned an error status included in query. SFQueryFailedWithError
-// 3, ErrQueryIsRunning, if the requested query is still running and might have
-// a complete result later, these statuses were listed in query. SFQueryRunning
-func (sc *snowflakeConn) checkQueryStatus(ctx context.Context, qid string) error {
-	headers := make(map[string]string)
-	param := make(url.Values)
-	param.Add(requestGUIDKey, uuid.New().String())
-	if tok, _, _ := sc.rest.TokenAccessor.GetTokens(); tok != "" {
-		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, tok)
-	}
-	resultPath := fmt.Sprintf("/monitoring/queries/%s", qid)
-	url := sc.rest.getFullURL(resultPath, &param)
-
-	res, err := sc.rest.FuncGet(ctx, sc.rest, url, headers, sc.rest.RequestTimeout)
-	if err != nil {
-		logger.WithContext(ctx).Errorf("failed to get response. err: %v", err)
-		return err
-	}
-	var statusResp = statusResponse{}
-
-	err = json.NewDecoder(res.Body).Decode(&statusResp)
-	if err != nil {
-		logger.WithContext(ctx).Errorf("failed to decode JSON. err: %v", err)
-		return err
-	}
-
-	if !statusResp.Success || len(statusResp.Data.Queries) == 0 {
-		logger.WithContext(ctx).Errorf("status query returned not-success or no status returned.")
-		return &SnowflakeError{
-			Number:  ErrQueryStatus,
-			Message: "status query returned not-success or no status returned. Please retry"}
-	}
-
-	var queryRet = statusResp.Data.Queries[0]
-	if queryRet.ErrorCode != 0 {
-		return &SnowflakeError{
-			Number: ErrQueryStatus,
-			Message: fmt.Sprintf("server ErrorCode=%d, ErrorMessage=%s",
-				queryRet.ErrorCode, queryRet.ErrorMessage),
-			IncludeQueryID: true,
-			QueryID:        qid,
-		}
-	}
-
-	// returned errorCode is 0. Now check what is the returned status of the query.
-	var qstatus = strToSFQueryStatus(queryRet.Status)
-	if sfqStatusIsAnError(qstatus) {
-		return &SnowflakeError{
-			Number: ErrQueryReportedError,
-			Message: fmt.Sprintf("%s: status from server: [%s]",
-				queryRet.ErrorMessage, queryRet.Status),
-			IncludeQueryID: true,
-			QueryID:        qid,
-		}
-	}
-
-	if sfqStatusIsStillRunning(qstatus) {
-		return &SnowflakeError{
-			Number: ErrQueryIsRunning,
-			Message: fmt.Sprintf("%s: status from server: [%s]",
-				queryRet.ErrorMessage, queryRet.Status),
-			IncludeQueryID: true,
-			QueryID:        qid,
-		}
-	}
-
-	//success
-	return nil
-}
-
-// Fetch query result for a query id from /queries/<qid>/result endpoint.
-func (sc *snowflakeConn) rowsForRunningQuery(ctx context.Context, qid string, rows *snowflakeRows) error {
-	resultPath := fmt.Sprintf(urlQueriesResultFmt, qid)
-	resp, err := sc.getQueryResultResp(ctx, resultPath)
-	if err != nil {
-		logger.WithContext(ctx).Errorf("error: %v", err)
-		if resp != nil {
-			code, err := strconv.Atoi(resp.Code)
-			if err != nil {
-				return err
-			}
-			return &SnowflakeError{
-				Number:   code,
-				SQLState: resp.Data.SQLState,
-				Message:  err.Error(),
-				QueryID:  resp.Data.QueryID}
-		}
-		return err
-	}
-	rows.addDownloader(populateChunkDownloader(ctx, sc, resp.Data))
-	return nil
-}
-
-// prepare a Rows object to return for query of 'qid'
-func (sc *snowflakeConn) buildRowsForRunningQuery(ctx context.Context, qid string) (driver.Rows, error) {
-	rows := new(snowflakeRows)
-	rows.sc = sc
-	rows.queryID = qid
-	err := sc.rowsForRunningQuery(ctx, qid, rows)
-	if err != nil {
-		return nil, err
-	}
-	rows.ChunkDownloader.start()
-	return rows, err
-}
-
-func (sc *snowflakeConn) setupOCSPPrivatelink(app string, host string) error {
-	ocspCacheServer := fmt.Sprintf("http://ocsp.%v/ocsp_response_cache.json", host)
-	if err := os.Setenv(cacheServerURLEnv, ocspCacheServer); err != nil {
-		return err
-	}
-	ocspRetryHost := fmt.Sprintf("http://ocsp.%v/retry/", host) + "%v/%v"
-	if err := os.Setenv(ocspRetryURLEnv, ocspRetryHost); err != nil {
-		return err
-	}
-	return nil
-}
-
-func isAsyncMode(ctx context.Context) bool {
-	val := ctx.Value(asyncMode)
-	if val == nil {
-		return false
-	}
-	a, ok := val.(bool)
-	return ok && a
-}
-
-func getResumeQueryID(ctx context.Context) (string, error) {
-	val := ctx.Value(fetchResultByID)
-	if val == nil {
-		return "", nil
-	}
-	strVal, ok := val.(string)
-	if !ok {
-		return "", fmt.Errorf("failed to cast val %+v to string", val)
-	}
-	// so there is a queryID in context for which we want to fetch the result
-	if !queryIDRegexp.MatchString(strVal) {
-		return strVal, &SnowflakeError{
-			Number:  ErrQueryIDFormat,
-			Message: "Invalid QID",
-			QueryID: strVal}
-	}
-	return strVal, nil
-}
-
 func getAsync(
 	ctx context.Context,
 	sr *snowflakeRestful,
@@ -848,9 +564,9 @@ func getAsync(
 	if respd.Success {
 		if resType == execResultType {
 			res.insertID = -1
-			if sc.isDml(respd.Data.StatementTypeID) {
+			if isDml(respd.Data.StatementTypeID) {
 				res.affectedRows, _ = updateRows(respd.Data)
-			} else if sc.isMultiStmt(&respd.Data) {
+			} else if isMultiStmt(&respd.Data) {
 				r, err := sc.handleMultiExec(ctx, respd.Data)
 				if err != nil {
 					res.errChannel <- err
@@ -869,7 +585,7 @@ func getAsync(
 		} else {
 			rows.sc = sc
 			rows.queryID = respd.Data.QueryID
-			if sc.isMultiStmt(&respd.Data) {
+			if isMultiStmt(&respd.Data) {
 				err = sc.handleMultiQuery(ctx, respd.Data, rows)
 				if err != nil {
 					rows.errChannel <- err
@@ -901,84 +617,6 @@ func getAsync(
 	}
 }
 
-func getQueryIDChan(ctx context.Context) chan<- string {
-	v := ctx.Value(queryIDChannel)
-	if v == nil {
-		return nil
-	}
-	c, _ := v.(chan<- string)
-	return c
-}
-
-func getFileStream(ctx context.Context) *bytes.Buffer {
-	s := ctx.Value(fileStreamFile)
-	r, ok := s.(io.Reader)
-	if !ok {
-		return nil
-	}
-	buf := new(bytes.Buffer)
-	buf.ReadFrom(r)
-	return buf
-}
-
-func getFileTransferOptions(ctx context.Context) *SnowflakeFileTransferOptions {
-	v := ctx.Value(fileTransferOptions)
-	if v == nil {
-		return nil
-	}
-	o, ok := v.(*SnowflakeFileTransferOptions)
-	if !ok {
-		return nil
-	}
-	return o
-}
-
-func isDescribeOnly(ctx context.Context) bool {
-	v := ctx.Value(describeOnly)
-	if v == nil {
-		return false
-	}
-	d, ok := v.(bool)
-	return ok && d
-}
-
-// returns snowflake chunk downloader by default or stream based chunk
-// downloader if option provided through context
-func populateChunkDownloader(ctx context.Context, sc *snowflakeConn, data execResponseData) chunkDownloader {
-	if useStreamDownloader(ctx) && data.QueryResultFormat == "json" {
-		// stream chunk downloading only works for row based data formats, i.e. json
-		fetcher := &httpStreamChunkFetcher{
-			ctx:      ctx,
-			client:   sc.rest.Client,
-			clientIP: sc.cfg.ClientIP,
-			headers:  data.ChunkHeaders,
-			qrmk:     data.Qrmk,
-		}
-		return newStreamChunkDownloader(ctx, fetcher, data.Total, data.RowType, data.RowSet, data.Chunks)
-	}
-
-	return &snowflakeChunkDownloader{
-		sc:                 sc,
-		ctx:                ctx,
-		CurrentChunk:       make([]chunkRowType, len(data.RowSet)),
-		ChunkMetas:         data.Chunks,
-		Total:              data.Total,
-		TotalRowIndex:      int64(-1),
-		CellCount:          len(data.RowType),
-		Qrmk:               data.Qrmk,
-		QueryResultFormat:  data.QueryResultFormat,
-		ChunkHeader:        data.ChunkHeaders,
-		FuncDownload:       downloadChunk,
-		FuncDownloadHelper: downloadChunkHelper,
-		FuncGet:            getChunk,
-		RowSet: rowSetType{
-			RowType:      data.RowType,
-			JSON:         data.RowSet,
-			RowSetBase64: data.RowSetBase64,
-		},
-	}
-}
-
 func buildSnowflakeConn(ctx context.Context, config Config) (*snowflakeConn, error) {
 	sc := &snowflakeConn{
 		SequenceCounter: 0,
diff --git a/connection_util.go b/connection_util.go
new file mode 100644
index 0000000..31a3a59
--- /dev/null
+++ b/connection_util.go
@@ -0,0 +1,278 @@
+// Copyright (c) 2021 Snowflake Computing Inc. All right reserved.
+
+package gosnowflake
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"os"
+	"strconv"
+	"strings"
+)
+
+func (sc *snowflakeConn) isClientSessionKeepAliveEnabled() bool {
+	v, ok := sc.cfg.Params[sessionClientSessionKeepAlive]
+	if !ok {
+		return false
+	}
+	return strings.Compare(*v, "true") == 0
+}
+
+func (sc *snowflakeConn) startHeartBeat() {
+	if !sc.isClientSessionKeepAliveEnabled() {
+		return
+	}
+	sc.rest.HeartBeat = &heartbeat{
+		restful: sc.rest,
+	}
+	sc.rest.HeartBeat.start()
+}
+
+func (sc *snowflakeConn) stopHeartBeat() {
+	if !sc.isClientSessionKeepAliveEnabled() {
+		return
+	}
+	sc.rest.HeartBeat.stop()
+}
+
+func (sc *snowflakeConn) getArrayBindStageThreshold() int {
+	v, ok := sc.cfg.Params[sessionArrayBindStageThreshold]
+	if !ok {
+		return 0
+	}
+	num, err := strconv.Atoi(*v)
+	if err != nil {
+		return 0
+	}
+	return num
+}
+
+// processFileTransfer creates a snowflakeFileTransferAgent object to process
+// any PUT/GET commands with their specified options
+func (sc *snowflakeConn) processFileTransfer(
+	ctx context.Context,
+	data *execResponse,
+	query string,
+	isInternal bool) (
+	*execResponse, error) {
+	sfa := snowflakeFileTransferAgent{
+		sc:      sc,
+		data:    &data.Data,
+		command: query,
+		options: new(SnowflakeFileTransferOptions),
+	}
+	if fs := getFileStream(ctx); fs != nil {
+		sfa.sourceStream = fs
+		if isInternal {
+			sfa.data.AutoCompress = false
+		}
+	}
+	if op := getFileTransferOptions(ctx); op != nil {
+		sfa.options = op
+	}
+	if sfa.options.multiPartThreshold == 0 {
+		sfa.options.multiPartThreshold = dataSizeThreshold
+	}
+	if err := sfa.execute(); err != nil {
+		return nil, err
+	}
+	data, err := sfa.result()
+	if err != nil {
+		return nil, err
+	}
+	return data, nil
+}
+
+func getFileStream(ctx context.Context) *bytes.Buffer {
+	s := ctx.Value(fileStreamFile)
+	r, ok := s.(io.Reader)
+	if !ok {
+		return nil
+	}
+	buf := new(bytes.Buffer)
+	buf.ReadFrom(r)
+	return buf
+}
+
+func getFileTransferOptions(ctx context.Context) *SnowflakeFileTransferOptions {
+	v := ctx.Value(fileTransferOptions)
+	if v == nil {
+		return nil
+	}
+	o, ok := v.(*SnowflakeFileTransferOptions)
+	if !ok {
+		return nil
+	}
+	return o
+}
+
+func (sc *snowflakeConn) populateSessionParameters(parameters []nameValueParameter) {
+	// other session parameters (not all)
+	logger.WithContext(sc.ctx).Infof("params: %#v", parameters)
+	for _, param := range parameters {
+		v := ""
+		switch param.Value.(type) {
+		case int64:
+			if vv, ok := param.Value.(int64); ok {
+				v = strconv.FormatInt(vv, 10)
+			}
+		case float64:
+			if vv, ok := param.Value.(float64); ok {
+				v = strconv.FormatFloat(vv, 'g', -1, 64)
+			}
+		case bool:
+			if vv, ok := param.Value.(bool); ok {
+				v = strconv.FormatBool(vv)
+			}
+		default:
+			if vv, ok := param.Value.(string); ok {
+				v = vv
+			}
+		}
+		logger.Debugf("parameter. name: %v, value: %v", param.Name, v)
+		sc.cfg.Params[strings.ToLower(param.Name)] = &v
+	}
+}
+
+func isAsyncMode(ctx context.Context) bool {
+	val := ctx.Value(asyncMode)
+	if val == nil {
+		return false
+	}
+	a, ok := val.(bool)
+	return ok && a
+}
+
+func isDescribeOnly(ctx context.Context) bool {
+	v := ctx.Value(describeOnly)
+	if v == nil {
+		return false
+	}
+	d, ok := v.(bool)
+	return ok && d
+}
+
+func setResultType(ctx context.Context, resType resultType) context.Context {
+	return context.WithValue(ctx, snowflakeResultType, resType)
+}
+
+func getResultType(ctx context.Context) resultType {
+	return ctx.Value(snowflakeResultType).(resultType)
+}
+
+// isDml returns true if the statement type code is in the range of DML.
+func isDml(v int64) bool {
+	return statementTypeIDDml <= v && v <= statementTypeIDMultiTableInsert
+}
+
+func updateRows(data execResponseData) (int64, error) {
+	var count int64
+	for i, n := 0, len(data.RowType); i < n; i++ {
+		v, err := strconv.ParseInt(*data.RowSet[0][i], 10, 64)
+		if err != nil {
+			return -1, err
+		}
+		count += v
+	}
+	return count, nil
+}
+
+// isMultiStmt returns true if the statement code is of type multistatement
+// Note that the statement type code is also equivalent to type INSERT, so an
+// additional check of the name is required
+func isMultiStmt(data *execResponseData) bool {
+	return data.StatementTypeID == statementTypeIDMulti &&
+		data.RowType[0].Name == "multiple statement execution"
+}
+
+func getResumeQueryID(ctx context.Context) (string, error) {
+	val := ctx.Value(fetchResultByID)
+	if val == nil {
+		return "", nil
+	}
+	strVal, ok := val.(string)
+	if !ok {
+		return "", fmt.Errorf("failed to cast val %+v to string", val)
+	}
+	// so there is a queryID in context for which we want to fetch the result
+	if !queryIDRegexp.MatchString(strVal) {
+		return strVal, &SnowflakeError{
+			Number:  ErrQueryIDFormat,
+			Message: "Invalid QID",
+			QueryID: strVal}
+	}
+	return strVal, nil
+}
+
+type childResult struct {
+	id  string
+	typ string
+}
+
+func getChildResults(IDs string, types string) []childResult {
+	if IDs == "" {
+		return nil
+	}
+	queryIDs := strings.Split(IDs, ",")
+	resultTypes := strings.Split(types, ",")
+	res := make([]childResult, len(queryIDs))
+	for i, id := range queryIDs {
+		res[i] = childResult{id, resultTypes[i]}
+	}
+	return res
+}
+
+// returns snowflake chunk downloader by default or stream based chunk
+// downloader if option provided through context
+func populateChunkDownloader(
+	ctx context.Context,
+	sc *snowflakeConn,
+	data execResponseData) chunkDownloader {
+	if useStreamDownloader(ctx) && data.QueryResultFormat == "json" {
+		// stream chunk downloading only works for row based data formats, i.e. json
+		fetcher := &httpStreamChunkFetcher{
+			ctx:      ctx,
+			client:   sc.rest.Client,
+			clientIP: sc.cfg.ClientIP,
+			headers:  data.ChunkHeaders,
+			qrmk:     data.Qrmk,
+		}
+		return newStreamChunkDownloader(ctx, fetcher, data.Total, data.RowType,
+			data.RowSet, data.Chunks)
+	}
+
+	return &snowflakeChunkDownloader{
+		sc:                 sc,
+		ctx:                ctx,
+		CurrentChunk:       make([]chunkRowType, len(data.RowSet)),
+		ChunkMetas:         data.Chunks,
+		Total:              data.Total,
+		TotalRowIndex:      int64(-1),
+		CellCount:          len(data.RowType),
+		Qrmk:               data.Qrmk,
+		QueryResultFormat:  data.QueryResultFormat,
+		ChunkHeader:        data.ChunkHeaders,
+		FuncDownload:       downloadChunk,
+		FuncDownloadHelper: downloadChunkHelper,
+		FuncGet:            getChunk,
+		RowSet: rowSetType{
+			RowType:      data.RowType,
+			JSON:         data.RowSet,
+			RowSetBase64: data.RowSetBase64,
+		},
+	}
+}
+
+func (sc *snowflakeConn) setupOCSPPrivatelink(app string, host string) error {
+	ocspCacheServer := fmt.Sprintf("http://ocsp.%v/ocsp_response_cache.json", host)
+	if err := os.Setenv(cacheServerURLEnv, ocspCacheServer); err != nil {
+		return err
+	}
+	ocspRetryHost := fmt.Sprintf("http://ocsp.%v/retry/", host) + "%v/%v"
+	if err := os.Setenv(ocspRetryURLEnv, ocspRetryHost); err != nil {
+		return err
+	}
+	return nil
+}
diff --git a/monitoring.go b/monitoring.go
new file mode 100644
index 0000000..600bff4
--- /dev/null
+++ b/monitoring.go
@@ -0,0 +1,267 @@
+// Copyright (c) 2021 Snowflake Computing Inc. All right reserved.
+
+package gosnowflake
+
+import (
+	"context"
+	"database/sql/driver"
+	"encoding/json"
+	"fmt"
+	"net/url"
+	"strconv"
+	"time"
+
+	"github.com/google/uuid"
+)
+
+const urlQueriesResultFmt = "/queries/%s/result"
+
+// QueryStatus is status returned from server
+type queryResultStatus int
+
+// Query Status defined at server side
+const (
+	SFQueryRunning queryResultStatus = iota
+	SFQueryAborting
+	SFQuerySuccess
+	SFQueryFailedWithError
+	SFQueryAborted
+	SFQueryQueued
+	SFQueryFailedWithIncident
+	SFQueryDisconnected
+	SFQueryResumingWarehouse
+	// SFQueryQueueRepairingWarehouse present in QueryDTO.java.
+	SFQueryQueueRepairingWarehouse
+	SFQueryRestarted
+	// SFQueryBlocked is when a statement is waiting on a lock on resource held
+	// by another statement.
+	SFQueryBlocked
+	SFQueryNoData
+)
+
+func (qs queryResultStatus) String() string {
+	return [...]string{"RUNNING", "ABORTING", "SUCCESS", "FAILED_WITH_ERROR",
+		"ABORTED", "QUEUED", "FAILED_WITH_INCIDENT", "DISCONNECTED",
+		"RESUMING_WAREHOUSE", "QUEUED_REPAIRING_WAREHOUSE", "RESTARTED",
+		"BLOCKED", "NO_DATA"}[qs]
+}
+
+func (qs queryResultStatus) isRunning() bool {
+	switch qs {
+	case SFQueryRunning, SFQueryResumingWarehouse, SFQueryQueued,
+		SFQueryQueueRepairingWarehouse, SFQueryNoData:
+		return true
+	default:
+		return false
+	}
+}
+
+func (qs queryResultStatus) isError() bool {
+	switch qs {
+	case SFQueryAborting, SFQueryFailedWithError, SFQueryAborted,
+		SFQueryFailedWithIncident, SFQueryDisconnected, SFQueryBlocked:
+		return true
+	default:
+		return false
+	}
+}
+
+var strQueryStatusMap = map[string]queryResultStatus{"RUNNING": SFQueryRunning,
+	"ABORTING": SFQueryAborting, "SUCCESS": SFQuerySuccess,
+	"FAILED_WITH_ERROR": SFQueryFailedWithError, "ABORTED": SFQueryAborted,
+	"QUEUED": SFQueryQueued, "FAILED_WITH_INCIDENT": SFQueryFailedWithIncident,
+	"DISCONNECTED":               SFQueryDisconnected,
+	"RESUMING_WAREHOUSE":         SFQueryResumingWarehouse,
+	"QUEUED_REPAIRING_WAREHOUSE": SFQueryQueueRepairingWarehouse,
+	"RESTARTED":                  SFQueryRestarted,
+	"BLOCKED":                    SFQueryBlocked, "NO_DATA": SFQueryNoData}
+
+type retStatus struct {
+	Status       string   `json:"status"`
+	SQLText      string   `json:"sqlText"`
+	StartTime    int64    `json:"startTime"`
+	EndTime      int64    `json:"endTime"`
+	ErrorCode    int      `json:"errorCode"`
+	ErrorMessage string   `json:"errorMessage"`
+	Stats        retStats `json:"stats"`
+}
+
+type retStats struct {
+	ScanBytes    int64 `json:"scanBytes"`
+	ProducedRows int64 `json:"producedRows"`
+}
+
+type statusResponse struct {
+	Data struct {
+		Queries []retStatus `json:"queries"`
+	} `json:"data"`
+	Message string `json:"message"`
+	Code    string `json:"code"`
+	Success bool   `json:"success"`
+}
+
+func strToQueryStatus(in string) queryResultStatus {
+	return strQueryStatusMap[in]
+}
+
+// SnowflakeQueryStatus is the query status metadata of a snowflake query
+type SnowflakeQueryStatus struct {
+	SQLText      string
+	StartTime    int64
+	EndTime      int64
+	ErrorCode    int
+	ErrorMessage string
+	ScanBytes    int64
+	ProducedRows int64
+}
+
+// SnowflakeConnection is a wrapper to snowflakeConn that exposes API functions
+type SnowflakeConnection interface {
+	GetQueryStatus(ctx context.Context, queryID string) (*SnowflakeQueryStatus, error)
+}
+
+// checkQueryStatus returns the status given the query ID. If successful,
+// the error will be nil, indicating there is a complete query result to fetch.
+// Other than nil, there are three error types that can be returned:
+// 1. ErrQueryStatus, if GS cannot return any kind of status due to any reason,
+// i.e. connection, permission, if a query was just submitted, etc.
+// 2, ErrQueryReportedError, if the requested query was terminated or aborted
+// and GS returned an error status included in query. SFQueryFailedWithError
+// 3, ErrQueryIsRunning, if the requested query is still running and might have
+// a complete result later, these statuses were listed in query. SFQueryRunning
+func (sc *snowflakeConn) checkQueryStatus(
+	ctx context.Context,
+	qid string) (
+	*retStatus, error) {
+	headers := make(map[string]string)
+	param := make(url.Values)
+	param.Add(requestGUIDKey, uuid.New().String())
+	if tok, _, _ := sc.rest.TokenAccessor.GetTokens(); tok != "" {
+		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, tok)
+	}
+	resultPath := fmt.Sprintf("/monitoring/queries/%s", qid)
+	url := sc.rest.getFullURL(resultPath, &param)
+
+	res, err := sc.rest.FuncGet(ctx, sc.rest, url, headers, sc.rest.RequestTimeout)
+	if err != nil {
+		logger.WithContext(ctx).Errorf("failed to get response. err: %v", err)
+		return nil, err
+	}
+	var statusResp = statusResponse{}
+	if err = json.NewDecoder(res.Body).Decode(&statusResp); err != nil {
+		logger.WithContext(ctx).Errorf("failed to decode JSON. err: %v", err)
+		return nil, err
+	}
+
+	if !statusResp.Success || len(statusResp.Data.Queries) == 0 {
+		logger.WithContext(ctx).Errorf("status query returned not-success or no status returned.")
+		return nil, &SnowflakeError{
+			Number:  ErrQueryStatus,
+			Message: "status query returned not-success or no status returned. Please retry"}
+	}
+
+	queryRet := statusResp.Data.Queries[0]
+	if queryRet.ErrorCode != 0 {
+		return &queryRet, &SnowflakeError{
+			Number: ErrQueryStatus,
+			Message: fmt.Sprintf("server ErrorCode=%d, ErrorMessage=%s",
+				queryRet.ErrorCode, queryRet.ErrorMessage),
+			IncludeQueryID: true,
+			QueryID:        qid,
+		}
+	}
+
+	// returned errorCode is 0. Now check what is the returned status of the query.
+	qStatus := strToQueryStatus(queryRet.Status)
+	if qStatus.isError() {
+		return &queryRet, &SnowflakeError{
+			Number: ErrQueryReportedError,
+			Message: fmt.Sprintf("%s: status from server: [%s]",
+				queryRet.ErrorMessage, queryRet.Status),
+			IncludeQueryID: true,
+			QueryID:        qid,
+		}
+	}
+
+	if qStatus.isRunning() {
+		return &queryRet, &SnowflakeError{
+			Number: ErrQueryIsRunning,
+			Message: fmt.Sprintf("%s: status from server: [%s]",
+				queryRet.ErrorMessage, queryRet.Status),
+			IncludeQueryID: true,
+			QueryID:        qid,
+		}
+	}
+	//success
+	return &queryRet, nil
+}
+
+func (sc *snowflakeConn) getQueryResultResp(
+	ctx context.Context,
+	resultPath string) (
+	*execResponse, error) {
+	headers := getHeaders()
+	if serviceName, ok := sc.cfg.Params[serviceName]; ok {
+		headers[httpHeaderServiceName] = *serviceName
+	}
+	param := make(url.Values)
+	param.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
+	param.Add("clientStartTime", strconv.FormatInt(time.Now().Unix(), 10))
+	param.Add(requestGUIDKey, uuid.New().String())
+	token, _, _ := sc.rest.TokenAccessor.GetTokens()
+	if token != "" {
+		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
+	}
+	url := sc.rest.getFullURL(resultPath, &param)
+	res, err := sc.rest.FuncGet(ctx, sc.rest, url, headers, sc.rest.RequestTimeout)
+	if err != nil {
+		logger.WithContext(ctx).Errorf("failed to get response. err: %v", err)
+		return nil, err
+	}
+	var respd *execResponse
+	if err = json.NewDecoder(res.Body).Decode(&respd); err != nil {
+		logger.WithContext(ctx).Errorf("failed to decode JSON. err: %v", err)
+		return nil, err
+	}
+	return respd, nil
+}
+
+// Fetch query result for a query id from /queries/<qid>/result endpoint.
+func (sc *snowflakeConn) rowsForRunningQuery(
+	ctx context.Context, qid string,
+	rows *snowflakeRows) error {
+	resultPath := fmt.Sprintf(urlQueriesResultFmt, qid)
+	resp, err := sc.getQueryResultResp(ctx, resultPath)
+	if err != nil {
+		logger.WithContext(ctx).Errorf("error: %v", err)
+		if resp != nil {
+			code, err := strconv.Atoi(resp.Code)
+			if err != nil {
+				return err
+			}
+			return &SnowflakeError{
+				Number:   code,
+				SQLState: resp.Data.SQLState,
+				Message:  err.Error(),
+				QueryID:  resp.Data.QueryID}
+		}
+		return err
+	}
+	rows.addDownloader(populateChunkDownloader(ctx, sc, resp.Data))
+	return nil
+}
+
+// prepare a Rows object to return for query of 'qid'
+func (sc *snowflakeConn) buildRowsForRunningQuery(
+	ctx context.Context,
+	qid string) (
+	driver.Rows, error) {
+	rows := new(snowflakeRows)
+	rows.sc = sc
+	rows.queryID = qid
+	if err := sc.rowsForRunningQuery(ctx, qid, rows); err != nil {
+		return nil, err
+	}
+	rows.ChunkDownloader.start()
+	return rows, nil
+}
diff --git a/query.go b/query.go
index 4950ea8..54a0391 100644
--- a/query.go
+++ b/query.go
@@ -129,71 +129,3 @@ type execResponse struct {
 	Code    string           `json:"code"`
 	Success bool             `json:"success"`
 }
-
-// QueryStatusFromServer status returned from server
-type QueryStatusFromServer int
-
-// Query Status defined at server side
-const (
-	SFQueryRunning = iota
-	SFQueryAborting
-	SFQuerySuccess
-	SFQueryFailedWithError
-	SFQueryAborted
-	SFQueryQueued
-	SFQueryFailedWithIncident
-	SFQueryDisconnected
-	SFQueryResumingWarehouse
-	// SFQueryQueueRepairingWarehouse present in QueryDTO.java.
-	SFQueryQueueRepairingWarehouse
-	SFQueryRestarted
-
-	// SFQueryBlocked the state when a statement is waiting on a lock on resource held by another statement.
-	SFQueryBlocked
-	SFQueryNoData
-)
-
-var sfQueryStrStatusMap = map[string]QueryStatusFromServer{
-	"RUNNING": SFQueryRunning, "ABORTING": SFQueryAborting, "SUCCESS": SFQuerySuccess,
-	"FAILED_WITH_ERROR": SFQueryFailedWithError, "ABORTED": SFQueryAborted, "QUEUED": SFQueryQueued,
-	"FAILED_WITH_INCIDENT": SFQueryFailedWithIncident, "DISCONNECTED": SFQueryDisconnected,
-	"RESUMING_WAREHOUSE": SFQueryResumingWarehouse, "QUEUED_REPAIRING_WAREHOUSE": SFQueryQueueRepairingWarehouse,
-	"RESTARTED": SFQueryRestarted, "BLOCKED": SFQueryBlocked, "NO_DATA": SFQueryNoData}
-
-var dummy = struct{}{}
-var sfqueryStatusRunning = map[QueryStatusFromServer]struct{}{
-	SFQueryRunning: dummy, SFQueryResumingWarehouse: dummy, SFQueryQueued: dummy,
-	SFQueryQueueRepairingWarehouse: dummy, SFQueryNoData: dummy}
-
-var sfqueryStatusError = map[QueryStatusFromServer]struct{}{
-	SFQueryAborting: dummy, SFQueryFailedWithError: dummy, SFQueryAborted: dummy,
-	SFQueryFailedWithIncident: dummy, SFQueryDisconnected: dummy, SFQueryBlocked: dummy}
-
-type retStatus struct {
-	Status       string `json:"status"`
-	ErrorMessage string `json:"errorMessage"`
-	ErrorCode    int    `json:"errorCode"`
-}
-
-type statusResponse struct {
-	Data struct {
-		Queries []retStatus `json:"queries"`
-	} `json:"data"`
-	Message string `json:"message"`
-	Code    string `json:"code"`
-	Success bool   `json:"success"`
-}
-
-func strToSFQueryStatus(in string) QueryStatusFromServer {
-	return sfQueryStrStatusMap[in]
-}
-
-func sfqStatusIsAnError(status QueryStatusFromServer) bool {
-	_, exist := sfqueryStatusError[status]
-	return exist
-}
-
-func sfqStatusIsStillRunning(status QueryStatusFromServer) bool {
-	_, exist := sfqueryStatusRunning[status]
-	return exist
-}
diff --git a/restful.go b/restful.go
index 19e2706..81f040d 100644
--- a/restful.go
+++ b/restful.go
@@ -506,3 +506,12 @@ func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID uuid.UUID,
 		MessageArgs: []interface{}{resp.StatusCode, fullURL},
 	}
 }
+
+func getQueryIDChan(ctx context.Context) chan<- string {
+	v := ctx.Value(queryIDChannel)
+	if v == nil {
+		return nil
+	}
+	c, _ := v.(chan<- string)
+	return c
+}
diff --git a/result.go b/result.go
index 0f0ba06..8699e07 100644
--- a/result.go
+++ b/result.go
@@ -5,8 +5,6 @@ package gosnowflake
 type queryStatus string
 
 const (
-	/* query status */
-
 	// QueryStatusWaiting denotes a query execution waiting to happen
 	QueryStatusWaiting queryStatus = "queryStatusWaiting"
 	// QueryStatusInProgress denotes a query execution in progress


Test patch:
diff --git a/connection_test.go b/connection_test.go
index bded99a..1e4a10c 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -75,7 +75,7 @@ func TestExecWithEmptyRequestID(t *testing.T) {
 		cfg:  &Config{Params: map[string]*string{}},
 		rest: sr,
 	}
-	if _, err := sc.exec(ctx, "", false /* noResult */, false /* isInternal */,
+	if _, err := sc.exec(ctx, "", false /* noResult */, false, /* isInternal */
 		false /* describeOnly */, nil); err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -146,16 +146,17 @@ func TestExecWithSpecificRequestID(t *testing.T) {
 		cfg:  &Config{Params: map[string]*string{}},
 		rest: sr,
 	}
-	if _, err := sc.exec(ctx, "", false /* noResult */, false /* isInternal */,
+	if _, err := sc.exec(ctx, "", false /* noResult */, false, /* isInternal */
 		false /* describeOnly */, nil); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 }
 
 // TestServiceName tests two things:
-// 1. request header would contain X-Snowflake-Service if the cfg parameters contains SERVICE_NAME
-// 2. SERVICE_NAME would be update by response payload
-// It is achieved through an interactive postQueryMock that would generate response based on header
+// 1. request header contains X-Snowflake-Service if the cfg parameters
+// contains SERVICE_NAME
+// 2. SERVICE_NAME is updated by response payload
+// Uses interactive postQueryMock that generates a response based on header
 func TestServiceName(t *testing.T) {
 	sr := &snowflakeRestful{
 		FuncPostQuery: postQueryMock,
@@ -168,7 +169,7 @@ func TestServiceName(t *testing.T) {
 
 	expectServiceName := serviceNameStub
 	for i := 0; i < 5; i++ {
-		sc.exec(context.TODO(), "", false /* noResult */,
+		sc.exec(context.TODO(), "", false, /* noResult */
 			false /* isInternal */, false /* describeOnly */, nil)
 		if actualServiceName, ok := sc.cfg.Params[serviceName]; ok {
 			if *actualServiceName != expectServiceName {
@@ -178,7 +179,6 @@ func TestServiceName(t *testing.T) {
 		} else {
 			t.Error("No service name in the response")
 		}
-
 		expectServiceName += serviceNameAppend
 	}
 }
@@ -257,7 +257,7 @@ func customGetQuery(ctx context.Context, rest *snowflakeRestful, url *url.URL,
 
 func returnQueryIsRunningStatus(ctx context.Context, rest *snowflakeRestful, fullURL *url.URL,
 	vals map[string]string, duration time.Duration) (*http.Response, error) {
-	var jsonStr = `{"data" : { "queries" : [{"status" : "RUNNING", "state" :
+	jsonStr := `{"data" : { "queries" : [{"status" : "RUNNING", "state" :
 		"FILE_SET_INITIALIZATION", "errorCode" : 0, "errorMessage" : null}] },
 		"code" : null, "message" : null, "success" : true }`
 	return customGetQuery(ctx, rest, fullURL, vals, duration, jsonStr)
@@ -265,9 +265,9 @@ func returnQueryIsRunningStatus(ctx context.Context, rest *snowflakeRestful, ful
 
 func returnQueryIsErrStatus(ctx context.Context, rest *snowflakeRestful, fullURL *url.URL,
 	vals map[string]string, duration time.Duration) (*http.Response, error) {
-	var jsonStr = `{"data" : { "queries" : [{"status" : "FAILED_WITH_ERROR",
-		"errorCode" : 0, "errorMessage" : ""}] }, "code" : null,
-		"message" : null, "success" : true }`
+	jsonStr := `{"data" : { "queries" : [{"status" : "FAILED_WITH_ERROR",
+		"errorCode" : 0, "errorMessage" : ""}] }, "code" : null, "message" :
+		null, "success" : true }`
 	return customGetQuery(ctx, rest, fullURL, vals, duration, jsonStr)
 }
 
@@ -276,13 +276,13 @@ func returnQueryIsErrStatus(ctx context.Context, rest *snowflakeRestful, fullURL
 // of that query.
 func fetchResultByQueryID(
 	t *testing.T,
-	customget FuncGetType,
+	customGet FuncGetType,
 	expectedFetchErr *SnowflakeError) error {
 	config, _ := ParseDSN(dsn)
 	ctx := context.Background()
 	sc, err := buildSnowflakeConn(ctx, *config)
-	if customget != nil {
-		sc.rest.FuncGet = customget
+	if customGet != nil {
+		sc.rest.FuncGet = customGet
 	}
 	if err != nil {
 		return err
@@ -353,3 +353,39 @@ func TestPrivateLink(t *testing.T) {
 		t.Errorf("expected: %v, got: %v", expectedURL, retryURL)
 	}
 }
+
+func TestGetQueryStatus(t *testing.T) {
+	config, _ := ParseDSN(dsn)
+	ctx := context.Background()
+	sc, err := buildSnowflakeConn(ctx, *config)
+	if err != nil {
+		t.Error(err)
+	}
+	if err = authenticateWithConfig(sc); err != nil {
+		t.Error(err)
+	}
+
+	if _, err = sc.Exec(`create or replace table ut_conn(c1 number, c2 string)
+		as (select seq4() as seq, concat('str',to_varchar(seq)) as str1 from
+		table(generator(rowcount => 100)))`, nil); err != nil {
+		t.Error(err)
+	}
+
+	rows, err := sc.QueryContext(ctx, "select min(c1) as ms, sum(c1) from ut_conn group by (c1 % 10) order by ms", nil)
+	if err != nil {
+		t.Error(err)
+	}
+	qid := rows.(SnowflakeResult).GetQueryID()
+
+	// use conn as type holder for SnowflakeConnection placeholder
+	var conn interface{} = sc
+	qStatus, err := conn.(SnowflakeConnection).GetQueryStatus(ctx, qid)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if qStatus.ErrorCode != 0 || qStatus.ScanBytes != 1536 || qStatus.ProducedRows != 10 {
+		t.Errorf("expected no error. got: %v, scan bytes: %v, produced rows: %v",
+			qStatus.ErrorCode, qStatus.ScanBytes, qStatus.ProducedRows)
+	}
+}