Repository: snowflakedb/gosnowflake
Pull request #372 (2021-01-28 05:16:51)

SNOW-209084 Stream Based Chunk Downloader
### Description
Replace traditional worker based downloader with stream based downloader to be used for JSON format queries. The worker based downloader is still the default for Arrow out of necessity.

Moved a bulk of the chunk downloading/fetching logic as well as the structs to a separate file for modularity. Added multiple unit tests involving downloading and fetching using the stream based downloader. Added integration test enabling stream based downloader. Accessible via context returned from `WithStreamDownloader`.

### 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



Patch:
diff --git a/chunk_downloader.go b/chunk_downloader.go
new file mode 100644
index 0000000..ab88feb
--- /dev/null
+++ b/chunk_downloader.go
@@ -0,0 +1,617 @@
+// Copyright (c) 2021 Snowflake Computing Inc. All right reserved.
+
+package gosnowflake
+
+import (
+	"bufio"
+	"compress/gzip"
+	"context"
+	"encoding/json"
+	"fmt"
+	"github.com/apache/arrow/go/arrow/ipc"
+	"github.com/apache/arrow/go/arrow/memory"
+	"io"
+	"io/ioutil"
+	"math/rand"
+	"net"
+	"net/http"
+	"net/url"
+	"strings"
+	"sync"
+	"time"
+)
+
+type chunkDownloader interface {
+	totalUncompressedSize() (acc int64)
+	hasNextResultSet() bool
+	nextResultSet() error
+	start() error
+	next() (chunkRowType, error)
+	reset()
+	getChunkMetas() []execResponseChunk
+	getQueryResultFormat() resultFormat
+	setNextChunkDownloader(downloader chunkDownloader)
+	getNextChunkDownloader() chunkDownloader
+}
+
+type snowflakeChunkDownloader struct {
+	sc                 *snowflakeConn
+	ctx                context.Context
+	Total              int64
+	TotalRowIndex      int64
+	CellCount          int
+	CurrentChunk       []chunkRowType
+	CurrentChunkIndex  int
+	CurrentChunkSize   int
+	CurrentIndex       int
+	ChunkHeader        map[string]string
+	ChunkMetas         []execResponseChunk
+	Chunks             map[int][]chunkRowType
+	ChunksChan         chan int
+	ChunksError        chan *chunkError
+	ChunksErrorCounter int
+	ChunksFinalErrors  []*chunkError
+	ChunksMutex        *sync.Mutex
+	DoneDownloadCond   *sync.Cond
+	NextDownloader     chunkDownloader
+	Qrmk               string
+	QueryResultFormat  string
+	RowSet             rowSetType
+	FuncDownload       func(context.Context, *snowflakeChunkDownloader, int)
+	FuncDownloadHelper func(context.Context, *snowflakeChunkDownloader, int) error
+	FuncGet            func(context.Context, *snowflakeChunkDownloader, string, map[string]string, time.Duration) (*http.Response, error)
+}
+
+func (scd *snowflakeChunkDownloader) totalUncompressedSize() (acc int64) {
+	for _, c := range scd.ChunkMetas {
+		acc += c.UncompressedSize
+	}
+	return
+}
+
+func (scd *snowflakeChunkDownloader) hasNextResultSet() bool {
+	if len(scd.ChunkMetas) == 0 && scd.NextDownloader == nil {
+		return false // no extra chunk
+	}
+	// next result set exists if current chunk has remaining result sets or there is another downloader
+	return scd.CurrentChunkIndex < len(scd.ChunkMetas) || scd.NextDownloader != nil
+}
+
+func (scd *snowflakeChunkDownloader) nextResultSet() error {
+	// no error at all times as the next chunk/resultset is automatically read
+	if scd.CurrentChunkIndex < len(scd.ChunkMetas) {
+		return nil
+	}
+	return io.EOF
+}
+
+func (scd *snowflakeChunkDownloader) start() error {
+	scd.CurrentChunkSize = len(scd.RowSet.JSON) // cache the size
+	scd.CurrentIndex = -1                       // initial chunks idx
+	scd.CurrentChunkIndex = -1                  // initial chunk
+
+	scd.CurrentChunk = make([]chunkRowType, scd.CurrentChunkSize)
+	populateJSONRowSet(scd.CurrentChunk, scd.RowSet.JSON)
+
+	if scd.getQueryResultFormat() == arrowFormat && scd.RowSet.RowSetBase64 != "" {
+		// if the rowsetbase64 retrieved from the server is empty, move on to downloading chunks
+		var err error
+		firstArrowChunk := buildFirstArrowChunk(scd.RowSet.RowSetBase64)
+		scd.CurrentChunk, err = firstArrowChunk.decodeArrowChunk(scd.RowSet.RowType)
+		scd.CurrentChunkSize = firstArrowChunk.rowCount
+		if err != nil {
+			return err
+		}
+	}
+
+	// start downloading chunks if exists
+	chunkMetaLen := len(scd.ChunkMetas)
+	if chunkMetaLen > 0 {
+		logger.Debugf("MaxChunkDownloadWorkers: %v", MaxChunkDownloadWorkers)
+		logger.Debugf("chunks: %v, total bytes: %d", chunkMetaLen, scd.totalUncompressedSize())
+		scd.ChunksMutex = &sync.Mutex{}
+		scd.DoneDownloadCond = sync.NewCond(scd.ChunksMutex)
+		scd.Chunks = make(map[int][]chunkRowType)
+		scd.ChunksChan = make(chan int, chunkMetaLen)
+		scd.ChunksError = make(chan *chunkError, MaxChunkDownloadWorkers)
+		for i := 0; i < chunkMetaLen; i++ {
+			logger.Debugf("add chunk to channel ChunksChan: %v", i+1)
+			scd.ChunksChan <- i
+		}
+		for i := 0; i < intMin(MaxChunkDownloadWorkers, chunkMetaLen); i++ {
+			scd.schedule()
+		}
+	}
+	return nil
+}
+
+func (scd *snowflakeChunkDownloader) schedule() {
+	select {
+	case nextIdx := <-scd.ChunksChan:
+		logger.Infof("schedule chunk: %v", nextIdx+1)
+		go scd.FuncDownload(scd.ctx, scd, nextIdx)
+	default:
+		// no more download
+		logger.Info("no more download")
+	}
+}
+
+func (scd *snowflakeChunkDownloader) checkErrorRetry() (err error) {
+	select {
+	case errc := <-scd.ChunksError:
+		if scd.ChunksErrorCounter < maxChunkDownloaderErrorCounter && errc.Error != context.Canceled {
+			// add the index to the chunks channel so that the download will be retried.
+			go scd.FuncDownload(scd.ctx, scd, errc.Index)
+			scd.ChunksErrorCounter++
+			logger.Warningf("chunk idx: %v, err: %v. retrying (%v/%v)...",
+				errc.Index, errc.Error, scd.ChunksErrorCounter, maxChunkDownloaderErrorCounter)
+		} else {
+			scd.ChunksFinalErrors = append(scd.ChunksFinalErrors, errc)
+			logger.Warningf("chunk idx: %v, err: %v. no further retry", errc.Index, errc.Error)
+			return errc.Error
+		}
+	default:
+		logger.Info("no error is detected.")
+	}
+	return nil
+}
+
+func (scd *snowflakeChunkDownloader) next() (chunkRowType, error) {
+	for {
+		scd.CurrentIndex++
+		if scd.CurrentIndex < scd.CurrentChunkSize {
+			return scd.CurrentChunk[scd.CurrentIndex], nil
+		}
+		scd.CurrentChunkIndex++ // next chunk
+		scd.CurrentIndex = -1   // reset
+		if scd.CurrentChunkIndex >= len(scd.ChunkMetas) {
+			break
+		}
+
+		scd.ChunksMutex.Lock()
+		if scd.CurrentChunkIndex > 1 {
+			scd.Chunks[scd.CurrentChunkIndex-1] = nil // detach the previously used chunk
+		}
+
+		for scd.Chunks[scd.CurrentChunkIndex] == nil {
+			logger.Debugf("waiting for chunk idx: %v/%v",
+				scd.CurrentChunkIndex+1, len(scd.ChunkMetas))
+
+			err := scd.checkErrorRetry()
+			if err != nil {
+				scd.ChunksMutex.Unlock()
+				return chunkRowType{}, err
+			}
+
+			// wait for chunk downloader goroutine to broadcast the event,
+			// 1) one chunk download finishes or 2) an error occurs.
+			scd.DoneDownloadCond.Wait()
+		}
+		logger.Debugf("ready: chunk %v", scd.CurrentChunkIndex+1)
+		scd.CurrentChunk = scd.Chunks[scd.CurrentChunkIndex]
+		scd.ChunksMutex.Unlock()
+		scd.CurrentChunkSize = len(scd.CurrentChunk)
+
+		// kick off the next download
+		scd.schedule()
+	}
+
+	logger.Debugf("no more data")
+	if len(scd.ChunkMetas) > 0 {
+		close(scd.ChunksError)
+		close(scd.ChunksChan)
+	}
+	return chunkRowType{}, io.EOF
+}
+
+func (scd *snowflakeChunkDownloader) reset() {
+	scd.Chunks = nil // detach all chunks. No way to go backward without reinitialize it.
+}
+
+func (scd *snowflakeChunkDownloader) getChunkMetas() []execResponseChunk {
+	return scd.ChunkMetas
+}
+
+func (scd *snowflakeChunkDownloader) getQueryResultFormat() resultFormat {
+	if scd.QueryResultFormat == "arrow" {
+		return arrowFormat
+	}
+	return jsonFormat
+}
+
+func (scd *snowflakeChunkDownloader) setNextChunkDownloader(nextDownloader chunkDownloader) {
+	scd.NextDownloader = nextDownloader
+}
+
+func (scd *snowflakeChunkDownloader) getNextChunkDownloader() chunkDownloader {
+	return scd.NextDownloader
+}
+
+func getChunk(
+	ctx context.Context,
+	scd *snowflakeChunkDownloader,
+	fullURL string,
+	headers map[string]string,
+	timeout time.Duration) (
+	*http.Response, error) {
+	u, err := url.Parse(fullURL)
+	if err != nil {
+		return nil, err
+	}
+	return newRetryHTTP(ctx, scd.sc.rest.Client, http.NewRequest, u, headers, timeout).execute()
+}
+
+/* largeResultSetReader is a reader that wraps the large result set with leading and tailing brackets. */
+type largeResultSetReader struct {
+	status int
+	body   io.Reader
+}
+
+func (r *largeResultSetReader) Read(p []byte) (n int, err error) {
+	if r.status == 0 {
+		p[0] = 0x5b // initial 0x5b ([)
+		r.status = 1
+		return 1, nil
+	}
+	if r.status == 1 {
+		var len int
+		len, err = r.body.Read(p)
+		if err == io.EOF {
+			r.status = 2
+			return len, nil
+		}
+		if err != nil {
+			return 0, err
+		}
+		return len, nil
+	}
+	if r.status == 2 {
+		p[0] = 0x5d // tail 0x5d (])
+		r.status = 3
+		return 1, nil
+	}
+	// ensure no data and EOF
+	return 0, io.EOF
+}
+
+func downloadChunk(ctx context.Context, scd *snowflakeChunkDownloader, idx int) {
+	logger.Infof("download start chunk: %v", idx+1)
+	defer scd.DoneDownloadCond.Broadcast()
+
+	if err := scd.FuncDownloadHelper(ctx, scd, idx); err != nil {
+		logger.Errorf(
+			"failed to extract HTTP response body. URL: %v, err: %v", scd.ChunkMetas[idx].URL, err)
+		scd.ChunksError <- &chunkError{Index: idx, Error: err}
+	} else if scd.ctx.Err() == context.Canceled || scd.ctx.Err() == context.DeadlineExceeded {
+		scd.ChunksError <- &chunkError{Index: idx, Error: scd.ctx.Err()}
+	}
+}
+
+func downloadChunkHelper(ctx context.Context, scd *snowflakeChunkDownloader, idx int) error {
+	headers := make(map[string]string)
+	if len(scd.ChunkHeader) > 0 {
+		logger.Debug("chunk header is provided.")
+		for k, v := range scd.ChunkHeader {
+			headers[k] = v
+		}
+	} else {
+		headers[headerSseCAlgorithm] = headerSseCAes
+		headers[headerSseCKey] = scd.Qrmk
+	}
+
+	resp, err := scd.FuncGet(ctx, scd, scd.ChunkMetas[idx].URL, headers, scd.sc.rest.RequestTimeout)
+	if err != nil {
+		return err
+	}
+	bufStream := bufio.NewReader(resp.Body)
+	defer resp.Body.Close()
+	logger.Infof("response returned chunk: %v, resp: %v", idx+1, resp)
+	if resp.StatusCode != http.StatusOK {
+		b, err := ioutil.ReadAll(bufStream)
+		if err != nil {
+			return err
+		}
+		logger.Infof("HTTP: %v, URL: %v, Body: %v", resp.StatusCode, scd.ChunkMetas[idx].URL, b)
+		logger.Infof("Header: %v", resp.Header)
+		return &SnowflakeError{
+			Number:      ErrFailedToGetChunk,
+			SQLState:    SQLStateConnectionFailure,
+			Message:     errMsgFailedToGetChunk,
+			MessageArgs: []interface{}{idx},
+		}
+	}
+	return decodeChunk(scd, idx, bufStream)
+}
+
+func decodeChunk(scd *snowflakeChunkDownloader, idx int, bufStream *bufio.Reader) (err error) {
+	gzipMagic, err := bufStream.Peek(2)
+	if err != nil {
+		return err
+	}
+	start := time.Now()
+	var source io.Reader
+	if gzipMagic[0] == 0x1f && gzipMagic[1] == 0x8b {
+		// detects and uncompresses Gzip format data
+		bufStream0, err := gzip.NewReader(bufStream)
+		if err != nil {
+			return err
+		}
+		defer bufStream0.Close()
+		source = bufStream0
+	} else {
+		source = bufStream
+	}
+	st := &largeResultSetReader{
+		status: 0,
+		body:   source,
+	}
+	var respd []chunkRowType
+	if scd.getQueryResultFormat() != arrowFormat {
+		var decRespd [][]*string
+		if !CustomJSONDecoderEnabled {
+			dec := json.NewDecoder(st)
+			for {
+				if err := dec.Decode(&decRespd); err == io.EOF {
+					break
+				} else if err != nil {
+					return err
+				}
+			}
+		} else {
+			decRespd, err = decodeLargeChunk(st, scd.ChunkMetas[idx].RowCount, scd.CellCount)
+			if err != nil {
+				return err
+			}
+		}
+		respd = make([]chunkRowType, len(decRespd))
+		populateJSONRowSet(respd, decRespd)
+	} else {
+		ipcReader, err := ipc.NewReader(source)
+		if err != nil {
+			return err
+		}
+		arc := arrowResultChunk{
+			*ipcReader,
+			0,
+			int(scd.totalUncompressedSize()),
+			memory.NewGoAllocator(),
+		}
+		respd, err = arc.decodeArrowChunk(scd.RowSet.RowType)
+		if err != nil {
+			return err
+		}
+	}
+	logger.Debugf(
+		"decoded %d rows w/ %d bytes in %s (chunk %v)",
+		scd.ChunkMetas[idx].RowCount,
+		scd.ChunkMetas[idx].UncompressedSize,
+		time.Since(start), idx+1,
+	)
+
+	scd.ChunksMutex.Lock()
+	defer scd.ChunksMutex.Unlock()
+	scd.Chunks[idx] = respd
+	return nil
+}
+
+func populateJSONRowSet(dst []chunkRowType, src [][]*string) {
+	// populate string rowset from src to dst's chunkRowType struct's RowSet field
+	for i, row := range src {
+		dst[i].RowSet = row
+	}
+}
+
+type streamChunkDownloader struct {
+	ctx            context.Context
+	id             int64
+	fetcher        streamChunkFetcher
+	readErr        error
+	rowStream      chan []*string
+	Total          int64
+	ChunkMetas     []execResponseChunk
+	NextDownloader chunkDownloader
+	RowSet         rowSetType
+}
+
+func (scd *streamChunkDownloader) totalUncompressedSize() (acc int64) {
+	return -1
+}
+
+func (scd *streamChunkDownloader) hasNextResultSet() bool {
+	return scd.readErr == nil
+}
+
+func (scd *streamChunkDownloader) nextResultSet() error {
+	return scd.readErr
+}
+
+func (scd *streamChunkDownloader) start() error {
+	go func() {
+		var readErr = io.EOF
+
+		logger.WithContext(scd.ctx).Infof(
+			"start downloading. downloader id: %v, %v/%v rows, %v chunks",
+			scd.id, len(scd.RowSet.RowType), scd.Total, len(scd.ChunkMetas))
+		t := time.Now()
+
+		defer func() {
+			if readErr == io.EOF {
+				logger.WithContext(scd.ctx).Infof("downloading done. downloader id: %v", scd.id)
+			} else {
+				logger.WithContext(scd.ctx).Debugf("downloading error. downloader id: %v", scd.id)
+			}
+			scd.readErr = readErr
+			close(scd.rowStream)
+
+			if r := recover(); r != nil {
+				if err, ok := r.(error); ok {
+					readErr = err
+				} else {
+					readErr = fmt.Errorf("%v", r)
+				}
+			}
+		}()
+
+		logger.WithContext(scd.ctx).Infof("sending initial set of rows in %vms", time.Since(t).Microseconds())
+		t = time.Now()
+		for _, row := range scd.RowSet.JSON {
+			scd.rowStream <- row
+		}
+		scd.RowSet.JSON = nil
+
+		// Download and parse one chunk at a time. The fetcher will send each
+		// parsed row to the row stream. When an error occurs, the fetcher will
+		// stop writing to the row stream so we can stop processing immediately
+		for i, chunk := range scd.ChunkMetas {
+			logger.WithContext(scd.ctx).Infof("starting chunk fetch %d (%d rows)", i, chunk.RowCount)
+			if err := scd.fetcher.fetch(chunk.URL, scd.rowStream); err != nil {
+				logger.WithContext(scd.ctx).Infof(
+					"failed chunk fetch %d: %#v, downloader id: %v, %v/%v rows, %v chunks",
+					i, err, scd.id, len(scd.RowSet.RowType), scd.Total, len(scd.ChunkMetas))
+				readErr = fmt.Errorf("chunk fetch: %w", err)
+				break
+			}
+			logger.WithContext(scd.ctx).Infof("fetched chunk %d (%d rows) in %vms", i, chunk.RowCount, time.Since(t).Microseconds())
+			t = time.Now()
+		}
+	}()
+	return nil
+}
+
+func (scd *streamChunkDownloader) next() (chunkRowType, error) {
+	if row, ok := <-scd.rowStream; ok {
+		return chunkRowType{RowSet: row}, nil
+	}
+	return chunkRowType{}, scd.readErr
+}
+
+func (scd *streamChunkDownloader) reset() {}
+
+func (scd *streamChunkDownloader) getChunkMetas() []execResponseChunk {
+	return scd.ChunkMetas
+}
+
+func (scd *streamChunkDownloader) getQueryResultFormat() resultFormat {
+	return jsonFormat
+}
+
+func (scd *streamChunkDownloader) setNextChunkDownloader(nextDownloader chunkDownloader) {
+	scd.NextDownloader = nextDownloader
+}
+
+func (scd *streamChunkDownloader) getNextChunkDownloader() chunkDownloader {
+	return scd.NextDownloader
+}
+
+func useStreamDownloader(ctx context.Context) bool {
+	val := ctx.Value(streamChunkDownload)
+	if val == nil {
+		return false
+	}
+	boolVal, ok := val.(bool)
+	if !ok {
+		return false
+	}
+	return boolVal
+}
+
+type streamChunkFetcher interface {
+	fetch(url string, rows chan<- []*string) error
+}
+
+type httpStreamChunkFetcher struct {
+	ctx      context.Context
+	client   *http.Client
+	clientIP net.IP
+	headers  map[string]string
+	qrmk     string
+}
+
+func newStreamChunkDownloader(
+	ctx context.Context,
+	fetcher streamChunkFetcher,
+	total int64,
+	firstRows [][]*string,
+	chunks []execResponseChunk) *streamChunkDownloader {
+	return &streamChunkDownloader{
+		ctx:        ctx,
+		id:         rand.Int63(),
+		fetcher:    fetcher,
+		RowSet:     rowSetType{JSON: firstRows},
+		ChunkMetas: chunks,
+		readErr:    nil,
+		Total:      total,
+		rowStream:  make(chan []*string),
+	}
+}
+
+func (f *httpStreamChunkFetcher) fetch(URL string, rows chan<- []*string) error {
+	headers := f.headers
+	if len(headers) == 0 {
+		headers = map[string]string{
+			headerSseCAlgorithm: headerSseCAes,
+			headerSseCKey:       f.qrmk,
+		}
+	}
+
+	fullURL, _ := url.Parse(URL)
+	res, err := newRetryHTTP(context.Background(), f.client, http.NewRequest, fullURL, f.headers, 0).execute()
+	if err != nil {
+		return err
+	}
+	defer res.Body.Close()
+	if res.StatusCode != http.StatusOK {
+		b, _ := ioutil.ReadAll(res.Body)
+		return fmt.Errorf("status (%d): %s", res.StatusCode, string(b))
+	}
+	if err := copyChunkStream(res.Body, rows); err != nil {
+		return fmt.Errorf("read: %w", err)
+	}
+	return nil
+}
+
+func copyChunkStream(body io.Reader, rows chan<- []*string) error {
+	bufStream := bufio.NewReader(body)
+	gzipMagic, err := bufStream.Peek(2)
+	if err != nil {
+		return err
+	}
+	var source io.Reader
+	if gzipMagic[0] == 0x1f && gzipMagic[1] == 0x8b {
+		// detect and decompress Gzip format data
+		bufStream0, err := gzip.NewReader(bufStream)
+		if err != nil {
+			return err
+		}
+		defer bufStream0.Close()
+		source = bufStream0
+	} else {
+		source = bufStream
+	}
+	r := io.MultiReader(strings.NewReader("["), source, strings.NewReader("]"))
+	dec := json.NewDecoder(r)
+	openToken := json.Delim('[')
+	closeToken := json.Delim(']')
+	for {
+		if t, err := dec.Token(); err == io.EOF {
+			break
+		} else if err != nil {
+			return fmt.Errorf("delim open: %w", err)
+		} else if t != openToken {
+			return fmt.Errorf("delim open: got %T", t)
+		}
+		for dec.More() {
+			var row []*string
+			if err := dec.Decode(&row); err != nil {
+				return fmt.Errorf("decode: %w", err)
+			}
+			rows <- row
+		}
+		if t, err := dec.Token(); err != nil {
+			return fmt.Errorf("delim close: %w", err)
+		} else if t != closeToken {
+			return fmt.Errorf("delim close: got %T", t)
+		}
+	}
+	return nil
+}
diff --git a/connection.go b/connection.go
index e8e93ac..58b50d1 100644
--- a/connection.go
+++ b/connection.go
@@ -135,7 +135,7 @@ func (sc *snowflakeConn) exec(
 	var data *execResponse
 
 	requestID := getOrGenerateRequestIDFromContext(ctx)
-	data, err = sc.rest.FuncPostQuery(ctx, sc.rest, &url.Values{}, headers, jsonBody, sc.rest.RequestTimeout, requestID)
+	data, err = sc.rest.FuncPostQuery(ctx, sc.rest, &url.Values{}, headers, jsonBody, sc.rest.RequestTimeout, requestID, sc.cfg)
 	if err != nil {
 		return data, err
 	}
@@ -282,10 +282,9 @@ func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []d
 		return data.Data.AsyncResult, nil
 	}
 
