Repository: snowflakedb/gosnowflake
Pull request #379 (2021-02-09 16:17:21)

SNOW-212951 Fetch query result with query Id
With a query Id provided in context to fetch the results of that query.

### Checklist
- [ ] Code compiles correctly
- [ ] Run ``make fmt`` to fix inconsistent formats
- [ ] Run ``make lint`` to get lint errors and fix all of them
- [ ] Created tests which fail without the change (if possible)
- [ ] All tests passing
- [ ] Extended the README / documentation, if necessary


sfc-gh-dwu: What is the behavior of running `result_scan` over a in-progress QID that still hasn't had completed?
Does it:
1. error out due to QID not completed
2. block until completion
3. read in-progress result

I would like to clarify whether the approach of using `select * from result_scan(qid)` solves the need we have with resuming read result of in-progress query.
sfc-gh-dwu: Nvm, looks like we distinguish between in-progress query and completed queries.
Completed queries use `result_scan` and in progress uses `buildRowsForRunningQuery`

Only remaining question would be what if we hit race condition and call `buildRowsForRunningQuery` for a query that is completed?

Patch:
diff --git a/auth.go b/auth.go
index 3c48aab..81970ef 100644
--- a/auth.go
+++ b/auth.go
@@ -415,3 +415,52 @@ func prepareJWTToken(config *Config) (string, error) {
 
 	return tokenString, err
 }
