Repository: snowflakedb/gosnowflake
Pull request #993 (2023-12-07 10:49:54)

SNOW-985356 Do not start chunk downloader when first batch causes error
### Description
When error occurred when parsing first arrow batch, the error was just swallowed and empty struct was returned. This caused that when the struct was used, caused panic on nil field access.

### 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
- [x] Extended the README / documentation, if necessary


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/993?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
Attention: `3 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`e799467`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/e79946798ead5939e451909501229101013b0c5e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) 82.79% compared to head [(`49fe0a8`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/993?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) 82.10%.
> Report is 1 commits behind head on master.

| [Files](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/993?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | Patch % | Lines |
|---|---|---|
| [chunk\_downloader.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/993?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-Y2h1bmtfZG93bmxvYWRlci5nbw==) | 75.00% | [2 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/993?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) |
| [arrow\_chunk.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/993?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-YXJyb3dfY2h1bmsuZ28=) | 75.00% | [1 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/993?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master     #993      +/-   ##
==========================================
- Coverage   82.79%   82.10%   -0.69%     
==========================================
  Files          50       49       -1     
  Lines        9687     9696       +9     
==========================================
- Hits         8020     7961      -59     
- Misses       1667     1735      +68     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/993?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).


Patch:
diff --git a/arrow_chunk.go b/arrow_chunk.go
index 15851a8..27dea9a 100644
--- a/arrow_chunk.go
+++ b/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
 }
diff --git a/chunk_downloader.go b/chunk_downloader.go
index 626b5c4..9107184 100644
--- a/chunk_downloader.go
+++ b/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,


Test patch:
diff --git a/chunk_downloader_test.go b/chunk_downloader_test.go
new file mode 100644
index 0000000..5d55f9f
--- /dev/null
+++ b/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)
+		})
+	}
+}