-	var updatedRows int64
 	if sc.isDml(data.Data.StatementTypeID) {
 		// collects all values from the returned row sets
-		updatedRows, err = updateRows(data.Data)
+		updatedRows, err := updateRows(data.Data)
 		if err != nil {
 			return nil, err
 		}
@@ -296,51 +295,7 @@ func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []d
 			queryID:      sc.QueryID,
 		}, nil // last insert id is not supported by Snowflake
 	} else if sc.isMultiStmt(data.Data) {
-		childResults := getChildResults(data.Data.ResultIDs, data.Data.ResultTypes)
-		for _, child := range childResults {
-			resultPath := fmt.Sprintf("/queries/%s/result", child.id)
-			childData, err := sc.getQueryResult(ctx, resultPath)
-			if err != nil {
-				logger.Errorf("error: %v", err)
-				code, err := strconv.Atoi(childData.Code)
-				if err != nil {
-					return nil, err
-				}
-				if childData != nil {
-					return nil, &SnowflakeError{
-						Number:   code,
-						SQLState: childData.Data.SQLState,
-						Message:  err.Error(),
-						QueryID:  childData.Data.QueryID}
-				}
-				return nil, err
-			}
-			if sc.isDml(childData.Data.StatementTypeID) {
-				count, err := updateRows(childData.Data)
-				if err != nil {
-					logger.WithContext(ctx).Errorf("error: %v", err)
-					if childData != nil {
-						code, err := strconv.Atoi(childData.Code)
-						if err != nil {
-							return nil, err
-						}
-						return nil, &SnowflakeError{
-							Number:   code,
-							SQLState: childData.Data.SQLState,
-							Message:  err.Error(),
-							QueryID:  childData.Data.QueryID}
-					}
-					return nil, err
-				}
-				updatedRows += count
-			}
-		}
-		logger.WithContext(ctx).Infof("number of updated rows: %#v", updatedRows)
-		return &snowflakeResult{
-			affectedRows: updatedRows,
-			insertID:     -1,
-			queryID:      sc.QueryID,
-		}, nil
+		return sc.handleMultiExec(ctx, data.Data)
 	}
 	logger.Debug("DDL")
 	return driver.ResultNoRows, nil