+
+// Authenticate with sc.cfg
+func authenticateWithConfig(sc *snowflakeConn) error {
+	var authData *authResponseMain
+	var samlResponse []byte
+	var proofKey []byte
+	var err error
+	logger.Infof("Authenticating via %v", sc.cfg.Authenticator.String())
+	switch sc.cfg.Authenticator {
+	case AuthTypeExternalBrowser:
+		samlResponse, proofKey, err = authenticateByExternalBrowser(
+			sc.ctx,
+			sc.rest,
+			sc.cfg.Authenticator.String(),
+			sc.cfg.Application,
+			sc.cfg.Account,
+			sc.cfg.User,
+			sc.cfg.Password)
+		if err != nil {
+			sc.cleanup()
+			return err
+		}
+	case AuthTypeOkta:
+		samlResponse, err = authenticateBySAML(
+			sc.ctx,
+			sc.rest,
+			sc.cfg.OktaURL,
+			sc.cfg.Application,
+			sc.cfg.Account,
+			sc.cfg.User,
+			sc.cfg.Password)
+		if err != nil {
+			sc.cleanup()
+			return err
+		}
+	}
+	authData, err = authenticate(
+		sc.ctx,
+		sc,
+		samlResponse,
+		proofKey)
+	if err != nil {
+		sc.cleanup()
+		return err
+	}
+	sc.populateSessionParameters(authData.Parameters)
+	sc.ctx = context.WithValue(sc.ctx, SFSessionIDKey, authData.SessionID)
+	return nil
+}
diff --git a/chunk_downloader.go b/chunk_downloader.go
index ae8ae3d..6103ad4 100644
--- a/chunk_downloader.go
+++ b/chunk_downloader.go
@@ -30,9 +30,9 @@ type chunkDownloader interface {
 	reset()
 	getChunkMetas() []execResponseChunk
 	getQueryResultFormat() resultFormat
+	getRowType() []execResponseRowType
 	setNextChunkDownloader(downloader chunkDownloader)
 	getNextChunkDownloader() chunkDownloader
-	getRowType() []execResponseRowType
 }
 
 type snowflakeChunkDownloader struct {
diff --git a/connection.go b/connection.go
index fe65c10..a5f3a86 100644
--- a/connection.go
+++ b/connection.go
@@ -8,7 +8,9 @@ import (
 	"database/sql/driver"
 	"encoding/json"
 	"fmt"
+	"net/http"
 	"net/url"
+	"regexp"
 	"strconv"
 	"strings"
 	"sync/atomic"
@@ -62,6 +64,12 @@ type snowflakeConn struct {
 	internal        InternalClient
 }
 
+var queryIDPattern = `[\w\-_]+`
+var queryIDRegexp = regexp.MustCompile(queryIDPattern)
+
+const (
+	urlQueriesResultFmt string = "/queries/%s/result"
+)
 // 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
@@ -303,6 +311,25 @@ func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []d
 }
 
 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
+	}
+	if qid == "" {
+		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)
+	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) {
 	logger.WithContext(ctx).Infof("Query: %#v, %v", query, args)
 	if sc.rest == nil {
 		return nil, driver.ErrBadConn
@@ -338,14 +365,16 @@ func (sc *snowflakeConn) QueryContext(ctx context.Context, query string, args []
 
 	rows := new(snowflakeRows)
 	rows.sc = sc
-	rows.ChunkDownloader = populateChunkDownloader(ctx, sc, data.Data)
 	rows.queryID = sc.QueryID
 
 	if sc.isMultiStmt(&data.Data) {
+		// handleMultiQuery is responsible to fill rows with childResults
 		err := sc.handleMultiQuery(ctx, data.Data, rows)
 		if err != nil {
 			return nil, err
 		}
+	} else {
+		rows.addDownloader(populateChunkDownloader(ctx, sc, data.Data))
 	}
 
 	rows.ChunkDownloader.start()
@@ -462,8 +491,8 @@ func (sc *snowflakeConn) handleMultiExec(ctx context.Context, data execResponseD
 	var updatedRows int64
 	childResults := getChildResults(data.ResultIDs, data.ResultTypes)
 	for _, child := range childResults {
-		resultPath := fmt.Sprintf("/queries/%s/result", child.id)
-		childData, err := sc.getQueryResult(ctx, resultPath)
+		resultPath := fmt.Sprintf(urlQueriesResultFmt, child.id)
+		childData, err := sc.getQueryResultResp(ctx, resultPath)
 		if err != nil {
 			logger.Errorf("error: %v", err)
 			code, err := strconv.Atoi(childData.Code)
@@ -507,38 +536,15 @@ 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 {
 	childResults := getChildResults(data.ResultIDs, data.ResultTypes)
-	var nextChunkDownloader chunkDownloader
-	firstResultSet := false
 
 	for _, child := range childResults {
-		resultPath := fmt.Sprintf("/queries/%s/result", child.id)
-		childData, err := sc.getQueryResult(ctx, resultPath)
+		err := sc.rowsForRunningQuery(ctx, child.id, rows)
 		if err != nil {
-			logger.WithContext(ctx).Errorf("error: %v", err)
-			if childData != nil {
-				code, err := strconv.Atoi(childData.Code)
-				if err != nil {
-					return err
-				}
-				return &SnowflakeError{
-					Number:   code,
-					SQLState: childData.Data.SQLState,
-					Message:  err.Error(),
-					QueryID:  childData.Data.QueryID}
-			}
 			return err
 		}
-		if !firstResultSet {
-			// populate rows.ChunkDownloader with the first child
-			rows.ChunkDownloader = populateChunkDownloader(ctx, sc, childData.Data)
-			nextChunkDownloader = rows.ChunkDownloader
-			firstResultSet = true
-		} else {
-			nextChunkDownloader.setNextChunkDownloader(populateChunkDownloader(ctx, sc, childData.Data))
-			nextChunkDownloader = nextChunkDownloader.getNextChunkDownloader()
-		}
 	}
 	return nil
 }
@@ -581,7 +587,7 @@ func getChildResults(IDs string, types string) []childResult {
 	return res
 }
 
-func (sc *snowflakeConn) getQueryResult(ctx context.Context, resultPath string) (*execResponse, error) {
+func (sc *snowflakeConn) getQueryResultResp(ctx context.Context, resultPath string) (*execResponse, error) {
 	headers := getHeaders()
 	if serviceName, ok := sc.cfg.Params[serviceName]; ok {
 		headers[httpHeaderServiceName] = *serviceName
@@ -609,6 +615,119 @@ func (sc *snowflakeConn) getQueryResult(ctx context.Context, resultPath string)
 	return respd, nil
 }
 
+// checkQueryStatus return error==nil means the query completed successfully and there is complete query result to fetch.
+// when the GS could not return a status (when a query was just submitted GS might not be able to return a status)
+// an ErrQueryStatus will be returned.
+// Other than error==nil, There are three error types will be returned:
+// 1. ErrQueryStatus, when GS could not return any status or report error due to any reason - connection,
+// permission, etc.
+// 2, ErrQueryReportedError, if the requested query was terminated, aborted, and GS returned
+// an error status included in query.sfqueryStatusError
+// 3, ErrQueryIsRunning, if the requested query is still running and might have complete result later, these statuses
+// were listed in query.sfqueryStatusRunning
+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 isAsyncMode(ctx context.Context) (bool, error) {
 	val := ctx.Value(asyncMode)
 	if val == nil {
@@ -621,6 +740,25 @@ func isAsyncMode(ctx context.Context) (bool, error) {
 	return boolVal, nil
 }
 
+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,
@@ -693,7 +831,6 @@ func getAsync(
 			res.errChannel <- nil // mark exec status complete
 		} else {
 			rows.sc = sc
-			rows.ChunkDownloader = populateChunkDownloader(ctx, sc, respd.Data)
 			rows.queryID = respd.Data.QueryID
 			if sc.isMultiStmt(&respd.Data) {
 				err = sc.handleMultiQuery(ctx, respd.Data, rows)
@@ -702,6 +839,8 @@ func getAsync(
 					close(errChannel)
 					return
 				}
+			} else {
+				rows.addDownloader(populateChunkDownloader(ctx, sc, respd.Data))
 			}
 			rows.ChunkDownloader.start()
 			rows.errChannel <- nil // mark query status complete
@@ -769,3 +908,58 @@ func populateChunkDownloader(ctx context.Context, sc *snowflakeConn, data execRe
 		},
 	}
 }
+
+func buildSnowFlakeConn(ctx context.Context, config Config) (*snowflakeConn, error) {
+	sc := &snowflakeConn{
+		SequenceCounter: 0,
+		ctx:             ctx,
+		cfg:             &config,
+	}
+	var st http.RoundTripper = SnowflakeTransport
+	if sc.cfg.Transporter == nil {
+		if sc.cfg.InsecureMode {
+			// no revocation check with OCSP. Think twice when you want to enable this option.
+			st = snowflakeInsecureTransport
+		} else {
+			// set OCSP fail open mode
+			ocspResponseCacheLock.Lock()
+			atomic.StoreUint32((*uint32)(&ocspFailOpen), uint32(sc.cfg.OCSPFailOpen))
+			ocspResponseCacheLock.Unlock()
+		}
+	} else {
+		// use the custom transport
+		st = sc.cfg.Transporter
+	}
+	var tokenAccessor TokenAccessor
+	if sc.cfg.TokenAccessor != nil {
+		tokenAccessor = sc.cfg.TokenAccessor
+	} else {
+		tokenAccessor = getSimpleTokenAccessor()
+	}
+	// authenticate
+	sc.rest = &snowflakeRestful{
+		Host:     sc.cfg.Host,
+		Port:     sc.cfg.Port,
+		Protocol: sc.cfg.Protocol,
+		Client: &http.Client{
+			// request timeout including reading response body
+			Timeout:   sc.cfg.ClientTimeout,
+			Transport: st,
+		},
+		TokenAccessor:       tokenAccessor,
+		LoginTimeout:        sc.cfg.LoginTimeout,
+		RequestTimeout:      sc.cfg.RequestTimeout,
+		FuncPost:            postRestful,
+		FuncGet:             getRestful,
+		FuncPostQuery:       postRestfulQuery,
+		FuncPostQueryHelper: postRestfulQueryHelper,
+		FuncRenewSession:    renewRestfulSession,
+		FuncPostAuth:        postAuth,
+		FuncCloseSession:    closeSession,
+		FuncCancelQuery:     cancelQuery,
+		FuncPostAuthSAML:    postAuthSAML,
+		FuncPostAuthOKTA:    postAuthOKTA,
+		FuncGetSSO:          getSSO,
+	}
+	return sc, nil
+}
diff --git a/driver.go b/driver.go
index 70ba159..c465cc7 100644
--- a/driver.go
+++ b/driver.go
@@ -6,9 +6,7 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
-	"net/http"
 	"os"
-	"sync/atomic"
 )
 
 // SnowflakeDriver is a context of Go Driver
@@ -30,102 +28,15 @@ func (d SnowflakeDriver) Open(dsn string) (driver.Conn, error) {
 func (d SnowflakeDriver) OpenWithConfig(ctx context.Context, config Config) (driver.Conn, error) {
 	logger.Info("OpenWithConfig")
 	var err error
-	sc := &snowflakeConn{
-		SequenceCounter: 0,
-		ctx:             ctx,
-		cfg:             &config,
-	}
-	var st http.RoundTripper = SnowflakeTransport
-	if sc.cfg.Transporter == nil {
-		if sc.cfg.InsecureMode {
-			// no revocation check with OCSP. Think twice when you want to enable this option.
-			st = snowflakeInsecureTransport
-		} else {
-			// set OCSP fail open mode
-			ocspResponseCacheLock.Lock()
-			atomic.StoreUint32((*uint32)(&ocspFailOpen), uint32(sc.cfg.OCSPFailOpen))
-			ocspResponseCacheLock.Unlock()
-		}
-	} else {
-		// use the custom transport
-		st = sc.cfg.Transporter
-	}
-	var tokenAccessor TokenAccessor
-	if sc.cfg.TokenAccessor != nil {
-		tokenAccessor = sc.cfg.TokenAccessor
-	} else {
-		tokenAccessor = getSimpleTokenAccessor()
-	}
-	// authenticate
-	sc.rest = &snowflakeRestful{
-		Host:     sc.cfg.Host,
-		Port:     sc.cfg.Port,
-		Protocol: sc.cfg.Protocol,
-		Client: &http.Client{
-			// request timeout including reading response body
-			Timeout:   sc.cfg.ClientTimeout,
-			Transport: st,
-		},
-		TokenAccessor:       tokenAccessor,
-		LoginTimeout:        sc.cfg.LoginTimeout,
-		RequestTimeout:      sc.cfg.RequestTimeout,
-		FuncPost:            postRestful,
-		FuncGet:             getRestful,
-		FuncPostQuery:       postRestfulQuery,
-		FuncPostQueryHelper: postRestfulQueryHelper,
-		FuncRenewSession:    renewRestfulSession,
-		FuncPostAuth:        postAuth,
-		FuncCloseSession:    closeSession,
-		FuncCancelQuery:     cancelQuery,
-		FuncPostAuthSAML:    postAuthSAML,
-		FuncPostAuthOKTA:    postAuthOKTA,
-		FuncGetSSO:          getSSO,
+	sc, err := buildSnowFlakeConn(ctx, config)
+	if err != nil {
+		return nil, err
 	}
-	var authData *authResponseMain
-	var samlResponse []byte
-	var proofKey []byte
 
-	logger.Infof("Authenticating via %v", sc.cfg.Authenticator.String())
-	switch sc.cfg.Authenticator {
-	case AuthTypeExternalBrowser:
-		samlResponse, proofKey, err = authenticateByExternalBrowser(
-			sc.ctx,
-			sc.rest,
-			sc.cfg.Authenticator.String(),
-			sc.cfg.Application,
-			sc.cfg.Account,
-			sc.cfg.User,
-			sc.cfg.Password)
-		if err != nil {
-			sc.cleanup()
-			return nil, err
-		}
-	case AuthTypeOkta:
-		samlResponse, err = authenticateBySAML(
-			sc.ctx,
-			sc.rest,
-			sc.cfg.OktaURL,
-			sc.cfg.Application,
-			sc.cfg.Account,
-			sc.cfg.User,
-			sc.cfg.Password)
-		if err != nil {
-			sc.cleanup()
-			return nil, err
-		}
-	}
-	authData, err = authenticate(
-		sc.ctx,
-		sc,
-		samlResponse,
-		proofKey)
+	err = authenticateWithConfig(sc)
 	if err != nil {
-		sc.cleanup()
 		return nil, err
 	}
-
-	sc.populateSessionParameters(authData.Parameters)
-	sc.ctx = context.WithValue(sc.ctx, SFSessionIDKey, authData.SessionID)
 	sc.startHeartBeat()
 	sc.internal = &httpClient{sr: sc.rest}
 	return sc, nil
diff --git a/errors.go b/errors.go
index 625a51c..7cd54da 100644
--- a/errors.go
+++ b/errors.go
@@ -143,6 +143,17 @@ const (
 	// ErrOCSPNoOCSPResponderURL is an error code for the case where the OCSP responder URL is not attached.
 	ErrOCSPNoOCSPResponderURL = 269004
 
+	/* Query Status*/
+
+	// ErrQueryStatus when check the status of a query, receive error or no status
+	ErrQueryStatus = 279001
+	// ErrQueryIDFormat the query ID given to fetch its result is not valid
+	ErrQueryIDFormat = 279101
+	// ErrQueryReportedError server side reports the query failed with error
+	ErrQueryReportedError = 279201
+	// ErrQueryIsRunning the query is still running
+	ErrQueryIsRunning = 279301
+
 	/* GS error code */
 
 	// ErrSessionGone is an GS error code for the case that session is already closed
diff --git a/query.go b/query.go
index 736c583..cacdcb3 100644
--- a/query.go
+++ b/query.go
@@ -129,3 +129,71 @@ 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 7cee0e9..e057fc2 100644
--- a/restful.go
+++ b/restful.go
@@ -42,6 +42,12 @@ const (
 	heartBeatPath            = "/session/heartbeat"
 )
 
+// FuncGetType httpclient GET method to return http.Response
+type FuncGetType func(context.Context, *snowflakeRestful, *url.URL, map[string]string, time.Duration) (*http.Response, error)
+
+// FuncPostType httpclient POST method to return http.Response
+type FuncPostType func(context.Context, *snowflakeRestful, *url.URL, map[string]string, []byte, time.Duration, bool) (*http.Response, error)
+
 type snowflakeRestful struct {
 	Host           string
 	Port           int
@@ -53,11 +59,12 @@ type snowflakeRestful struct {
 	TokenAccessor TokenAccessor
 	HeartBeat     *heartbeat
 
-	Connection          *snowflakeConn
+	Connection *snowflakeConn
+
 	FuncPostQuery       func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, uuid.UUID, *Config) (*execResponse, error)
 	FuncPostQueryHelper func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, uuid.UUID, *Config) (*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)
+	FuncPost            FuncPostType
+	FuncGet             FuncGetType
 	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
diff --git a/rows.go b/rows.go
index a9778bf..cbdc7e6 100644
--- a/rows.go
+++ b/rows.go
@@ -28,12 +28,13 @@ var (
 )
 
 type snowflakeRows struct {
-	sc              *snowflakeConn
-	ChunkDownloader chunkDownloader
-	queryID         string
-	status          queryStatus
-	err             error
-	errChannel      chan error
+	sc                  *snowflakeConn
+	ChunkDownloader     chunkDownloader
+	tailChunkDownloader chunkDownloader
+	queryID             string
+	status              queryStatus
+	err                 error
+	errChannel          chan error
 }
 
 type snowflakeValue interface{}
@@ -209,3 +210,13 @@ func (rows *snowflakeRows) waitForAsyncQueryStatus() error {
 	}
 	return nil
 }
+
+func (rows *snowflakeRows) addDownloader(newDL chunkDownloader) {
+	if rows.ChunkDownloader == nil {
+		rows.ChunkDownloader = newDL
+		rows.tailChunkDownloader = newDL
+		return
+	}
+	rows.tailChunkDownloader.setNextChunkDownloader(newDL)
+	rows.tailChunkDownloader = newDL
+}
diff --git a/util.go b/util.go
index c61da36..c72f530 100644
--- a/util.go
+++ b/util.go
@@ -25,6 +25,8 @@ const (
 	snowflakeRequestIDKey contextKey = "SNOWFLAKE_REQUEST_ID"
 	// streamChunkDownload determines whether to use a stream based chunk downloader
 	streamChunkDownload contextKey = "STREAM_CHUNK_DOWNLOAD"
+	// FetchResultByID the queryID of query result to fetch
+	FetchResultByID contextKey = "SF_FETCH_RESULT_BY_ID"
 )
 
 // WithMultiStatement returns a context that allows the user to execute the desired number of sql queries in one query


Test patch:
diff --git a/connection_test.go b/connection_test.go
index 708ba62..74dd7b4 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -4,10 +4,14 @@ package gosnowflake
 
 import (
 	"context"
+	"database/sql/driver"
 	"encoding/json"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"net/http"
 	"net/url"
+	"strings"
 	"testing"
 	"time"
 
@@ -101,7 +105,7 @@ func TestGetQueryResultUsesTokenFromTokenAccessor(t *testing.T) {
 		cfg:  &Config{Params: map[string]*string{}},
 		rest: sr,
 	}
-	_, err := sc.getQueryResult(context.Background(), "")
+	_, err := sc.getQueryResultResp(context.Background(), "")
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -186,3 +190,104 @@ func TestCloseIgnoreSessionGone(t *testing.T) {
 		t.Error("Close should let go session gone error")
 	}
 }
+
+// test
+func TestFetchResultByQueryID(t *testing.T) {
+	fetchResultByQueryID(t, nil, nil)
+}
+
+func TestFetchRunningQueryByID(t *testing.T) {
+	fetchResultByQueryID(t, returnQueryIsRunningStatus, nil)
+}
+
+func TestFetchErrorQueryByID(t *testing.T) {
+	err := &SnowflakeError{
+		Number: ErrQueryReportedError}
+	fetchResultByQueryID(t, returnQueryIsErrStatus, err)
+}
+
+func customGetQuery(ctx context.Context, rest *snowflakeRestful, url *url.URL,
+	vals map[string]string, _ time.Duration, jsonStr string) (*http.Response, error) {
+	if strings.Contains(url.Path, "/monitoring/queries/") {
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Body:       ioutil.NopCloser(strings.NewReader(jsonStr)),
+		}, nil
+	}
+	return getRestful(ctx, rest, url, vals, rest.RequestTimeout)
+}
+
+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" : "FILE_SET_INITIALIZATION",
+        "errorCode" : 0, "errorMessage" : null}] }, "code" : null, "message" : null, "success" : true }`
+	return customGetQuery(ctx, rest, fullURL, vals, duration, jsonStr)
+}
+
+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 }`
+	return customGetQuery(ctx, rest, fullURL, vals, duration, jsonStr)
+}
+
+// this function is going to: 1, create a table, 2, query on this table,
+//      3, fetch result of query in step 2, mock running status and error status of that query.
+func fetchResultByQueryID(t *testing.T, customget FuncGetType, expectedFetchErr *SnowflakeError) error {
+
+	config, _ := ParseDSN(dsn)
+	ctx := context.Background()
+	sc, err := buildSnowFlakeConn(ctx, *config)
+	if customget != nil {
+		sc.rest.FuncGet = customget
+	}
+	if err != nil {
+		return err
+	}
+	err = authenticateWithConfig(sc)
+	if err != nil {
+		return err
+	}
+
+	_, 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)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	
+	rows1, 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.Fatalf("Query failed: %v", err)
+	}
+
+	qid := rows1.(SnowflakeResult).GetQueryID()
+	newCtx := context.WithValue(ctx, FetchResultByID, qid)
+
+	rows2, err := sc.QueryContext(newCtx, "", nil)
+	if err != nil {
+		if expectedFetchErr != nil { // got expected error number
+			if expectedFetchErr.Number == err.(*SnowflakeError).Number {
+				return nil
+			}
+		}
+		t.Fatalf("Fetch Query Result by ID failed: %v", err)
+	}
+
+	dest := make([]driver.Value, 2)
+	cnt := 0
+	for {
+		err = rows2.Next(dest)
+		if err != nil {
+			if err == io.EOF {
+				break
+			} else {
+				t.Fatalf("unexpected error: %v", err)
+			}
+		}
+		cnt++
+	}
+	if cnt != 10 {
+		t.Fatalf("rowcount is not expected 10: %v", cnt)
+	}
+	return nil
+}
diff --git a/statement_test.go b/statement_test.go
index e506906..dbf696d 100644
--- a/statement_test.go
+++ b/statement_test.go
@@ -58,6 +58,16 @@ func TestGetQueryID(t *testing.T) {
 	}
 }
 
+func getDB(t *testing.T) *sql.DB {
+	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)
+	}
+	return db
+}
+
 func TestEmitQueryID(t *testing.T) {
 	var db *sql.DB
 	var err error
@@ -101,3 +111,54 @@ func TestEmitQueryID(t *testing.T) {
 		t.Errorf("number of rows didn't match. expected: %v, got: %v", numrows, cnt)
 	}
 }
