Skip to content

Commit

Permalink
SNOW-988793 Do not read arrow chunk when it is empty for arrow batches
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-pfus committed Dec 11, 2023
1 parent a4c3557 commit b90df21
Show file tree
Hide file tree
Showing 3 changed files with 91 additions and 0 deletions.
4 changes: 4 additions & 0 deletions assert_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,10 @@ func assertTrueF(t *testing.T, actual bool, descriptions ...string) {
fatalOnNonEmpty(t, validateEqual(actual, true, descriptions...))
}

func assertFalseF(t *testing.T, actual bool, descriptions ...string) {
fatalOnNonEmpty(t, validateEqual(actual, false, descriptions...))
}

func assertStringContainsE(t *testing.T, actual string, expectedToContain string, descriptions ...string) {
errorOnNonEmpty(t, validateStringContains(actual, expectedToContain, descriptions...))
}
Expand Down
3 changes: 3 additions & 0 deletions chunk_downloader.go
Original file line number Diff line number Diff line change
Expand Up @@ -271,6 +271,9 @@ func getChunk(
}

func (scd *snowflakeChunkDownloader) startArrowBatches() error {
if scd.RowSet.RowSetBase64 == "" {
return nil
}
var err error
chunkMetaLen := len(scd.ChunkMetas)
var loc *time.Location
Expand Down
84 changes: 84 additions & 0 deletions chunk_downloader_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package gosnowflake

import (
"context"
"database/sql/driver"
"testing"
)

Expand All @@ -26,3 +27,86 @@ func TestChunkDownloaderDoesNotStartWhenArrowParsingCausesError(t *testing.T) {
})
}
}

func TestWithArrowBatchesWhenQueryReturnsNoRowsWhenUsingNativeGoSQLInterface(t *testing.T) {
runDBTest(t, func(dbt *DBTest) {
var rows driver.Rows
var err error
err = dbt.conn.Raw(func(x interface{}) error {
rows, err = x.(driver.QueryerContext).QueryContext(WithArrowBatches(context.Background()), "SELECT 1 WHERE 0 = 1", nil)
return err
})
assertNilF(t, err)
defer rows.Close()
})
}

func TestWithArrowBatchesWhenQueryReturnsNoRows(t *testing.T) {
runDBTest(t, func(dbt *DBTest) {
rows := dbt.mustQueryContext(WithArrowBatches(context.Background()), "SELECT 1")
defer rows.Close()
assertFalseF(t, rows.Next())
})
}

func TestWithArrowBatchesWhenQueryReturnsSomeRowsInGivenFormatUsingNativeGoSQLInterface(t *testing.T) {
for _, tc := range []struct {
useJson bool

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.19 on Ubuntu

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.19 on Ubuntu

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.20 on Ubuntu

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.20 on Ubuntu

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.20 on Ubuntu

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.19 on Mac

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.19 on Ubuntu

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.20 on Mac

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.20 on Mac

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.19 on Mac

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.20 on Mac

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.19 on Mac

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.20 on Windows

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.20 on Windows

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.19 on Windows

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.20 on Windows

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.19 on Windows

struct field useJson should be useJSON

Check failure on line 54 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.19 on Windows

struct field useJson should be useJSON
desc string
}{
{
useJson: true,
desc: "json",
},
{
useJson: false,
desc: "arrow",
},
} {
t.Run(tc.desc, func(t *testing.T) {
runDBTest(t, func(dbt *DBTest) {
if tc.useJson {
dbt.mustExec(forceJSON)
}
var rows driver.Rows
var err error
err = dbt.conn.Raw(func(x interface{}) error {
rows, err = x.(driver.QueryerContext).QueryContext(WithArrowBatches(context.Background()), "SELECT 1", nil)
return err
})
assertNilF(t, err)
defer rows.Close()
values := make([]driver.Value, 1)
rows.Next(values)
assertEqualE(t, values[0], nil)
})
})
}
}

func TestWithArrowBatchesWhenQueryReturnsSomeRowsInGivenFormat(t *testing.T) {
for _, tc := range []struct {
useJson bool

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.19 on Ubuntu

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.19 on Ubuntu

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.20 on Ubuntu

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.20 on Ubuntu

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.20 on Ubuntu

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.19 on Mac

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.19 on Ubuntu

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.20 on Mac

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.20 on Mac

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.19 on Mac

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.20 on Mac

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.19 on Mac

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.20 on Windows

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.20 on Windows

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.19 on Windows

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.20 on Windows

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.19 on Windows

struct field useJson should be useJSON

Check failure on line 89 in chunk_downloader_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.19 on Windows

struct field useJson should be useJSON
desc string
}{
{
useJson: true,
desc: "json",
},
{
useJson: false,
desc: "arrow",
},
} {
t.Run(tc.desc, func(t *testing.T) {
runDBTest(t, func(dbt *DBTest) {
if tc.useJson {
dbt.mustExec(forceJSON)
}
rows := dbt.mustQueryContext(WithArrowBatches(context.Background()), "SELECT 1")
defer rows.Close()
assertFalseF(t, rows.Next())
})
})
}
}

0 comments on commit b90df21

Please sign in to comment.