This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new f64da3ca fix(go/adbc/driver/snowflake): handle result sets without
Arrow data (#864)
f64da3ca is described below
commit f64da3ca953d2bc78e1390a1593c82ff9fdf319d
Author: David Li <[email protected]>
AuthorDate: Fri Jun 30 11:49:03 2023 -0400
fix(go/adbc/driver/snowflake): handle result sets without Arrow data (#864)
Prevents #863 from crashing, but doesn't fix the underlying error.
---
go/adbc/driver/snowflake/driver_test.go | 75 +++++++++++++++++++++++++++++--
go/adbc/driver/snowflake/record_reader.go | 19 ++++++++
2 files changed, 90 insertions(+), 4 deletions(-)
diff --git a/go/adbc/driver/snowflake/driver_test.go
b/go/adbc/driver/snowflake/driver_test.go
index 7ac1f27f..bdbf3791 100644
--- a/go/adbc/driver/snowflake/driver_test.go
+++ b/go/adbc/driver/snowflake/driver_test.go
@@ -247,7 +247,7 @@ func dropTempSchema(uri, schema string) {
}
}
-func TestADBCSnowflake(t *testing.T) {
+func withQuirks(t *testing.T, fn func(*SnowflakeQuirks)) {
uri := os.Getenv("SNOWFLAKE_URI")
if uri == "" {
@@ -258,7 +258,74 @@ func TestADBCSnowflake(t *testing.T) {
// dropping that schema when we're done.
q := &SnowflakeQuirks{dsn: uri, schemaName: createTempSchema(uri)}
defer dropTempSchema(uri, q.schemaName)
- suite.Run(t, &validation.DatabaseTests{Quirks: q})
- suite.Run(t, &validation.ConnectionTests{Quirks: q})
- suite.Run(t, &validation.StatementTests{Quirks: q})
+
+ fn(q)
+}
+
+func TestValidation(t *testing.T) {
+ withQuirks(t, func(q *SnowflakeQuirks) {
+ suite.Run(t, &validation.DatabaseTests{Quirks: q})
+ suite.Run(t, &validation.ConnectionTests{Quirks: q})
+ suite.Run(t, &validation.StatementTests{Quirks: q})
+ })
+}
+
+func TestSnowflake(t *testing.T) {
+ withQuirks(t, func(q *SnowflakeQuirks) {
+ suite.Run(t, &SnowflakeTests{Quirks: q})
+ })
+}
+
+// ---- Additional Tests --------------------
+
+type SnowflakeTests struct {
+ suite.Suite
+
+ Quirks *SnowflakeQuirks
+
+ ctx context.Context
+ driver adbc.Driver
+ db adbc.Database
+ cnxn adbc.Connection
+ stmt adbc.Statement
+}
+
+func (suite *SnowflakeTests) SetupSuite() {
+ var err error
+ suite.ctx = context.Background()
+ suite.driver = suite.Quirks.SetupDriver(suite.T())
+ suite.db, err = suite.driver.NewDatabase(suite.Quirks.DatabaseOptions())
+ suite.NoError(err)
+}
+
+func (suite *SnowflakeTests) SetupTest() {
+ var err error
+ suite.cnxn, err = suite.db.Open(suite.ctx)
+ suite.NoError(err)
+
+ suite.stmt, err = suite.cnxn.NewStatement()
+ suite.NoError(err)
+}
+
+func (suite *SnowflakeTests) TearDownTest() {
+ suite.NoError(suite.stmt.Close())
+ suite.NoError(suite.cnxn.Close())
+}
+
+func (suite *SnowflakeTests) TearDownSuite() {
+ suite.db = nil
+}
+
+func (suite *SnowflakeTests) TestStatementEmptyResultSet() {
+ // Regression test for https://github.com/apache/arrow-adbc/issues/863
+ suite.NoError(suite.stmt.SetSqlQuery("SHOW WAREHOUSES"))
+
+ // XXX: there IS data in this result set, but Snowflake doesn't
+ // appear to support getting the results as Arrow
+ _, _, err := suite.stmt.ExecuteQuery(suite.ctx)
+ var adbcErr adbc.Error
+ suite.ErrorAs(err, &adbcErr)
+
+ suite.Equal(adbc.StatusInternal, adbcErr.Code)
+ suite.Contains(adbcErr.Msg, "Cannot get Arrow data from this result
set")
}
diff --git a/go/adbc/driver/snowflake/record_reader.go
b/go/adbc/driver/snowflake/record_reader.go
index 32169bb7..081127d7 100644
--- a/go/adbc/driver/snowflake/record_reader.go
+++ b/go/adbc/driver/snowflake/record_reader.go
@@ -245,6 +245,25 @@ func newRecordReader(ctx context.Context, alloc
memory.Allocator, ld gosnowflake
return nil, errToAdbcErr(adbc.StatusInternal, err)
}
+ if len(batches) == 0 {
+ if ld.TotalRows() != 0 {
+ //
XXX(https://github.com/apache/arrow-adbc/issues/863): Snowflake won't return
Arrow data for certain queries
+ return nil, adbc.Error{
+ Msg: "[Snowflake] Cannot get Arrow data from
this result set (see apache/arrow-adbc#863)",
+ Code: adbc.StatusInternal,
+ }
+ }
+ schema := arrow.NewSchema([]arrow.Field{}, nil)
+ reader, err := array.NewRecordReader(schema, []arrow.Record{})
+ if err != nil {
+ return nil, adbc.Error{
+ Msg: err.Error(),
+ Code: adbc.StatusInternal,
+ }
+ }
+ return reader, nil
+ }
+
ch := make(chan arrow.Record, bufferSize)
r, err := batches[0].GetStream(ctx)
if err != nil {