@@ -383,59 +338,13 @@ func (sc *snowflakeConn) QueryContext(ctx context.Context, query string, args []
 	rows := new(snowflakeRows)
 	rows.sc = sc
 	rows.RowType = data.Data.RowType
-	rows.ChunkDownloader = &snowflakeChunkDownloader{
-		sc:                 sc,
-		ctx:                ctx,
-		CurrentChunk:       make([]chunkRowType, len(data.Data.RowSet)),
-		ChunkMetas:         data.Data.Chunks,
-		Total:              data.Data.Total,
-		TotalRowIndex:      int64(-1),
-		CellCount:          len(data.Data.RowType),
-		Qrmk:               data.Data.Qrmk,
-		QueryResultFormat:  data.Data.QueryResultFormat,
-		ChunkHeader:        data.Data.ChunkHeaders,
-		FuncDownload:       downloadChunk,
-		FuncDownloadHelper: downloadChunkHelper,
-		FuncGet:            getChunk,
-		RowSet: rowSetType{RowType: data.Data.RowType,
-			JSON:         data.Data.RowSet,
-			RowSetBase64: data.Data.RowSetBase64,
-		},
-	}
+	rows.ChunkDownloader = populateChunkDownloader(ctx, sc, data.Data)
 	rows.queryID = sc.QueryID
 
 	if sc.isMultiStmt(data.Data) {
-		childResults := getChildResults(data.Data.ResultIDs, data.Data.ResultTypes)
-		var nextChunkDownloader *snowflakeChunkDownloader
-		firstResultSet := false
-
-		for _, child := range childResults {
-			resultPath := fmt.Sprintf("/queries/%s/result", child.id)
-			childData, err := sc.getQueryResult(ctx, resultPath)
-			if err != nil {
-				logger.WithContext(ctx).Errorf("error: %v", err)
-				if childData != nil {
-					code, err := strconv.Atoi(childData.Code)
-					if err != nil {
-						return nil, err
-					}
-					return nil, &SnowflakeError{
-						Number:   code,
-						SQLState: childData.Data.SQLState,
-						Message:  err.Error(),
-						QueryID:  childData.Data.QueryID}
-				}
-				return nil, err
-			}
-			if !firstResultSet {
-				// populate rows.ChunkDownloader with the first child
-				rows.ChunkDownloader = populateChunkDownloader(ctx, sc, childData.Data)
-				nextChunkDownloader = rows.ChunkDownloader
-				firstResultSet = true
-			} else {
-				nextChunkDownloader.NextDownloader = populateChunkDownloader(ctx, sc, childData.Data)
-				nextChunkDownloader = nextChunkDownloader.NextDownloader
-			}
+		err := sc.handleMultiQuery(ctx, data.Data, rows)
+		if err != nil {
+			return nil, err
 		}
 	}
 
@@ -538,6 +447,91 @@ func (sc *snowflakeConn) stopHeartBeat() {
 	sc.rest.HeartBeat.stop()
 }
 
+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 {
+		resultPath := fmt.Sprintf("/queries/%s/result", child.id)
+		childData, err := sc.getQueryResult(ctx, resultPath)
+		if err != nil {
+			logger.Errorf("error: %v", err)
+			code, err := strconv.Atoi(childData.Code)
+			if err != nil {
+				return nil, err
+			}
+			if childData != nil {
+				return nil, &SnowflakeError{
+					Number:   code,
+					SQLState: childData.Data.SQLState,
+					Message:  err.Error(),
+					QueryID:  childData.Data.QueryID}
+			}
+			return nil, err
+		}
+		if sc.isDml(childData.Data.StatementTypeID) {
+			count, err := updateRows(childData.Data)
+			if err != nil {
+				logger.WithContext(ctx).Errorf("error: %v", err)
+				if childData != nil {
+					code, err := strconv.Atoi(childData.Code)
+					if err != nil {
+						return nil, err
+					}
+					return nil, &SnowflakeError{
+						Number:   code,
+						SQLState: childData.Data.SQLState,
+						Message:  err.Error(),
+						QueryID:  childData.Data.QueryID}
+				}
+				return nil, err
+			}
+			updatedRows += count
+		}
+	}
+	logger.WithContext(ctx).Infof("number of updated rows: %#v", updatedRows)
+	return &snowflakeResult{
+		affectedRows: updatedRows,
+		insertID:     -1,
+		queryID:      sc.QueryID,
+	}, nil
+}
+
+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)
+		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
+}
+
 func setResultType(ctx context.Context, resType resultType) context.Context {
 	return context.WithValue(ctx, snowflakeResultType, resType)
 }
