Repository: snowflakedb/gosnowflake
Pull request #1039 (2024-01-25 08:57:12)

SNOW-1016278 Fix panic on empty arrow batches
### Description
[SNOW-1016278](https://snowflakecomputing.atlassian.net/browse/SNOW-1016278) Fixes getting arrow batches when query returns no results at all.

### Checklist
- [ ] Code compiles correctly
- [ ] Run ``make fmt`` to fix inconsistent formats
- [ ] Run ``make lint`` to get lint errors and fix all of them
- [ ] Created tests which fail without the change (if possible)
- [ ] All tests passing
- [ ] Extended the README / documentation, if necessary


[SNOW-1016278]: https://snowflakecomputing.atlassian.net/browse/SNOW-1016278?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1039?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`611fe9d`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/611fe9ddb38e826016373c44a52ad76b5f8400a6?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) 83.12% compared to head [(`1b0b843`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1039?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) 81.60%.


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


```diff
@@            Coverage Diff             @@
##           master    #1039      +/-   ##
==========================================
- Coverage   83.12%   81.60%   -1.53%     
==========================================
  Files          50       49       -1     
  Lines        9759     9757       -2     
==========================================
- Hits         8112     7962     -150     
- Misses       1647     1795     +148     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1039?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/chunk_downloader.go b/chunk_downloader.go
index b68f9ec..a265f2f 100644
--- a/chunk_downloader.go
+++ b/chunk_downloader.go
@@ -249,7 +249,7 @@ func (scd *snowflakeChunkDownloader) getRowType() []execResponseRowType {
 }
 
 func (scd *snowflakeChunkDownloader) getArrowBatches() []*ArrowBatch {
-	if scd.FirstBatch.rec == nil {
+	if scd.FirstBatch == nil || scd.FirstBatch.rec == nil {
 		return scd.ArrowBatches
 	}
 	return append([]*ArrowBatch{scd.FirstBatch}, scd.ArrowBatches...)


Test patch:
diff --git a/assert_test.go b/assert_test.go
index b492136..91a69c3 100644
--- a/assert_test.go
+++ b/assert_test.go
@@ -58,6 +58,10 @@ func assertBetweenE(t *testing.T, value float64, min float64, max float64, descr
 	errorOnNonEmpty(t, validateValueBetween(value, min, max, descriptions...))
 }
 
+func assertEmptyE[T any](t *testing.T, actual []T, descriptions ...string) {
+	errorOnNonEmpty(t, validateEmpty(actual, descriptions...))
+}
+
 func fatalOnNonEmpty(t *testing.T, errMsg string) {
 	if errMsg != "" {
 		t.Fatal(formatErrorMessage(errMsg))
@@ -122,6 +126,14 @@ func validateValueBetween(value float64, min float64, max float64, descriptions
 	return fmt.Sprintf("expected \"%f\" should be between \"%f\" and  \"%f\" but did not. %s", value, min, max, desc)
 }
 
+func validateEmpty[T any](value []T, descriptions ...string) string {
+	if len(value) == 0 {
+		return ""
+	}
+	desc := joinDescriptions(descriptions...)
+	return fmt.Sprintf("expected \"%v\" to be empty. %s", value, desc)
+}
+
 func joinDescriptions(descriptions ...string) string {
 	return strings.Join(descriptions, " ")
 }
diff --git a/chunk_downloader_test.go b/chunk_downloader_test.go
index b63ef57..ea8192d 100644
--- a/chunk_downloader_test.go
+++ b/chunk_downloader_test.go
@@ -41,7 +41,7 @@ func TestWithArrowBatchesWhenQueryReturnsNoRowsWhenUsingNativeGoSQLInterface(t *
 	})
 }
 
-func TestWithArrowBatchesWhenQueryReturnsNoRows(t *testing.T) {
+func TestWithArrowBatchesWhenQueryReturnsRowsAndReadingRows(t *testing.T) {
 	runDBTest(t, func(dbt *DBTest) {
 		rows := dbt.mustQueryContext(WithArrowBatches(context.Background()), "SELECT 1")
 		defer rows.Close()
@@ -49,6 +49,30 @@ func TestWithArrowBatchesWhenQueryReturnsNoRows(t *testing.T) {
 	})
 }
 
+func TestWithArrowBatchesWhenQueryReturnsNoRowsAndReadingRows(t *testing.T) {
+	runDBTest(t, func(dbt *DBTest) {
+		rows := dbt.mustQueryContext(WithArrowBatches(context.Background()), "SELECT 1 WHERE 1 = 0")
+		defer rows.Close()
+		assertFalseF(t, rows.Next())
+	})
+}
+
+func TestWithArrowBatchesWhenQueryReturnsNoRowsAndReadingArrowBatches(t *testing.T) {
+	runDBTest(t, func(dbt *DBTest) {
+		var rows driver.Rows
+		var err error
+		err = dbt.conn.Raw(func(x any) error {
+			rows, err = x.(driver.QueryerContext).QueryContext(WithArrowBatches(context.Background()), "SELECT 1 WHERE 1 = 0", nil)
+			return err
+		})
+		assertNilF(t, err)
+		defer rows.Close()
+		batches, err := rows.(SnowflakeRows).GetArrowBatches()
+		assertNilF(t, err)
+		assertEmptyE(t, batches)
+	})
+}
+
 func TestWithArrowBatchesWhenQueryReturnsSomeRowsInGivenFormatUsingNativeGoSQLInterface(t *testing.T) {
 	for _, tc := range []struct {
 		useJSON bool