SNOW-985356 Do not start chunk downloader when first batch causes error (snowflakedb#993)

Author: sfc-gh-pfus

arrow_chunk.go

@@ -68,15 +68,15 @@ func (arc *arrowResultChunk) decodeArrowBatch(scd *snowflakeChunkDownloader) (*[
 }
 
 // Build arrow chunk based on RowSet of base64
-func buildFirstArrowChunk(rowsetBase64 string, loc *time.Location, alloc memory.Allocator) arrowResultChunk {
+func buildFirstArrowChunk(rowsetBase64 string, loc *time.Location, alloc memory.Allocator) (arrowResultChunk, error) {
 	rowSetBytes, err := base64.StdEncoding.DecodeString(rowsetBase64)
 	if err != nil {
-		return arrowResultChunk{}
+		return arrowResultChunk{}, err
 	}
 	rr, err := ipc.NewReader(bytes.NewReader(rowSetBytes), ipc.WithAllocator(alloc))
 	if err != nil {
-		return arrowResultChunk{}
+		return arrowResultChunk{}, err
 	}
 
-	return arrowResultChunk{rr, 0, loc, alloc}
+	return arrowResultChunk{rr, 0, loc, alloc}, nil
 }

chunk_downloader.go

@@ -109,7 +109,10 @@ func (scd *snowflakeChunkDownloader) start() error {
 		if scd.sc != nil && scd.sc.cfg != nil {
 			loc = getCurrentLocation(scd.sc.cfg.Params)
 		}
-		firstArrowChunk := buildFirstArrowChunk(scd.RowSet.RowSetBase64, loc, scd.pool)
+		firstArrowChunk, err := buildFirstArrowChunk(scd.RowSet.RowSetBase64, loc, scd.pool)
+		if err != nil {
+			return err
+		}
 		higherPrecision := higherPrecisionEnabled(scd.ctx)
 		scd.CurrentChunk, err = firstArrowChunk.decodeArrowChunk(scd.RowSet.RowType, higherPrecision)
 		scd.CurrentChunkSize = firstArrowChunk.rowCount
@@ -274,7 +277,10 @@ func (scd *snowflakeChunkDownloader) startArrowBatches() error {
 	if scd.sc != nil && scd.sc.cfg != nil {
 		loc = getCurrentLocation(scd.sc.cfg.Params)
 	}
-	firstArrowChunk := buildFirstArrowChunk(scd.RowSet.RowSetBase64, loc, scd.pool)
+	firstArrowChunk, err := buildFirstArrowChunk(scd.RowSet.RowSetBase64, loc, scd.pool)
+	if err != nil {
+		return err
+	}
 	scd.FirstBatch = &ArrowBatch{
 		idx:                0,
 		scd:                scd,

chunk_downloader_test.go

@@ -0,0 +1,28 @@
+package gosnowflake
+
+import (
+	"context"
+	"testing"
+)
+
+func TestChunkDownloaderDoesNotStartWhenArrowParsingCausesError(t *testing.T) {
+	tcs := []string{
+		"invalid base64",
+		"aW52YWxpZCBhcnJvdw==", // valid base64, but invalid arrow
+	}
+	for _, tc := range tcs {
+		t.Run(tc, func(t *testing.T) {
+			scd := snowflakeChunkDownloader{
+				ctx:               context.Background(),
+				QueryResultFormat: "arrow",
+				RowSet: rowSetType{
+					RowSetBase64: tc,
+				},
+			}
+
+			err := scd.start()
+
+			assertNotNilF(t, err)
+		})
+	}
+}