@@ -607,7 +601,7 @@ func (sc *snowflakeConn) getQueryResult(ctx context.Context, resultPath string)
 }
 
 func isAsyncMode(ctx context.Context) (bool, error) {
-	val := ctx.Value(AsyncMode)
+	val := ctx.Value(asyncMode)
 	if val == nil {
 		return false, nil
 	}
@@ -626,6 +620,7 @@ func getAsync(
 	timeout time.Duration,
 	res *snowflakeResult,
 	rows *snowflakeRows,
+	cfg *Config,
 ) {
 	resType := getResultType(ctx)
 	var errChannel chan error
@@ -664,33 +659,40 @@ func getAsync(
 		return
 	}
 
+	sc := &snowflakeConn{rest: sr, cfg: cfg}
 	if respd.Success {
 		if resType == execResultType {
-			res.affectedRows, _ = updateRows(respd.Data)
 			res.insertID = -1
+			if sc.isDml(respd.Data.StatementTypeID) {
+				res.affectedRows, _ = updateRows(respd.Data)
+			} else if sc.isMultiStmt(respd.Data) {
+				r, err := sc.handleMultiExec(ctx, respd.Data)
+				if err != nil {
+					res.errChannel <- err
+					close(errChannel)
+					return
+				}
+				res.affectedRows, err = r.RowsAffected()
+				if err != nil {
+					res.errChannel <- err
+					close(errChannel)
+					return
+				}
+			}
+			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.ChunkDownloader = populateChunkDownloader(ctx, sc, respd.Data)
+			rows.queryID = respd.Data.QueryID
+			if sc.isMultiStmt(respd.Data) {
+				err = sc.handleMultiQuery(ctx, respd.Data, rows)
+				if err != nil {
+					rows.errChannel <- err
+					close(errChannel)
+					return
+				}
 			}
 			rows.ChunkDownloader.start()
 			rows.errChannel <- nil // mark query status complete
@@ -723,7 +725,20 @@ func getQueryIDChan(ctx context.Context) chan<- string {
 	return c
 }
 
-func populateChunkDownloader(ctx context.Context, sc *snowflakeConn, data execResponseData) *snowflakeChunkDownloader {
+// 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) {
+		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.RowSet, data.Chunks)
+	}
+
 	return &snowflakeChunkDownloader{
 		sc:                 sc,
 		ctx:                ctx,
@@ -738,7 +753,8 @@ func populateChunkDownloader(ctx context.Context, sc *snowflakeConn, data execRe
 		FuncDownload:       downloadChunk,
 		FuncDownloadHelper: downloadChunkHelper,
 		FuncGet:            getChunk,
-		RowSet: rowSetType{RowType: data.RowType,
+		RowSet: rowSetType{
+			RowType:      data.RowType,
 			JSON:         data.RowSet,
 			RowSetBase64: data.RowSetBase64,
 		},
diff --git a/dsn.go b/dsn.go
index 56c71c3..3084c57 100644
--- a/dsn.go
+++ b/dsn.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
 
@@ -6,6 +6,7 @@ import (
 	"crypto/rsa"
 	"encoding/base64"
 	"fmt"
+	"net"
 	"net/http"
 	"net/url"
 	"strconv"
@@ -49,6 +50,7 @@ type Config struct {
 
 	Params map[string]*string // other connection parameters
 
+	ClientIP net.IP // IP address for network check
 	Protocol string // http or https (optional)
 	Host     string // hostname (optional)
 	Port     int    // port (optional)
diff --git a/query.go b/query.go
index 76f1f07..ac9d6eb 100644
--- a/query.go
+++ b/query.go
@@ -6,7 +6,12 @@ import (
 	"time"
 )
 
-const arrowFormat = "arrow"
+type resultFormat string
+
+const (
+	jsonFormat  resultFormat = "json"
+	arrowFormat resultFormat = "arrow"
+)
 
 type execBindParameter struct {
 	Type  string      `json:"type"`
diff --git a/restful.go b/restful.go
index 4e32941..a8efbb1 100644
--- a/restful.go
+++ b/restful.go
@@ -56,8 +56,8 @@ type snowflakeRestful struct {
 	HeartBeat   *heartbeat
 
 	Connection          *snowflakeConn
-	FuncPostQuery       func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, uuid.UUID) (*execResponse, error)
-	FuncPostQueryHelper func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, uuid.UUID) (*execResponse, error)
+	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)
 	FuncRenewSession    func(context.Context, *snowflakeRestful, time.Duration) error
@@ -142,10 +142,11 @@ func postRestfulQuery(
 	headers map[string]string,
 	body []byte,
 	timeout time.Duration,
-	requestID uuid.UUID) (
+	requestID uuid.UUID,
+	cfg *Config) (
 	data *execResponse, err error) {
 
-	data, err = sr.FuncPostQueryHelper(ctx, sr, params, headers, body, timeout, requestID)
+	data, err = sr.FuncPostQueryHelper(ctx, sr, params, headers, body, timeout, requestID, cfg)
 
 	// errors other than context timeout and cancel would be returned to upper layers
 	if err != context.Canceled && err != context.DeadlineExceeded {
@@ -166,7 +167,8 @@ func postRestfulQueryHelper(
 	headers map[string]string,
 	body []byte,
 	timeout time.Duration,
-	requestID uuid.UUID) (
+	requestID uuid.UUID,
+	cfg *Config) (
 	data *execResponse, err error) {
 	logger.Infof("params: %v", params)
 	params.Add(requestIDKey, requestID.String())
@@ -198,7 +200,7 @@ func postRestfulQueryHelper(
 			if err != nil {
 				return nil, err
 			}
-			return sr.FuncPostQuery(ctx, sr, params, headers, body, timeout, requestID)
+			return sr.FuncPostQuery(ctx, sr, params, headers, body, timeout, requestID, cfg)
 		}
 
 		if queryIDChan := getQueryIDChan(ctx); queryIDChan != nil {
@@ -232,7 +234,7 @@ func postRestfulQueryHelper(
 			}
 
 			// spawn goroutine to retrieve asynchronous results
-			go getAsync(ctx, sr, headers, sr.getFullURL(respd.Data.GetResultURL, nil), timeout, res, rows)
+			go getAsync(ctx, sr, headers, sr.getFullURL(respd.Data.GetResultURL, nil), timeout, res, rows, cfg)
 			return &respd, nil
 		}
 		for isSessionRenewed || respd.Code == queryInProgressCode ||
diff --git a/rows.go b/rows.go
index a371bb3..8a06735 100644
--- a/rows.go
+++ b/rows.go
@@ -3,21 +3,10 @@
 package gosnowflake
 
 import (
-	"bufio"
-	"compress/gzip"
-	"context"
 	"database/sql/driver"
-	"encoding/json"
-	"github.com/apache/arrow/go/arrow/ipc"
-	"github.com/apache/arrow/go/arrow/memory"
 	"io"
-	"io/ioutil"
-	"net/http"
-	"net/url"
 	"reflect"
 	"strings"
-	"sync"
-	"time"
 )
 
 const (
@@ -41,21 +30,13 @@ var (
 type snowflakeRows struct {
 	sc              *snowflakeConn
 	RowType         []execResponseRowType
-	ChunkDownloader *snowflakeChunkDownloader
+	ChunkDownloader chunkDownloader
 	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
-}
-
 type snowflakeValue interface{}
 
 type chunkRowType struct {
@@ -74,32 +55,12 @@ type chunkError struct {
 	Error error
 }
 
-type snowflakeChunkDownloader struct {
-	sc                 *snowflakeConn
-	ctx                context.Context
-	Total              int64
-	TotalRowIndex      int64
-	CellCount          int
-	CurrentChunk       []chunkRowType
-	CurrentChunkIndex  int
-	CurrentChunkSize   int
-	ChunksMutex        *sync.Mutex
-	ChunkMetas         []execResponseChunk
-	Chunks             map[int][]chunkRowType
-	ChunksChan         chan int
-	ChunksError        chan *chunkError
-	ChunksErrorCounter int
-	ChunksFinalErrors  []*chunkError
-	Qrmk               string
-	QueryResultFormat  string
-	RowSet             rowSetType
-	ChunkHeader        map[string]string
-	CurrentIndex       int
-	FuncDownload       func(context.Context, *snowflakeChunkDownloader, int)
-	FuncDownloadHelper func(context.Context, *snowflakeChunkDownloader, int) error
-	FuncGet            func(context.Context, *snowflakeChunkDownloader, string, map[string]string, time.Duration) (*http.Response, error)
-	DoneDownloadCond   *sync.Cond
-	NextDownloader     *snowflakeChunkDownloader
+func (rows *snowflakeRows) Close() (err error) {
+	if err := rows.waitForAsyncQueryStatus(); err != nil {
+		return err
+	}
+	logger.WithContext(rows.sc.ctx).Debugln("Rows.Close")
+	return nil
 }
 
 // ColumnTypeDatabaseTypeName returns the database column name.
@@ -184,16 +145,16 @@ func (rows *snowflakeRows) Next(dest []driver.Value) (err error) {
 	if err := rows.waitForAsyncQueryStatus(); err != nil {
 		return err
 	}
-	row, err := rows.ChunkDownloader.Next()
+	row, err := rows.ChunkDownloader.next()
 	if err != nil {
 		// includes io.EOF
 		if err == io.EOF {
-			rows.ChunkDownloader.Chunks = nil // detach all chunks. No way to go backward without reinitialize it.
+			rows.ChunkDownloader.reset()
 		}
 		return err
 	}
 
-	if rows.ChunkDownloader.QueryResultFormat == arrowFormat {
+	if rows.ChunkDownloader.getQueryResultFormat() == arrowFormat {
 		for i, n := 0, len(row.ArrowRow); i < n; i++ {
 			dest[i] = row.ArrowRow[i]
 		}
@@ -214,9 +175,6 @@ 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
-	}
 	return rows.ChunkDownloader.hasNextResultSet()
 }
 
@@ -224,11 +182,11 @@ func (rows *snowflakeRows) NextResultSet() error {
 	if err := rows.waitForAsyncQueryStatus(); err != nil {
 		return err
 	}
-	if len(rows.ChunkDownloader.ChunkMetas) == 0 {
-		if rows.ChunkDownloader.NextDownloader == nil {
+	if len(rows.ChunkDownloader.getChunkMetas()) == 0 {
+		if rows.ChunkDownloader.getNextChunkDownloader() == nil {
 			return io.EOF
 		}
-		rows.ChunkDownloader = rows.ChunkDownloader.NextDownloader
+		rows.ChunkDownloader = rows.ChunkDownloader.getNextChunkDownloader()
 		rows.ChunkDownloader.start()
 	}
 	return rows.ChunkDownloader.nextResultSet()
@@ -249,315 +207,3 @@ func (rows *snowflakeRows) waitForAsyncQueryStatus() error {
 	}
 	return nil
 }
-
-func (scd *snowflakeChunkDownloader) totalUncompressedSize() (acc int64) {
-	for _, c := range scd.ChunkMetas {
-		acc += c.UncompressedSize
-	}
-	return
-}
-
-func (scd *snowflakeChunkDownloader) hasNextResultSet() bool {
-	// next result set exists if current chunk has remaining result sets or there is another downloader
-	return scd.CurrentChunkIndex < len(scd.ChunkMetas) || scd.NextDownloader != nil
-}
-
-func (scd *snowflakeChunkDownloader) nextResultSet() error {
-	// no error at all times as the next chunk/resultset is automatically read
-	if scd.CurrentChunkIndex < len(scd.ChunkMetas) {
-		return nil
-	}
-	return io.EOF
-}
-
-func (scd *snowflakeChunkDownloader) start() error {
-	scd.CurrentChunkSize = len(scd.RowSet.JSON) // cache the size
-	scd.CurrentIndex = -1                       // initial chunks idx
-	scd.CurrentChunkIndex = -1                  // initial chunk
-
-	scd.CurrentChunk = make([]chunkRowType, scd.CurrentChunkSize)
-	populateJSONRowSet(scd.CurrentChunk, scd.RowSet.JSON)
-
-	if scd.QueryResultFormat == arrowFormat && scd.RowSet.RowSetBase64 != "" {
-		// if the rowsetbase64 retrieved from the server is empty, move on to downloading chunks
-		var err error
-		firstArrowChunk := buildFirstArrowChunk(scd.RowSet.RowSetBase64)
-		scd.CurrentChunk, err = firstArrowChunk.decodeArrowChunk(scd.RowSet.RowType)
-		scd.CurrentChunkSize = firstArrowChunk.rowCount
-		if err != nil {
-			return err
-		}
-	}
-
-	// start downloading chunks if exists
-	chunkMetaLen := len(scd.ChunkMetas)
-	if chunkMetaLen > 0 {
-		logger.Debugf("MaxChunkDownloadWorkers: %v", MaxChunkDownloadWorkers)
-		logger.Debugf("chunks: %v, total bytes: %d", chunkMetaLen, scd.totalUncompressedSize())
-		scd.ChunksMutex = &sync.Mutex{}
-		scd.DoneDownloadCond = sync.NewCond(scd.ChunksMutex)
-		scd.Chunks = make(map[int][]chunkRowType)
-		scd.ChunksChan = make(chan int, chunkMetaLen)
-		scd.ChunksError = make(chan *chunkError, MaxChunkDownloadWorkers)
-		for i := 0; i < chunkMetaLen; i++ {
-			logger.Debugf("add chunk to channel ChunksChan: %v", i+1)
-			scd.ChunksChan <- i
-		}
-		for i := 0; i < intMin(MaxChunkDownloadWorkers, chunkMetaLen); i++ {
-			scd.schedule()
-		}
-	}
-	return nil
-}
-
-func (scd *snowflakeChunkDownloader) schedule() {
-	select {
-	case nextIdx := <-scd.ChunksChan:
-		logger.Infof("schedule chunk: %v", nextIdx+1)
-		go scd.FuncDownload(scd.ctx, scd, nextIdx)
-	default:
-		// no more download
-		logger.Info("no more download")
-	}
-}
-
-func (scd *snowflakeChunkDownloader) checkErrorRetry() (err error) {
-	select {
-	case errc := <-scd.ChunksError:
-		if scd.ChunksErrorCounter < maxChunkDownloaderErrorCounter && errc.Error != context.Canceled {
-			// add the index to the chunks channel so that the download will be retried.
-			go scd.FuncDownload(scd.ctx, scd, errc.Index)
-			scd.ChunksErrorCounter++
-			logger.Warningf("chunk idx: %v, err: %v. retrying (%v/%v)...",
-				errc.Index, errc.Error, scd.ChunksErrorCounter, maxChunkDownloaderErrorCounter)
-		} else {
-			scd.ChunksFinalErrors = append(scd.ChunksFinalErrors, errc)
-			logger.Warningf("chunk idx: %v, err: %v. no further retry", errc.Index, errc.Error)
-			return errc.Error
-		}
-	default:
-		logger.Info("no error is detected.")
-	}
-	return nil
-}
-func (scd *snowflakeChunkDownloader) Next() (chunkRowType, error) {
-	for {
-		scd.CurrentIndex++
-		if scd.CurrentIndex < scd.CurrentChunkSize {
-			return scd.CurrentChunk[scd.CurrentIndex], nil
-		}
-		scd.CurrentChunkIndex++ // next chunk
-		scd.CurrentIndex = -1   // reset
-		if scd.CurrentChunkIndex >= len(scd.ChunkMetas) {
-			break
-		}
-
-		scd.ChunksMutex.Lock()
-		if scd.CurrentChunkIndex > 1 {
-			scd.Chunks[scd.CurrentChunkIndex-1] = nil // detach the previously used chunk
-		}
-
-		for scd.Chunks[scd.CurrentChunkIndex] == nil {
-			logger.Debugf("waiting for chunk idx: %v/%v",
-				scd.CurrentChunkIndex+1, len(scd.ChunkMetas))
-
-			err := scd.checkErrorRetry()
-			if err != nil {
-				scd.ChunksMutex.Unlock()
-				return chunkRowType{}, err
-			}
-
-			// wait for chunk downloader goroutine to broadcast the event,
-			// 1) one chunk download finishes or 2) an error occurs.
-			scd.DoneDownloadCond.Wait()
-		}
-		logger.Debugf("ready: chunk %v", scd.CurrentChunkIndex+1)
-		scd.CurrentChunk = scd.Chunks[scd.CurrentChunkIndex]
-		scd.ChunksMutex.Unlock()
-		scd.CurrentChunkSize = len(scd.CurrentChunk)
-
-		// kick off the next download
-		scd.schedule()
-	}
-
-	logger.Debugf("no more data")
-	if len(scd.ChunkMetas) > 0 {
-		close(scd.ChunksError)
-		close(scd.ChunksChan)
-	}
-	return chunkRowType{}, io.EOF
-}
-
-func getChunk(
-	ctx context.Context,
-	scd *snowflakeChunkDownloader,
-	fullURL string,
-	headers map[string]string,
-	timeout time.Duration) (
-	*http.Response, error) {
-	u, err := url.Parse(fullURL)
-	if err != nil {
-		return nil, err
-	}
-	return newRetryHTTP(ctx, scd.sc.rest.Client, http.NewRequest, u, headers, timeout).execute()
-}
-
-/* largeResultSetReader is a reader that wraps the large result set with leading and tailing brackets. */
-type largeResultSetReader struct {
-	status int
-	body   io.Reader
-}
-
-func (r *largeResultSetReader) Read(p []byte) (n int, err error) {
-	if r.status == 0 {
-		p[0] = 0x5b // initial 0x5b ([)
-		r.status = 1
-		return 1, nil
-	}
-	if r.status == 1 {
-		var len int
-		len, err = r.body.Read(p)
-		if err == io.EOF {
-			r.status = 2
-			return len, nil
-		}
-		if err != nil {
-			return 0, err
-		}
-		return len, nil
-	}
-	if r.status == 2 {
-		p[0] = 0x5d // tail 0x5d (])
-		r.status = 3
-		return 1, nil
-	}
-	// ensure no data and EOF
-	return 0, io.EOF
-}
-
-func downloadChunk(ctx context.Context, scd *snowflakeChunkDownloader, idx int) {
-	logger.Infof("download start chunk: %v", idx+1)
-	defer scd.DoneDownloadCond.Broadcast()
-
-	if err := scd.FuncDownloadHelper(ctx, scd, idx); err != nil {
-		logger.Errorf(
-			"failed to extract HTTP response body. URL: %v, err: %v", scd.ChunkMetas[idx].URL, err)
-		scd.ChunksError <- &chunkError{Index: idx, Error: err}
-	} else if scd.ctx.Err() == context.Canceled || scd.ctx.Err() == context.DeadlineExceeded {
-		scd.ChunksError <- &chunkError{Index: idx, Error: scd.ctx.Err()}
-	}
-}
-
-func downloadChunkHelper(ctx context.Context, scd *snowflakeChunkDownloader, idx int) error {
-	headers := make(map[string]string)
-	if len(scd.ChunkHeader) > 0 {
-		logger.Debug("chunk header is provided.")
-		for k, v := range scd.ChunkHeader {
-			headers[k] = v
-		}
-	} else {
-		headers[headerSseCAlgorithm] = headerSseCAes
-		headers[headerSseCKey] = scd.Qrmk
-	}
-
-	resp, err := scd.FuncGet(ctx, scd, scd.ChunkMetas[idx].URL, headers, scd.sc.rest.RequestTimeout)
-	if err != nil {
-		return err
-	}
-	bufStream := bufio.NewReader(resp.Body)
-	defer resp.Body.Close()
-	logger.Infof("response returned chunk: %v, resp: %v", idx+1, resp)
-	if resp.StatusCode != http.StatusOK {
-		b, err := ioutil.ReadAll(bufStream)
-		if err != nil {
-			return err
-		}
-		logger.Infof("HTTP: %v, URL: %v, Body: %v", resp.StatusCode, scd.ChunkMetas[idx].URL, b)
-		logger.Infof("Header: %v", resp.Header)
-		return &SnowflakeError{
-			Number:      ErrFailedToGetChunk,
-			SQLState:    SQLStateConnectionFailure,
-			Message:     errMsgFailedToGetChunk,
-			MessageArgs: []interface{}{idx},
-		}
-	}
-	return decodeChunk(scd, idx, bufStream)
-}
-
-func decodeChunk(scd *snowflakeChunkDownloader, idx int, bufStream *bufio.Reader) (err error) {
-	gzipMagic, err := bufStream.Peek(2)
-	if err != nil {
-		return err
-	}
-	start := time.Now()
-	var source io.Reader
-	if gzipMagic[0] == 0x1f && gzipMagic[1] == 0x8b {
-		// detects and uncompresses Gzip format data
-		bufStream0, err := gzip.NewReader(bufStream)
-		if err != nil {
-			return err
-		}
-		defer bufStream0.Close()
-		source = bufStream0
-	} else {
-		source = bufStream
-	}
-	st := &largeResultSetReader{
-		status: 0,
-		body:   source,
-	}
-	var respd []chunkRowType
-	if scd.QueryResultFormat != arrowFormat {
-		var decRespd [][]*string
-		if !CustomJSONDecoderEnabled {
-			dec := json.NewDecoder(st)
-			for {
-				if err := dec.Decode(&decRespd); err == io.EOF {
-					break
-				} else if err != nil {
-					return err
-				}
-			}
-		} else {
-			decRespd, err = decodeLargeChunk(st, scd.ChunkMetas[idx].RowCount, scd.CellCount)
-			if err != nil {
-				return err
-			}
-		}
-		respd = make([]chunkRowType, len(decRespd))
-		populateJSONRowSet(respd, decRespd)
-	} else {
-		ipcReader, err := ipc.NewReader(source)
-		if err != nil {
-			return err
-		}
-		arc := arrowResultChunk{
-			*ipcReader,
-			0,
-			int(scd.totalUncompressedSize()),
-			memory.NewGoAllocator(),
-		}
-		respd, err = arc.decodeArrowChunk(scd.RowSet.RowType)
-		if err != nil {
-			return err
-		}
-	}
-	logger.Debugf(
-		"decoded %d rows w/ %d bytes in %s (chunk %v)",
-		scd.ChunkMetas[idx].RowCount,
-		scd.ChunkMetas[idx].UncompressedSize,
-		time.Since(start), idx+1,
-	)
-
-	scd.ChunksMutex.Lock()
-	defer scd.ChunksMutex.Unlock()
-	scd.Chunks[idx] = respd
-	return nil
-}
-
-func populateJSONRowSet(dst []chunkRowType, src [][]*string) {
-	// populate string rowset from src to dst's chunkRowType struct's RowSet field
-	for i, row := range src {
-		dst[i].RowSet = row
-	}
-}
diff --git a/statement.go b/statement.go
index 73a5af3..5371ad0 100644
--- a/statement.go
+++ b/statement.go
@@ -9,15 +9,6 @@ import (
 
 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"
-	// QueryIDChan is the channel to receive the query ID from
-	QueryIDChan paramKey = "QUERY_ID_CHANNEL"
-)
-
 type snowflakeStmt struct {
 	sc    *snowflakeConn
 	query string
@@ -54,18 +45,3 @@ func (stmt *snowflakeStmt) Query(args []driver.Value) (driver.Rows, error) {
 	logger.WithContext(stmt.sc.ctx).Infoln("Stmt.Query")
 	return stmt.sc.Query(stmt.query, args)
 }
-
-// WithMultiStatement returns a context that allows the user to execute the desired number of sql queries in one query
-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
-}
-
-// WithQueryIDChan returns a context that contains the channel to receive the query ID
-func WithQueryIDChan(ctx context.Context, c chan<- string) context.Context {
-	return context.WithValue(ctx, QueryIDChan, c)
-}
diff --git a/util.go b/util.go
index ee8593d..1ac5a4d 100644
--- a/util.go
+++ b/util.go
@@ -12,14 +12,44 @@ import (
 
 type contextKey string
 
-// SnowflakeRequestIDKey is optional context key to specify request id
-const SnowflakeRequestIDKey contextKey = "SNOWFLAKE_REQUEST_ID"
+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"
+	// QueryIDChan is the channel to receive the query ID from
+	QueryIDChan paramKey = "QUERY_ID_CHANNEL"
+	// SnowflakeRequestIDKey is optional context key to specify request id
+	SnowflakeRequestIDKey contextKey = "SNOWFLAKE_REQUEST_ID"
+	// streamChunkDownload determines whether to use a stream based chunk downloader
+	streamChunkDownload paramKey = "STREAM_CHUNK_DOWNLOAD"
+)
+
+// WithMultiStatement returns a context that allows the user to execute the desired number of sql queries in one query
+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 {
+	return context.WithValue(ctx, asyncMode, true)
+}
+
+// WithQueryIDChan returns a context that contains the channel to receive the query ID
+func WithQueryIDChan(ctx context.Context, c chan<- string) context.Context {
+	return context.WithValue(ctx, QueryIDChan, c)
+}
 
 // WithRequestID returns a new context with the specified snowflake request id
 func WithRequestID(ctx context.Context, requestID uuid.UUID) context.Context {
 	return context.WithValue(ctx, SnowflakeRequestIDKey, requestID)
 }
 
+// WithStreamDownloader returns a context that allows the use of a stream based chunk downloader
+func WithStreamDownloader(ctx context.Context) context.Context {
+	return context.WithValue(ctx, streamChunkDownload, true)
+}
+
 // Get the request ID from the context if specified, otherwise generate one
 func getOrGenerateRequestIDFromContext(ctx context.Context) uuid.UUID {
 	requestID, ok := ctx.Value(SnowflakeRequestIDKey).(uuid.UUID)


Test patch:
diff --git a/async_test.go b/async_test.go
index e46164a..4c45d34 100644
--- a/async_test.go
+++ b/async_test.go
@@ -18,7 +18,7 @@ func TestAsyncMode(t *testing.T) {
 	}
 	defer db.Close()
 
-	ctx, _ := WithAsyncMode(context.Background())
+	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()
@@ -64,7 +64,7 @@ func TestAsyncQueryFail(t *testing.T) {
 	}
 	defer db.Close()
 
-	ctx, _ := WithAsyncMode(context.Background())
+	ctx := WithAsyncMode(context.Background())
 	rows, err := db.QueryContext(ctx, "selectt 1")
 	if err != nil {
 		t.Fatal("asynchronous query should always return nil error")
@@ -89,7 +89,7 @@ func TestMultipleAsyncQueries(t *testing.T) {
 	}
 	defer db.Close()
 
-	ctx, _ := WithAsyncMode(context.Background())
+	ctx := WithAsyncMode(context.Background())
 	s1 := "foo"
 	s2 := "bar"
 	rows1, _ := db.QueryContext(ctx, fmt.Sprintf("select distinct '%v' from table (generator(timelimit=>%v))", s1, 30))
diff --git a/chunk_test.go b/chunk_test.go
index aab0466..ab2cb00 100644
--- a/chunk_test.go
+++ b/chunk_test.go
@@ -2,7 +2,12 @@ package gosnowflake
 
 import (
 	"bytes"
+	"context"
+	"database/sql"
 	"encoding/json"
+	"fmt"
+	"io"
+	"math/rand"
 	"strings"
 	"testing"
 )
@@ -163,3 +168,221 @@ func testDecodeErr(t *testing.T, s string) {
 		t.Fatalf("expected decode to fail for input: %s", s)
 	}
 }
+
+type mockStreamChunkFetcher struct {
+	chunks map[string][][]*string
+}
+
+func (f *mockStreamChunkFetcher) fetch(url string, stream chan<- []*string) error {
+	for _, row := range f.chunks[url] {
+		stream <- row
+	}
+	return nil
+}
+
+func TestStreamChunkDownloaderFirstRows(t *testing.T) {
+	fetcher := &mockStreamChunkFetcher{}
+	firstRows := generateStreamChunkRows(10, 4)
+	downloader := newStreamChunkDownloader(context.Background(), fetcher, int64(len(firstRows)), firstRows, []execResponseChunk{})
+	if err := downloader.start(); err != nil {
+		t.Fatalf("chunk download start failed. err: %v", err)
+	}
+	for i := 0; i < len(firstRows); i++ {
+		if !downloader.hasNextResultSet() {
+			t.Error("failed to retrieve next result set")
+		}
+		err := downloader.nextResultSet()
+		if err != nil {
+			t.Fatalf("failed to retrieve data. err: %v", err)
+		}
+		row, err := downloader.next()
+		if err != nil {
+			t.Fatalf("failed to retrieve data. err: %v", err)
+		}
+		assertEqualRows(firstRows[i], row)
+	}
+	row, err := downloader.next()
+	if !assertEmptyChunkRow(row) {
+		t.Fatal("row should be empty")
+	}
+	if err != io.EOF {
+		t.Fatalf("failed to finish getting data. err: %v", err)
+	}
+	if downloader.hasNextResultSet() {
+		t.Error("downloader has next result set. expected none.")
+	}
+	if downloader.nextResultSet() != io.EOF {
+		t.Fatalf("failed to finish getting data. err: %v", err)
+	}
+}
+
+func TestStreamChunkDownloaderChunks(t *testing.T) {
+	chunks, responseChunks := generateStreamChunkDownloaderChunks([]string{"foo", "bar"}, 4, 4)
+	fetcher := &mockStreamChunkFetcher{chunks}
+	firstRows := generateStreamChunkRows(2, 4)
+	downloader := newStreamChunkDownloader(context.Background(), fetcher, int64(len(firstRows)), firstRows, responseChunks)
+	if err := downloader.start(); err != nil {
+		t.Fatalf("chunk download start failed. err: %v", err)
+	}
+	for i := 0; i < len(firstRows); i++ {
+		if !downloader.hasNextResultSet() {
+			t.Error("failed to retrieve next result set")
+		}
+		err := downloader.nextResultSet()
+		if err != nil {
+			t.Fatalf("failed to retrieve data. err: %v", err)
+		}
+		row, err := downloader.next()
+		if err != nil {
+			t.Fatalf("failed to retrieve data. err: %v", err)
+		}
+		assertEqualRows(firstRows[i], row)
+	}
+	for _, chunk := range responseChunks {
+		for _, chunkRow := range chunks[chunk.URL] {
+			if !downloader.hasNextResultSet() {
+				t.Error("failed to retrieve next result set")
+			}
+			row, err := downloader.next()
+			if err != nil {
+				t.Fatalf("failed to retrieve data. err: %v", err)
+			}
+			assertEqualRows(chunkRow, row)
+		}
+	}
+	row, err := downloader.next()
+	if !assertEmptyChunkRow(row) {
+		t.Fatal("row should be empty")
+	}
+	if err != io.EOF {
+		t.Fatalf("failed to finish getting data. err: %v", err)
+	}
+	if downloader.hasNextResultSet() {
+		t.Error("downloader has next result set. expected none.")
+	}
+	if downloader.nextResultSet() != io.EOF {
+		t.Fatalf("failed to finish getting data. err: %v", err)
+	}
+}
+
+func TestCopyChunkStream(t *testing.T) {
+	foo := "foo"
+	bar := "bar"
+
+	r := strings.NewReader(`["foo","bar",null],["bar",null,"foo"],[]`)
+	c := make(chan []*string, 3)
+	err := copyChunkStream(r, c)
+	if err != nil {
+		t.Fatalf("error while copying chunk stream. err: %v", err)
+	}
+	assertEqualRows([]*string{&foo, &bar, nil}, <-c)
+	assertEqualRows([]*string{&bar, nil, &foo}, <-c)
+	assertEqualRows([]*string{}, <-c)
+}
+
+func TestCopyChunkStreamInvalid(t *testing.T) {
+	var r io.Reader
+	var c chan []*string
+	var err error
+
+	r = strings.NewReader("oops")
+	c = make(chan []*string, 1)
+	err = copyChunkStream(r, c)
+	if err == nil {
+		t.Fatalf("should fail to retrieve data. err: %v", err)
+	}
+
+	r = strings.NewReader(`[["foo"], ["bar"]]`)
+	c = make(chan []*string, 1)
+	err = copyChunkStream(r, c)
+	if err == nil {
+		t.Fatalf("should fail to retrieve data. err: %v", err)
+	}
+
+	r = strings.NewReader(`{"foo": "bar"}`)
+	c = make(chan []*string, 1)
+	err = copyChunkStream(r, c)
+	if err == nil {
+		t.Fatalf("should fail to retrieve data. err: %v", err)
+	}
+}
+
+func generateStreamChunkDownloaderChunks(urls []string, numRows, numCols int) (map[string][][]*string, []execResponseChunk) {
+	chunks := map[string][][]*string{}
+	var responseChunks []execResponseChunk
+	for _, url := range urls {
+		rows := generateStreamChunkRows(numRows, numCols)
+		chunks[url] = rows
+		responseChunks = append(responseChunks, execResponseChunk{url, len(rows), -1, -1})
+	}
+	return chunks, responseChunks
+}
+
+func generateStreamChunkRows(numRows, numCols int) [][]*string {
+	rows := [][]*string{}
+	for i := 0; i < numRows; i++ {
+		cols := []*string{}
+		for j := 0; j < numCols; j++ {
+			col := fmt.Sprintf("%d", rand.Intn(1000))
+			cols = append(cols, &col)
+		}
+		rows = append(rows, cols)
+	}
+	return rows
+}
+
+func assertEqualRows(expected []*string, actual interface{}) bool {
+	switch v := actual.(type) {
+	case chunkRowType:
+		for i := range expected {
+			if expected[i] != v.RowSet[i] {
+				return false
+			}
+		}
+		return len(expected) == len(v.RowSet)
+	case []*string:
+		for i := range expected {
+			if expected[i] != v[i] {
+				return false
+			}
+		}
+		return len(expected) == len(v)
+	}
+	return false
+}
+
+func assertEmptyChunkRow(row chunkRowType) bool {
+	return assertEqualRows(make([]*string, len(row.RowSet)), row)
+}
+
+func TestWithStreamDownloader(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 := WithStreamDownloader(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)
+	}
+}
diff --git a/connection_test.go b/connection_test.go
index 09c3e34..db6b606 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -17,7 +17,7 @@ const serviceNameAppend = "a"
 // postQueryMock would generate a response based on the X-Snowflake-Service header, to generate a response
 // with the SERVICE_NAME field appending a character at the end of the header
 // This way it could test both the send and receive logic
-func postQueryMock(_ context.Context, _ *snowflakeRestful, _ *url.Values, headers map[string]string, _ []byte, _ time.Duration, _ uuid.UUID) (*execResponse, error) {
+func postQueryMock(_ context.Context, _ *snowflakeRestful, _ *url.Values, headers map[string]string, _ []byte, _ time.Duration, _ uuid.UUID, _ *Config) (*execResponse, error) {
 	var serviceName string
 	if serviceHeader, ok := headers["X-Snowflake-Service"]; ok {
 		serviceName = serviceHeader + serviceNameAppend
@@ -38,7 +38,7 @@ func postQueryMock(_ context.Context, _ *snowflakeRestful, _ *url.Values, header
 
 func TestExecWithEmptyRequestID(t *testing.T) {
 	ctx := WithRequestID(context.Background(), uuid.Nil)
-	postQueryMock := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid.UUID) (*execResponse, error) {
+	postQueryMock := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid.UUID, _ *Config) (*execResponse, error) {
 		// ensure the same requestID from context is used
 		if len(requestID) == 0 {
 			t.Fatal("requestID is empty")
@@ -69,7 +69,7 @@ func TestExecWithEmptyRequestID(t *testing.T) {
 func TestExecWithSpecificRequestID(t *testing.T) {
 	origRequestID := uuid.New()
 	ctx := WithRequestID(context.Background(), origRequestID)
-	postQueryMock := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid.UUID) (*execResponse, error) {
+	postQueryMock := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid.UUID, _ *Config) (*execResponse, error) {
 		// ensure the same requestID from context is used
 		if requestID != origRequestID {
 			t.Fatal("requestID doesn't match")
diff --git a/restful_test.go b/restful_test.go
index 5b7da90..cb8f90c 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -1,4 +1,5 @@
-// Copyright (c) 2017-2019 Snowflake Computing Inc. All right reserved.
+// Copyright (c) 2017-2021 Snowflake Computing Inc. All right reserved.
+
 package gosnowflake
 
 import (
@@ -96,19 +97,19 @@ func TestUnitPostQueryHelperError(t *testing.T) {
 	var err error
 	var requestID uuid.UUID
 	requestID = uuid.New()
-	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, requestID)
+	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, requestID, &Config{})
 	if err == nil {
 		t.Fatalf("should have failed to post")
 	}
 	sr.FuncPost = postTestAppBadGatewayError
 	requestID = uuid.New()
-	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, requestID)
+	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, requestID, &Config{})
 	if err == nil {
 		t.Fatalf("should have failed to post")
 	}
 	sr.FuncPost = postTestSuccessButInvalidJSON
 	requestID = uuid.New()
-	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, requestID)
+	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, requestID, &Config{})
 	if err == nil {
 		t.Fatalf("should have failed to post")
 	}
@@ -125,7 +126,7 @@ func renewSessionTestError(_ context.Context, _ *snowflakeRestful, _ time.Durati
 func TestUnitPostQueryHelperRenewSession(t *testing.T) {
 	var err error
 	origRequestID := uuid.New()
-	postQueryTest := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid.UUID) (*execResponse, error) {
+	postQueryTest := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid.UUID, _ *Config) (*execResponse, error) {
 		// ensure the same requestID is used after the session token is renewed.
 		if requestID != origRequestID {
 			t.Fatal("requestID doesn't match")
@@ -145,12 +146,12 @@ func TestUnitPostQueryHelperRenewSession(t *testing.T) {
 		FuncRenewSession: renewSessionTest,
 	}
 
-	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, origRequestID)
+	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, origRequestID, &Config{})
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
 	sr.FuncRenewSession = renewSessionTestError
-	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, origRequestID)
+	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, origRequestID, &Config{})
 	if err == nil {
 		t.Fatal("should have failed to renew session")
 	}
diff --git a/rows_test.go b/rows_test.go
index f70d394..100e6c5 100644
--- a/rows_test.go
+++ b/rows_test.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
 
diff --git a/statement_test.go b/statement_test.go
index 68211b9..e506906 100644
--- a/statement_test.go
+++ b/statement_test.go
@@ -69,7 +69,7 @@ func TestEmitQueryID(t *testing.T) {
 
 	queryIDChan := make(chan string, 1)
 	numrows := 100000
-	ctx, _ := WithAsyncMode(context.Background())
+	ctx := WithAsyncMode(context.Background())
 	ctx = WithQueryIDChan(ctx, queryIDChan)
 
 	goRoutineChan := make(chan string)