Repository: snowflakedb/gosnowflake
Pull request #380 (2021-02-09 23:25:44)

SNOW-257769 Multi Column Result Sets
### Description
Add support for processing multi-column result sets by modifying the RowType per chunk downloader, necessary for multi-statement executions/queries.

### 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
index ab88feb..ae8ae3d 100644
--- a/chunk_downloader.go
+++ b/chunk_downloader.go
@@ -32,6 +32,7 @@ type chunkDownloader interface {
 	getQueryResultFormat() resultFormat
 	setNextChunkDownloader(downloader chunkDownloader)
 	getNextChunkDownloader() chunkDownloader
+	getRowType() []execResponseRowType
 }
 
 type snowflakeChunkDownloader struct {
@@ -227,6 +228,10 @@ func (scd *snowflakeChunkDownloader) getNextChunkDownloader() chunkDownloader {
 	return scd.NextDownloader
 }
 
+func (scd *snowflakeChunkDownloader) getRowType() []execResponseRowType {
+	return scd.RowSet.RowType
+}
+
 func getChunk(
 	ctx context.Context,
 	scd *snowflakeChunkDownloader,
@@ -503,6 +508,10 @@ func (scd *streamChunkDownloader) getNextChunkDownloader() chunkDownloader {
 	return scd.NextDownloader
 }
 
+func (scd *streamChunkDownloader) getRowType() []execResponseRowType {
+	return scd.RowSet.RowType
+}
+
 func useStreamDownloader(ctx context.Context) bool {
 	val := ctx.Value(streamChunkDownload)
 	if val == nil {
@@ -531,17 +540,18 @@ func newStreamChunkDownloader(
 	ctx context.Context,
 	fetcher streamChunkFetcher,
 	total int64,
+	rowType []execResponseRowType,
 	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),
+		Total:      total,
+		ChunkMetas: chunks,
+		RowSet:     rowSetType{RowType: rowType, JSON: firstRows},
 	}
 }
 
diff --git a/connection.go b/connection.go
index 58b50d1..6b7964b 100644
--- a/connection.go
+++ b/connection.go
@@ -337,7 +337,6 @@ func (sc *snowflakeConn) QueryContext(ctx context.Context, query string, args []
 
 	rows := new(snowflakeRows)
 	rows.sc = sc
-	rows.RowType = data.Data.RowType
 	rows.ChunkDownloader = populateChunkDownloader(ctx, sc, data.Data)
 	rows.queryID = sc.QueryID
 
@@ -683,7 +682,6 @@ func getAsync(
 			res.errChannel <- nil // mark exec status complete
 		} else {
 			rows.sc = sc
-			rows.RowType = respd.Data.RowType
 			rows.ChunkDownloader = populateChunkDownloader(ctx, sc, respd.Data)
 			rows.queryID = respd.Data.QueryID
 			if sc.isMultiStmt(respd.Data) {
@@ -736,7 +734,7 @@ func populateChunkDownloader(ctx context.Context, sc *snowflakeConn, data execRe
 			headers:  data.ChunkHeaders,
 			qrmk:     data.Qrmk,
 		}
-		return newStreamChunkDownloader(ctx, fetcher, data.Total, data.RowSet, data.Chunks)
+		return newStreamChunkDownloader(ctx, fetcher, data.Total, data.RowType, data.RowSet, data.Chunks)
 	}
 
 	return &snowflakeChunkDownloader{
diff --git a/rows.go b/rows.go
index 8a06735..3367915 100644
--- a/rows.go
+++ b/rows.go
@@ -29,7 +29,6 @@ var (
 
 type snowflakeRows struct {
 	sc              *snowflakeConn
-	RowType         []execResponseRowType
 	ChunkDownloader chunkDownloader
 	queryID         string
 	status          queryStatus
@@ -68,7 +67,7 @@ func (rows *snowflakeRows) ColumnTypeDatabaseTypeName(index int) string {
 	if err := rows.waitForAsyncQueryStatus(); err != nil {
 		return err.Error()
 	}
-	return strings.ToUpper(rows.RowType[index].Type)
+	return strings.ToUpper(rows.ChunkDownloader.getRowType()[index].Type)
 }
 
 // ColumnTypeLength returns the length of the column
@@ -76,12 +75,12 @@ 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) {
+	if index < 0 || index > len(rows.ChunkDownloader.getRowType()) {
 		return 0, false
 	}
-	switch rows.RowType[index].Type {
+	switch rows.ChunkDownloader.getRowType()[index].Type {
 	case "text", "variant", "object", "array", "binary":
-		return rows.RowType[index].Length, true
+		return rows.ChunkDownloader.getRowType()[index].Length, true
 	}
 	return 0, false
 }
@@ -90,26 +89,27 @@ 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) {
+	if index < 0 || index > len(rows.ChunkDownloader.getRowType()) {
 		return false, false
 	}
-	return rows.RowType[index].Nullable, true
+	return rows.ChunkDownloader.getRowType()[index].Nullable, true
 }
 
 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) {
+	rowType := rows.ChunkDownloader.getRowType()
+	if index < 0 || index > len(rowType) {
 		return 0, 0, false
 	}
-	switch rows.RowType[index].Type {
+	switch rowType[index].Type {
 	case "fixed":
-		return rows.RowType[index].Precision, rows.RowType[index].Scale, true
+		return rowType[index].Precision, rowType[index].Scale, true
 	case "time":
-		return rows.RowType[index].Scale, 0, true
+		return rowType[index].Scale, 0, true
 	case "timestamp":
-		return rows.RowType[index].Scale, 0, true
+		return rowType[index].Scale, 0, true
 	}
 	return 0, 0, false
 }
@@ -119,9 +119,9 @@ func (rows *snowflakeRows) Columns() []string {
 		return make([]string, 0)
 	}
 	logger.Debug("Rows.Columns")
-	ret := make([]string, len(rows.RowType))
-	for i, n := 0, len(rows.RowType); i < n; i++ {
-		ret[i] = rows.RowType[i].Name
+	ret := make([]string, len(rows.ChunkDownloader.getRowType()))
+	for i, n := 0, len(rows.ChunkDownloader.getRowType()); i < n; i++ {
+		ret[i] = rows.ChunkDownloader.getRowType()[i].Name
 	}
 	return ret
 }
@@ -130,7 +130,7 @@ 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)
+	return snowflakeTypeToGo(rows.ChunkDownloader.getRowType()[index].Type, rows.ChunkDownloader.getRowType()[index].Scale)
 }
 
 func (rows *snowflakeRows) GetQueryID() string {
@@ -162,7 +162,7 @@ func (rows *snowflakeRows) Next(dest []driver.Value) (err error) {
 		for i, n := 0, len(row.RowSet); i < n; i++ {
 			// could move to chunk downloader so that each go routine
 			// can convert data
-			err := stringToValue(&dest[i], rows.RowType[i], row.RowSet[i])
+			err := stringToValue(&dest[i], rows.ChunkDownloader.getRowType()[i], row.RowSet[i])
 			if err != nil {
 				return err
 			}


Test patch:
diff --git a/chunk_test.go b/chunk_test.go
index ab2cb00..f7a33d8 100644
--- a/chunk_test.go
+++ b/chunk_test.go
@@ -183,7 +183,13 @@ func (f *mockStreamChunkFetcher) fetch(url string, stream chan<- []*string) erro
 func TestStreamChunkDownloaderFirstRows(t *testing.T) {
 	fetcher := &mockStreamChunkFetcher{}
 	firstRows := generateStreamChunkRows(10, 4)
-	downloader := newStreamChunkDownloader(context.Background(), fetcher, int64(len(firstRows)), firstRows, []execResponseChunk{})
+	downloader := newStreamChunkDownloader(
+		context.Background(),
+		fetcher,
+		int64(len(firstRows)),
+		[]execResponseRowType{},
+		firstRows,
+		[]execResponseChunk{})
 	if err := downloader.start(); err != nil {
 		t.Fatalf("chunk download start failed. err: %v", err)
 	}
@@ -220,7 +226,13 @@ 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)
+	downloader := newStreamChunkDownloader(
+		context.Background(),
+		fetcher,
+		int64(len(firstRows)),
+		[]execResponseRowType{},
+		firstRows,
+		responseChunks)
 	if err := downloader.start(); err != nil {
 		t.Fatalf("chunk download start failed. err: %v", err)
 	}
diff --git a/multistatement_test.go b/multistatement_test.go
index 39a2039..2c7ef57 100644
--- a/multistatement_test.go
+++ b/multistatement_test.go
@@ -441,3 +441,53 @@ func TestMultiStatementCountMismatch(t *testing.T) {
 		t.Fatal("should have failed to query multiple statements")
 	}
 }
+
+func TestMultiStatementVaryingColumnCount(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()
+
+	multiStmtQuery := "select c1 from test_tbl;\n" +
+		"select c1,c2 from test_tbl;"
+	db.Exec("create or replace table test_tbl(c1 int, c2 int)")
+	db.Exec("insert into test_tbl values(1, 0)")
+	ctx, _ := WithMultiStatement(context.Background(), 0)
+	rows, err := db.QueryContext(ctx, multiStmtQuery)
+	if err != nil {
+		t.Error(err)
+	}
+	defer rows.Close()
+
+	var v1, v2 int
+	if rows.Next() {
+		err = rows.Scan(&v1)
+		if err != nil {
+			t.Errorf("failed to scan: %#v", err)
+		}
+		if v1 != 1 {
+			t.Fatalf("failed to fetch. value: %v", v1)
+		}
+	} else {
+		t.Error("failed to query")
+	}
+
+	if !rows.NextResultSet() {
+		t.Error("failed to retrieve next result set")
+	}
+	if rows.Next() {
+		err = rows.Scan(&v1, &v2)
+		if err != nil {
+			t.Errorf("failed to scan: %#v", err)
+		}
+		if v1 != 1 || v2 != 0 {
+			t.Fatalf("failed to fetch. value: %v, %v", v1, v2)
+		}
+	} else {
+		t.Error("failed to query")
+	}
+	db.Exec("drop table if exists test_tbl")
+}
diff --git a/rows_test.go b/rows_test.go
index 100e6c5..ca3dc20 100644
--- a/rows_test.go
+++ b/rows_test.go
@@ -33,7 +33,6 @@ func TestRowsWithoutChunkDownloader(t *testing.T) {
 	cm := []execResponseChunk{}
 	rows := new(snowflakeRows)
 	rows.sc = nil
-	rows.RowType = rt
 	rows.ChunkDownloader = &snowflakeChunkDownloader{
 		sc:                 nil,
 		ctx:                context.Background(),
@@ -43,10 +42,9 @@ func TestRowsWithoutChunkDownloader(t *testing.T) {
 		Qrmk:               "",
 		FuncDownload:       nil,
 		FuncDownloadHelper: nil,
-		RowSet:             rowSetType{JSON: cc},
+		RowSet:             rowSetType{RowType: rt, JSON: cc},
 	}
 	rows.ChunkDownloader.start()
-	// var dest []driver.Value
 	dest := make([]driver.Value, 2)
 	for i = 0; i < len(cc); i++ {
 		err := rows.Next(dest)
@@ -105,7 +103,6 @@ func TestRowsWithChunkDownloader(t *testing.T) {
 	}
 	rows := new(snowflakeRows)
 	rows.sc = nil
-	rows.RowType = rt
 	rows.ChunkDownloader = &snowflakeChunkDownloader{
 		sc:            nil,
 		ctx:           context.Background(),
@@ -114,7 +111,7 @@ func TestRowsWithChunkDownloader(t *testing.T) {
 		TotalRowIndex: int64(-1),
 		Qrmk:          "HAHAHA",
 		FuncDownload:  downloadChunkTest,
-		RowSet:        rowSetType{JSON: cc},
+		RowSet:        rowSetType{RowType: rt, JSON: cc},
 	}
 	rows.ChunkDownloader.start()
 	cnt := 0
@@ -128,7 +125,6 @@ func TestRowsWithChunkDownloader(t *testing.T) {
 		if err != nil {
 			t.Fatalf("failed to get value. err: %v", err)
 		}
-		// fmt.Printf("data: %v\n", dest)
 		cnt++
 	}
 	if cnt != len(cc)+numChunks*rowsInChunk {
@@ -186,7 +182,6 @@ func TestRowsWithChunkDownloaderError(t *testing.T) {
 	}
 	rows := new(snowflakeRows)
 	rows.sc = nil
-	rows.RowType = rt
 	rows.ChunkDownloader = &snowflakeChunkDownloader{
 		sc:            nil,
 		ctx:           context.Background(),
@@ -195,7 +190,7 @@ func TestRowsWithChunkDownloaderError(t *testing.T) {
 		TotalRowIndex: int64(-1),
 		Qrmk:          "HOHOHO",
 		FuncDownload:  downloadChunkTestError,
-		RowSet:        rowSetType{JSON: cc},
+		RowSet:        rowSetType{RowType: rt, JSON: cc},
 	}
 	rows.ChunkDownloader.start()
 	cnt := 0
@@ -265,7 +260,6 @@ func TestRowsWithChunkDownloaderErrorFail(t *testing.T) {
 	}
 	rows := new(snowflakeRows)
 	rows.sc = nil
-	rows.RowType = rt
 	rows.ChunkDownloader = &snowflakeChunkDownloader{
 		sc:            nil,
 		ctx:           context.Background(),
@@ -274,7 +268,7 @@ func TestRowsWithChunkDownloaderErrorFail(t *testing.T) {
 		TotalRowIndex: int64(-1),
 		Qrmk:          "HOHOHO",
 		FuncDownload:  downloadChunkTestErrorFail,
-		RowSet:        rowSetType{JSON: cc},
+		RowSet:        rowSetType{RowType: rt, JSON: cc},
 	}
 	rows.ChunkDownloader.start()
 	cnt := 0
@@ -290,7 +284,6 @@ func TestRowsWithChunkDownloaderErrorFail(t *testing.T) {
 				"failure was expected by the number of rows is wrong. expected: %v, got: %v", 715, cnt)
 			break
 		}
-		// fmt.Printf("data: %v\n", dest)
 		cnt++
 	}
 }