+
+// End-to-end test to fetch result with queryID
+func TestE2EFetchResultByID(t *testing.T) {
+	var err error
+	db := getDB(t)
+
+	defer db.Close()
+	_, err = db.Exec("create or replace table test_fetch_result(c1 number, c2 string)" +
+		"as select 10, 'z'")
+	if err != nil {
+		t.Fatalf("failed to create table: %v", err)
+	}
+
+	ctx := context.Background()
+	conn, _ := db.Conn(ctx)
+	err = conn.Raw(func(x interface{}) error {
+		stmt, err := x.(driver.ConnPrepareContext).PrepareContext(ctx, "select * from test_fetch_result")
+		if err != nil {
+			return err
+		}
+
+		rows1, err := stmt.(driver.StmtQueryContext).QueryContext(ctx, nil)
+		if err != nil {
+			return err
+		}
+		qid := rows1.(SnowflakeResult).GetQueryID()
+
+		newCtx := context.WithValue(context.Background(), FetchResultByID, qid)
+		rows2, err := db.QueryContext(newCtx, "")
+		if err != nil {
+			t.Fatalf("Fetch Query Result by ID failed: %v", err)
+		}
+		var c1 sql.NullInt64
+		var c2 sql.NullString
+		for rows2.Next() {
+			err = rows2.Scan(&c1, &c2)
+		}
+		if c1.Int64 != 10 || c2.String != "z" {
+			t.Fatalf("Query result is not expected: %v", err)
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("failed to drop table: %v", err)
+	}
+
+	_, err = db.Exec("drop table if exists test_fetch_result")
+	if err != nil {
+		t.Fatalf("failed to drop table: %v", err)
+	}
+}