Repository: snowflakedb/gosnowflake
Pull request #331 (2020-09-09 22:23:27)

SNOW-160505 Async Support
### Description
This PR allows asynchronous querying by enabling the AsyncExec parameter to be set via the driver, which, when sent to the server, will not block waiting for the entire query to finish.

### 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 , please add more comments to the PR or the code itself to explain what change you are making to help the reviewers get context quickly.

Patch:
diff --git a/connection.go b/connection.go
index 78b19f8..302e528 100644
--- a/connection.go
+++ b/connection.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2020 Snowflake Computing Inc. All right reserved.
+// Copyright (c) 2017-2021 Snowflake Computing Inc. All right reserved.
 
 package gosnowflake
 
@@ -34,6 +34,14 @@ const (
 	serviceName                            = "service_name"
 )
 
+type resultType string
+
+const (
+	snowflakeResultType paramKey   = "snowflakeResultType"
+	execResultType      resultType = "exec"
+	queryResultType     resultType = "query"
+)
+
 type snowflakeConn struct {
 	ctx             context.Context
 	cfg             *Config
@@ -45,13 +53,7 @@ type snowflakeConn struct {
 
 // isDml returns true if the statement type code is in the range of DML.
 func (sc *snowflakeConn) isDml(v int64) bool {
-	switch v {
-	case statementTypeIDDml, statementTypeIDInsert,
-		statementTypeIDUpdate, statementTypeIDDelete,
-		statementTypeIDMerge, statementTypeIDMultiTableInsert:
-		return true
-	}
-	return false
+	return statementTypeIDDml <= v && v <= statementTypeIDMultiTableInsert
 }
 
 // isMultiStmt returns true if the statement type code is of type multistatement
@@ -73,9 +75,15 @@ func (sc *snowflakeConn) exec(
 	req := execRequest{
 		SQLText:    query,
 		AsyncExec:  noResult,
+		Parameters: map[string]interface{}{},
+		IsInternal: isInternal,
 		SequenceID: counter,
 	}
-	req.IsInternal = isInternal
+	if key := ctx.Value(MultiStatementCount); key != nil {
+		req.Parameters[string(MultiStatementCount)] = key
+	}
+	logger.WithContext(ctx).Infof("parameters: %v", req.Parameters)
+
 	tsmode := "TIMESTAMP_NTZ"
 	idx := 1
 	if len(bindings) > 0 {
@@ -106,12 +114,7 @@ func (sc *snowflakeConn) exec(
 			}
 		}
 	}
-	multiCount := ctx.Value(MultiStatementCount)
-	if multiCount != nil {
-		req.Parameters = map[string]interface{}{string(MultiStatementCount): multiCount}
-	}
 	logger.WithContext(ctx).Infof("bindings: %v", req.Bindings)
-	logger.WithContext(ctx).Infof("parameters: %v", req.Parameters)
 
 	headers := make(map[string]string)
 	headers["Content-Type"] = headerContentTypeApplicationJSON
@@ -232,8 +235,13 @@ func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []d
 	if sc.rest == nil {
 		return nil, driver.ErrBadConn
 	}
-	// TODO: handle noResult and isInternal
-	data, err := sc.exec(ctx, query, false, false, args)
+	noResult, err := isAsyncMode(ctx)
+	if err != nil {
+		return nil, err
+	}
+	// TODO handle isInternal
+	ctx = setResultType(ctx, execResultType)
+	data, err := sc.exec(ctx, query, noResult, false, args)
 	if err != nil {
 		logger.WithContext(ctx).Infof("error: %v", err)
 		if data != nil {
@@ -250,6 +258,11 @@ func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []d
 		return nil, err
 	}
 
+	// if async exec, return result object right away
+	if noResult {
+		return data.Data.AsyncResult, nil
+	}
+
 	var updatedRows int64
 	if sc.isDml(data.Data.StatementTypeID) {
 		// collects all values from the returned row sets
@@ -319,8 +332,13 @@ func (sc *snowflakeConn) QueryContext(ctx context.Context, query string, args []
 	if sc.rest == nil {
 		return nil, driver.ErrBadConn
 	}
-	// TODO: handle noResult and isInternal
-	data, err := sc.exec(ctx, query, false, false, args)
+
+	noResult, err := isAsyncMode(ctx)
+	if err != nil {
+		return nil, err
+	}
+	ctx = setResultType(ctx, queryResultType)
+	data, err := sc.exec(ctx, query, noResult, false, args)
 	if err != nil {
 		logger.WithContext(ctx).Errorf("error: %v", err)
 		if data != nil {
@@ -337,6 +355,11 @@ func (sc *snowflakeConn) QueryContext(ctx context.Context, query string, args []
 		return nil, err
 	}
 
+	// if async query, return row object right away
+	if noResult {
+		return data.Data.AsyncRows, nil
+	}
+
 	rows := new(snowflakeRows)
 	rows.sc = sc
 	rows.RowType = data.Data.RowType
@@ -419,8 +442,11 @@ func (sc *snowflakeConn) Ping(ctx context.Context) error {
 	if sc.rest == nil {
 		return driver.ErrBadConn
 	}
-	// TODO: handle noResult and isInternal
-	_, err := sc.exec(ctx, "SELECT 1", false, false, []driver.NamedValue{})
+	noResult, err := isAsyncMode(ctx)
+	if err != nil {
+		return err
+	}
+	_, err = sc.exec(ctx, "SELECT 1", noResult, false, []driver.NamedValue{})
 	return err
 }
 
@@ -487,6 +513,14 @@ func (sc *snowflakeConn) stopHeartBeat() {
 	sc.rest.HeartBeat.stop()
 }
 
+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++ {
@@ -547,6 +581,116 @@ func (sc *snowflakeConn) getQueryResult(ctx context.Context, resultPath string)
 	return respd, nil
 }
 
+func isAsyncMode(ctx context.Context) (bool, error) {
+	val := ctx.Value(AsyncMode)
+	if val == nil {
+		return false, nil
+	}
+	boolVal, ok := val.(bool)
+	if !ok {
+		return false, fmt.Errorf("failed to cast val %+v to bool", val)
+	}
+	return boolVal, nil
+}
+
+func getAsync(
+	ctx context.Context,
+	sr *snowflakeRestful,
+	headers map[string]string,
+	URL *url.URL,
+	timeout time.Duration,
+	res *snowflakeResult,
+	rows *snowflakeRows,
+) {
+	resType := getResultType(ctx)
+	var errChannel chan error
+	sfError := &SnowflakeError{
+		Number: -1,
+	}
+	if resType == execResultType {
+		errChannel = res.errChannel
+		sfError.QueryID = res.queryID
+	} else {
+		errChannel = rows.errChannel
+		sfError.QueryID = rows.queryID
+	}
+	defer close(errChannel)
+	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, sr.Token)
+	resp, err := sr.FuncGet(ctx, sr, URL, headers, timeout)
+	if err != nil {
+		logger.WithContext(ctx).Errorf("failed to get response. err: %v", err)
+		sfError.Message = err.Error()
+		errChannel <- sfError
+		close(errChannel)
+		return
+	}
+	if resp.Body != nil {
+		defer resp.Body.Close()
+	}
+
+	respd := execResponse{}
+	err = json.NewDecoder(resp.Body).Decode(&respd)
+	resp.Body.Close()
+	if err != nil {
+		logger.WithContext(ctx).Errorf("failed to decode JSON. err: %v", err)
+		sfError.Message = err.Error()
+		errChannel <- sfError
+		close(errChannel)
+		return
+	}
+
+	if respd.Success {
+		if resType == execResultType {
+			res.affectedRows, _ = updateRows(respd.Data)
+			res.insertID = -1
+			res.queryID = respd.Data.QueryID
+			res.errChannel <- nil // mark exec status complete
+		} else {
+			sc := &snowflakeConn{rest: sr}
+			rows.sc = sc
+			rows.RowType = respd.Data.RowType
+			rows.ChunkDownloader = &snowflakeChunkDownloader{
+				sc:                 sc,
+				ctx:                ctx,
+				CurrentChunk:       make([]chunkRowType, len(respd.Data.RowSet)),
+				ChunkMetas:         respd.Data.Chunks,
+				Total:              respd.Data.Total,
+				TotalRowIndex:      int64(-1),
+				CellCount:          len(respd.Data.RowType),
+				Qrmk:               respd.Data.Qrmk,
+				QueryResultFormat:  respd.Data.QueryResultFormat,
+				ChunkHeader:        respd.Data.ChunkHeaders,
+				FuncDownload:       downloadChunk,
+				FuncDownloadHelper: downloadChunkHelper,
+				FuncGet:            getChunk,
+				RowSet: rowSetType{RowType: respd.Data.RowType,
+					JSON:         respd.Data.RowSet,
+					RowSetBase64: respd.Data.RowSetBase64,
+				},
+			}
+			rows.queryID = respd.Data.QueryID
+			rows.ChunkDownloader.start()
+			rows.errChannel <- nil // mark query status complete
+		}
+	} else {
+		var code int
+		if respd.Code != "" {
+			code, err = strconv.Atoi(respd.Code)
+			if err != nil {
+				code = -1
+			}
+		} else {
+			code = -1
+		}
+		errChannel <- &SnowflakeError{
+			Number:   code,
+			SQLState: respd.Data.SQLState,
+			Message:  respd.Message,
+			QueryID:  respd.Data.QueryID,
+		}
+	}
+}
+
 func populateChunkDownloader(ctx context.Context, sc *snowflakeConn, data execResponseData) *snowflakeChunkDownloader {
 	return &snowflakeChunkDownloader{
 		sc:                 sc,
diff --git a/driver.go b/driver.go
index 149ecc1..2bc5f35 100644
--- a/driver.go
+++ b/driver.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2019 Snowflake Computing Inc. All right reserved.
+// Copyright (c) 2017-2021 Snowflake Computing Inc. All right reserved.
 
 package gosnowflake
 
@@ -31,8 +31,8 @@ func (d SnowflakeDriver) OpenWithConfig(ctx context.Context, config Config) (dri
 	var err error
 	sc := &snowflakeConn{
 		SequenceCounter: 0,
-		ctx: ctx,
-		cfg: &config,
+		ctx:             ctx,
+		cfg:             &config,
 	}
 	var st http.RoundTripper = SnowflakeTransport
 	if sc.cfg.Transporter == nil {
@@ -126,4 +126,5 @@ var logger = CreateDefaultLogger()
 
 func init() {
 	sql.Register("snowflake", &SnowflakeDriver{})
+	logger.SetLogLevel("error")
 }
diff --git a/query.go b/query.go
index 7f70ff0..9df7e30 100644
--- a/query.go
+++ b/query.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2019 Snowflake Computing Inc. All right reserved.
+// Copyright (c) 2017-2021 Snowflake Computing Inc. All right reserved.
 
 package gosnowflake
 
@@ -69,6 +69,10 @@ type execResponseData struct {
 	ResultIDs         string        `json:"resultIds,omitempty"`
 	ResultTypes       string        `json:"resultTypes,omitempty"`
 	QueryResultFormat string        `json:"queryResultFormat,omitempty"`
+
+	// async response placeholders
+	AsyncResult *snowflakeResult `json:"asyncResult,omitempty"`
+	AsyncRows   *snowflakeRows   `json:"asyncRows,omitempty"`
 }
 
 type execResponse struct {
diff --git a/restful.go b/restful.go
index eae316a..aac54a1 100644
--- a/restful.go
+++ b/restful.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2019 Snowflake Computing Inc. All right reserved.
+// Copyright (c) 2017-2021 Snowflake Computing Inc. All right reserved.
 
 package gosnowflake
 
@@ -175,12 +175,16 @@ func postRestfulQueryHelper(
 	if sr.Token != "" {
 		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, sr.Token)
 	}
+
+	var resp *http.Response
 	fullURL := sr.getFullURL(queryRequestPath, params)
-	resp, err := sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, false)
+	resp, err = sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, false)
+
 	if err != nil {
 		return nil, err
 	}
 	defer resp.Body.Close()
+
 	if resp.StatusCode == http.StatusOK {
 		logger.WithContext(ctx).Infof("postQuery: resp: %v", resp)
 		var respd execResponse
@@ -199,7 +203,32 @@ func postRestfulQueryHelper(
 
 		var resultURL string
 		isSessionRenewed := false
+		noResult, _ := isAsyncMode(ctx)
+
+		// if asynchronous query in progress, kick off retrieval but return object
+		if respd.Code == queryInProgressAsyncCode && noResult {
+			// placeholder object to return to user while retrieving results
+			rows := new(snowflakeRows)
+			res := new(snowflakeResult)
+			switch resType := getResultType(ctx); resType {
+			case execResultType:
+				res.queryID = respd.Data.QueryID
+				res.status = QueryStatusInProgress
+				res.errChannel = make(chan error)
+				respd.Data.AsyncResult = res
+			case queryResultType:
+				rows.queryID = respd.Data.QueryID
+				rows.status = QueryStatusInProgress
+				rows.errChannel = make(chan error)
+				respd.Data.AsyncRows = rows
+			default:
+				return &respd, nil
+			}
 
+			// spawn goroutine to retrieve asynchronous results
+			go getAsync(ctx, sr, headers, sr.getFullURL(respd.Data.GetResultURL, nil), timeout, res, rows)
+			return &respd, nil
+		}
 		for isSessionRenewed || respd.Code == queryInProgressCode ||
 			respd.Code == queryInProgressAsyncCode {
 			if !isSessionRenewed {
diff --git a/result.go b/result.go
index 6e271fc..3001056 100644
--- a/result.go
+++ b/result.go
@@ -1,26 +1,71 @@
-// Copyright (c) 2017-2019 Snowflake Computing Inc. All right reserved.
+// Copyright (c) 2017-2021 Snowflake Computing Inc. All right reserved.
 
 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
+	QueryStatusInProgress queryStatus = "queryStatusInProgress"
+	// QueryStatusComplete denotes a completed query execution
+	QueryStatusComplete queryStatus = "queryStatusComplete"
+	// QueryFailed denotes a failed query
+	QueryFailed queryStatus = "queryFailed"
+)
+
 // SnowflakeResult provides the associated query ID
 type SnowflakeResult interface {
-	QueryID() string
+	GetQueryID() string
+	GetStatus() queryStatus
 }
 
 type snowflakeResult struct {
 	affectedRows int64
 	insertID     int64 // Snowflake doesn't support last insert id
 	queryID      string
+	status       queryStatus
+	err          error
+	errChannel   chan error
 }
 
 func (res *snowflakeResult) LastInsertId() (int64, error) {
+	if err := res.waitForAsyncExecStatus(); err != nil {
+		return -1, err
+	}
 	return res.insertID, nil
 }
 
 func (res *snowflakeResult) RowsAffected() (int64, error) {
+	if err := res.waitForAsyncExecStatus(); err != nil {
+		return 0, err
+	}
 	return res.affectedRows, nil
 }
 
-func (res *snowflakeResult) QueryID() string {
+func (res *snowflakeResult) GetQueryID() string {
 	return res.queryID
 }
+
+func (res *snowflakeResult) GetStatus() queryStatus {
+	return res.status
+}
+
+func (res *snowflakeResult) waitForAsyncExecStatus() error {
+	// if async exec, block until execution is finished
+	if res.status == QueryStatusInProgress {
+		err := <-res.errChannel
+		res.status = QueryStatusComplete
+		if err != nil {
+			res.status = QueryFailed
+			res.err = err
+			return err
+		}
+	} else if res.status == QueryFailed {
+		return res.err
+	}
+	return nil
+}
diff --git a/rows.go b/rows.go
index f38e267..a371bb3 100644
--- a/rows.go
+++ b/rows.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2019 Snowflake Computing Inc. All right reserved.
+// Copyright (c) 2017-2021 Snowflake Computing Inc. All right reserved.
 
 package gosnowflake
 
@@ -43,9 +43,15 @@ type snowflakeRows struct {
 	RowType         []execResponseRowType
 	ChunkDownloader *snowflakeChunkDownloader
 	queryID         string
+	status          queryStatus
+	err             error
+	errChannel      chan error
 }
 
 func (rows *snowflakeRows) Close() (err error) {
+	if err := rows.waitForAsyncQueryStatus(); err != nil {
+		return err
+	}
 	logger.WithContext(rows.sc.ctx).Debugln("Rows.Close")
 	return nil
 }
@@ -98,11 +104,17 @@ type snowflakeChunkDownloader struct {
 
 // ColumnTypeDatabaseTypeName returns the database column name.
 func (rows *snowflakeRows) ColumnTypeDatabaseTypeName(index int) string {
+	if err := rows.waitForAsyncQueryStatus(); err != nil {
+		return err.Error()
+	}
 	return strings.ToUpper(rows.RowType[index].Type)
 }
 
 // ColumnTypeLength returns the length of the column
 func (rows *snowflakeRows) ColumnTypeLength(index int) (length int64, ok bool) {
+	if err := rows.waitForAsyncQueryStatus(); err != nil {
+		return 0, false
+	}
 	if index < 0 || index > len(rows.RowType) {
 		return 0, false
 	}
@@ -114,6 +126,9 @@ func (rows *snowflakeRows) ColumnTypeLength(index int) (length int64, ok bool) {
 }
 
 func (rows *snowflakeRows) ColumnTypeNullable(index int) (nullable, ok bool) {
+	if err := rows.waitForAsyncQueryStatus(); err != nil {
+		return false, false
+	}
 	if index < 0 || index > len(rows.RowType) {
 		return false, false
 	}
@@ -121,6 +136,9 @@ func (rows *snowflakeRows) ColumnTypeNullable(index int) (nullable, ok bool) {
 }
 
 func (rows *snowflakeRows) ColumnTypePrecisionScale(index int) (precision, scale int64, ok bool) {
+	if err := rows.waitForAsyncQueryStatus(); err != nil {
+		return 0, 0, false
+	}
 	if index < 0 || index > len(rows.RowType) {
 		return 0, 0, false
 	}
@@ -136,6 +154,9 @@ func (rows *snowflakeRows) ColumnTypePrecisionScale(index int) (precision, scale
 }
 
 func (rows *snowflakeRows) Columns() []string {
+	if err := rows.waitForAsyncQueryStatus(); err != nil {
+		return make([]string, 0)
+	}
 	logger.Debug("Rows.Columns")
 	ret := make([]string, len(rows.RowType))
 	for i, n := 0, len(rows.RowType); i < n; i++ {
@@ -145,14 +166,24 @@ func (rows *snowflakeRows) Columns() []string {
 }
 
 func (rows *snowflakeRows) ColumnTypeScanType(index int) reflect.Type {
+	if err := rows.waitForAsyncQueryStatus(); err != nil {
+		return nil
+	}
 	return snowflakeTypeToGo(rows.RowType[index].Type, rows.RowType[index].Scale)
 }
 
-func (rows *snowflakeRows) QueryID() string {
+func (rows *snowflakeRows) GetQueryID() string {
 	return rows.queryID
 }
 
+func (rows *snowflakeRows) GetStatus() queryStatus {
+	return rows.status
+}
+
 func (rows *snowflakeRows) Next(dest []driver.Value) (err error) {
+	if err := rows.waitForAsyncQueryStatus(); err != nil {
+		return err
+	}
 	row, err := rows.ChunkDownloader.Next()
 	if err != nil {
 		// includes io.EOF
@@ -180,6 +211,9 @@ func (rows *snowflakeRows) Next(dest []driver.Value) (err error) {
 }
 
 func (rows *snowflakeRows) HasNextResultSet() bool {
+	if err := rows.waitForAsyncQueryStatus(); err != nil {
+		return false
+	}
 	if len(rows.ChunkDownloader.ChunkMetas) == 0 && rows.ChunkDownloader.NextDownloader == nil {
 		return false // no extra chunk
 	}
@@ -187,6 +221,9 @@ func (rows *snowflakeRows) HasNextResultSet() bool {
 }
 
 func (rows *snowflakeRows) NextResultSet() error {
+	if err := rows.waitForAsyncQueryStatus(); err != nil {
+		return err
+	}
 	if len(rows.ChunkDownloader.ChunkMetas) == 0 {
 		if rows.ChunkDownloader.NextDownloader == nil {
 			return io.EOF
@@ -197,6 +234,22 @@ func (rows *snowflakeRows) NextResultSet() error {
 	return rows.ChunkDownloader.nextResultSet()
 }
 
+func (rows *snowflakeRows) waitForAsyncQueryStatus() error {
+	// if async query, block until query is finished
+	if rows.status == QueryStatusInProgress {
+		err := <-rows.errChannel
+		rows.status = QueryStatusComplete
+		if err != nil {
+			rows.status = QueryFailed
+			rows.err = err
+			return rows.err
+		}
+	} else if rows.status == QueryFailed {
+		return rows.err
+	}
+	return nil
+}
+
 func (scd *snowflakeChunkDownloader) totalUncompressedSize() (acc int64) {
 	for _, c := range scd.ChunkMetas {
 		acc += c.UncompressedSize
diff --git a/statement.go b/statement.go
index 224ea6c..cce8211 100644
--- a/statement.go
+++ b/statement.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2019 Snowflake Computing Inc. All right reserved.
+// Copyright (c) 2017-2021 Snowflake Computing Inc. All right reserved.
 
 package gosnowflake
 
@@ -12,6 +12,8 @@ type paramKey string
 const (
 	// MultiStatementCount controls the number of queries to execute in a single API call
 	MultiStatementCount paramKey = "MULTI_STATEMENT_COUNT"
+	// AsyncMode tells the server to not block the request on executing the entire query
+	AsyncMode paramKey = "ASYNC_MODE_QUERY"
 )
 
 type snowflakeStmt struct {
@@ -55,3 +57,8 @@ func (stmt *snowflakeStmt) Query(args []driver.Value) (driver.Rows, error) {
 func WithMultiStatement(ctx context.Context, num int) (context.Context, error) {
 	return context.WithValue(ctx, MultiStatementCount, num), nil
 }
+
+// WithAsyncMode returns a context that allows execution of query in async mode
+func WithAsyncMode(ctx context.Context) (context.Context, error) {
+	return context.WithValue(ctx, AsyncMode, true), nil
+}


Test patch:
diff --git a/statement_test.go b/statement_test.go
index 1f18f75..642fdc2 100644
--- a/statement_test.go
+++ b/statement_test.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2020 Snowflake Computing Inc. All right reserved.
+// Copyright (c) 2020-2021 Snowflake Computing Inc. All right reserved.
 
 package gosnowflake
 
@@ -6,6 +6,7 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"fmt"
 	"testing"
 )
 
@@ -465,7 +466,7 @@ func TestGetQueryID(t *testing.T) {
 			return err
 		}
 		defer rows.Close()
-		qid := rows.(SnowflakeResult).QueryID()
+		qid := rows.(SnowflakeResult).GetQueryID()
 		if qid == "" {
 			t.Fatal("should have returned a query ID string")
 		}
@@ -497,3 +498,135 @@ func TestGetQueryID(t *testing.T) {
 		t.Fatalf("failed to prepare statement. err: %v", err)
 	}
 }
+
+func TestAsyncMode(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, _ := WithAsyncMode(context.Background())
+	numrows := 100000
+	rows, _ := db.QueryContext(ctx, fmt.Sprintf("SELECT SEQ8(), RANDSTR(1000, RANDOM()) FROM TABLE(GENERATOR(ROWCOUNT=>%v))", numrows))
+	defer rows.Close()
+
+	cnt := 0
+	var idx int
+	var v string
+	// Next() will block and wait until results are available
+	for rows.Next() {
+		err := rows.Scan(&idx, &v)
+		if err != nil {
+			t.Fatal(err)
+		}
+		cnt++
+	}
+	logger.Infof("NextResultSet: %v", rows.NextResultSet())
+
+	if cnt != numrows {
+		t.Errorf("number of rows didn't match. expected: %v, got: %v", numrows, cnt)
+	}
+
+	db.Exec("create table test (value boolean)")
+	res, _ := db.ExecContext(ctx, "insert into test values (true)")
+	// RowsAffected() will block and wait until results are available
+	count, err := res.RowsAffected()
+	if err != nil || count != 1 {
+		t.Fatalf("count was invalid. err: %v, count: %v", err, count)
+	}
+}
+
+func TestAsyncQueryFail(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, _ := WithAsyncMode(context.Background())
+	rows, err := db.QueryContext(ctx, "selectt 1")
+	if err != nil {
+		t.Fatal("asynchronous query should always return nil error")
+	}
+	defer rows.Close()
+
+	if rows.Next() {
+		t.Fatal("should have no rows available")
+	} else {
+		if err = rows.Err(); err == nil {
+			t.Fatal("should return a syntax error")
+		}
+	}
+
+	conn, _ := db.Conn(ctx)
+	err = conn.Raw(func(x interface{}) error {
+		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("asynchronous query should always return nil error")
+		}
+		defer rows.Close()
+
+		dest := make([]driver.Value, 1)
+		err = rows.Next(dest)
+		if err == nil {
+			t.Fatal("should return a syntax error")
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("failed to query statement. err: %v", err)
+	}
+}
+
+func TestMultipleAsyncQueries(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, _ := WithAsyncMode(context.Background())
+	s1 := "foo"
+	s2 := "bar"
+	rows1, _ := db.QueryContext(ctx, fmt.Sprintf("select distinct '%v' from table (generator(timelimit=>%v))", s1, 30))
+	defer rows1.Close()
+	rows2, _ := db.QueryContext(ctx, fmt.Sprintf("select distinct '%v' from table (generator(timelimit=>%v))", s2, 10))
+	defer rows2.Close()
+
+	ch1 := make(chan string)
+	ch2 := make(chan string)
+	go retrieveRows(rows1, ch1)
+	go retrieveRows(rows2, ch2)
+	select {
+	case res := <-ch1:
+		t.Fatalf("value %v should not have been called earlier.", res)
+	case res := <-ch2:
+		if res != s2 {
+			t.Fatalf("query failed. expected: %v, got: %v", s2, res)
+		}
+	}
+}
+
+func retrieveRows(rows *sql.Rows, ch chan string) {
+	var s string
+	for rows.Next() {
+		if err := rows.Scan(&s); err != nil {
+			ch <- err.Error()
+			close(ch)
+			return
+		}
+	}
+	ch <- s
+	close(ch)
+}