[arrow-adbc] branch main updated: fix(go/adbc/driver/snowflake): handle result sets without Arrow data (#864)

[li...@apache.org]

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 <li...@gmail.com>
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 {