Support Snowflake ExecuteSchema

go/adbc/driver/snowflake/driver_test.go

@@ -194,7 +194,7 @@ func (s *SnowflakeQuirks) BindParameter(_ int) string            { return "?" }
 func (s *SnowflakeQuirks) SupportsBulkIngest(string) bool        { return true }
 func (s *SnowflakeQuirks) SupportsConcurrentStatements() bool    { return true }
 func (s *SnowflakeQuirks) SupportsCurrentCatalogSchema() bool    { return true }
-func (s *SnowflakeQuirks) SupportsExecuteSchema() bool           { return false }
+func (s *SnowflakeQuirks) SupportsExecuteSchema() bool           { return true }
 func (s *SnowflakeQuirks) SupportsGetSetOptions() bool           { return true }
 func (s *SnowflakeQuirks) SupportsPartitionedData() bool         { return false }
 func (s *SnowflakeQuirks) SupportsStatistics() bool              { return false }
@@ -236,15 +236,15 @@ func (s *SnowflakeQuirks) SampleTableSchemaMetadata(tblName string, dt arrow.Dat
 	return arrow.Metadata{}
 }
 
-func createTempSchema(uri string) string {
+func createTempSchema(database string, uri string) string {
 	db, err := sql.Open("snowflake", uri)
 	if err != nil {
 		panic(err)
 	}
 	defer db.Close()
 
 	schemaName := strings.ToUpper("ADBC_TESTING_" + strings.ReplaceAll(uuid.New().String(), "-", "_"))
-	_, err = db.Exec(`CREATE SCHEMA ADBC_TESTING.` + schemaName)
+	_, err = db.Exec(`CREATE SCHEMA ` + database + `.` + schemaName)
 	if err != nil {
 		panic(err)
 	}
@@ -277,7 +277,7 @@ func withQuirks(t *testing.T, fn func(*SnowflakeQuirks)) {
 
 	// avoid multiple runs clashing by operating in a fresh schema and then
 	// dropping that schema when we're done.
-	q := &SnowflakeQuirks{dsn: uri, catalogName: database, schemaName: createTempSchema(uri)}
+	q := &SnowflakeQuirks{dsn: uri, catalogName: database, schemaName: createTempSchema(database, uri)}
 	defer dropTempSchema(uri, q.schemaName)
 
 	fn(q)
@@ -679,6 +679,17 @@ func (suite *SnowflakeTests) TestUseHighPrecision() {
 	suite.Equal(9876543210.99, rec.Column(1).(*array.Float64).Value(1))
 }
 
+func (suite *SnowflakeTests) TestDescribeOnly() {
+	suite.Require().NoError(suite.stmt.SetOption(driver.OptionUseHighPrecision, adbc.OptionValueEnabled))
+	suite.Require().NoError(suite.stmt.SetSqlQuery("SELECT CAST('9999.99' AS NUMBER(6, 2)) AS RESULT"))
+	schema, err := suite.stmt.(adbc.StatementExecuteSchema).ExecuteSchema(suite.ctx)
+	suite.Require().NoError(err)
+
+	suite.Equal(1, len(schema.Fields()))
+	suite.Equal("RESULT", schema.Field(0).Name)
+	suite.Truef(arrow.TypeEqual(&arrow.Decimal128Type{Precision: 6, Scale: 2}, schema.Field(0).Type), "expected decimal(6, 2), got %s", schema.Field(0).Type)
+}
+
 func TestJwtAuthenticationUnencryptedValue(t *testing.T) {
 	// test doesn't participate in SnowflakeTests because
 	// JWT auth has a different behavior

go/adbc/driver/snowflake/statement.go

@@ -584,6 +584,41 @@ func (st *statement) ExecuteUpdate(ctx context.Context) (int64, error) {
 	return n, nil
 }
 
+// ExecuteQuery executes the current query or prepared statement
+// and returnes a RecordReader for the results along with the number
+// of rows affected if known, otherwise it will be -1.
+//
+// This invalidates any prior result sets on this statement.
+func (st *statement) ExecuteSchema(ctx context.Context) (*arrow.Schema, error) {
+	if st.targetTable != "" {
+		return nil, adbc.Error{
+			Msg:  "cannot execute schema for ingestion",
+			Code: adbc.StatusInvalidState,
+		}
+	}
+
+	if st.query == "" {
+		return nil, adbc.Error{
+			Msg:  "cannot execute without a query",
+			Code: adbc.StatusInvalidState,
+		}
+	}
+
+	if st.streamBind != nil || st.bound != nil {
+		return nil, adbc.Error{
+			Msg:  "executing schema with bound params not yet implemented",
+			Code: adbc.StatusNotImplemented,
+		}
+	}
+
+	loader, err := st.cnxn.cn.QueryArrowStream(gosnowflake.WithDescribeOnly(ctx), st.query)
+	if err != nil {
+		return nil, errToAdbcErr(adbc.StatusInternal, err)
+	}
+
+	return rowTypesToArrowSchema(ctx, loader, st.useHighPrecision)
+}
+
 // Prepare turns this statement into a prepared statement to be executed
 // multiple times. This invalidates any prior result sets.
 func (st *statement) Prepare(_ context.Context) error {

go/adbc/go.mod

@@ -104,4 +104,4 @@ require (
 	modernc.org/token v1.1.0 // indirect
 )
 
-replace github.com/snowflakedb/gosnowflake => github.com/snowflakedb/gosnowflake v1.6.23-0.20230717195239-fec38ba82d2a
+replace github.com/snowflakedb/gosnowflake => github.com/snowflakedb/gosnowflake v1.6.23-0.20231106142408-8445dcaf46d9

go/adbc/go.sum

@@ -136,6 +136,8 @@ github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ
 github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
 github.com/snowflakedb/gosnowflake v1.6.23-0.20230717195239-fec38ba82d2a h1:F7fKVj3t12jr3Bopzngsp/PZDm1or8zpk+29NN4YFGk=
 github.com/snowflakedb/gosnowflake v1.6.23-0.20230717195239-fec38ba82d2a/go.mod h1:KfO4F7bk+aXPUIvBqYxvPhxLlu2/w4TtSC8Rw/yr5Mg=
+github.com/snowflakedb/gosnowflake v1.6.23-0.20231106142408-8445dcaf46d9 h1:TCIqFsNrcVtups6vBXX2DHI+SjcXw22k0yf1JqzeZAM=
+github.com/snowflakedb/gosnowflake v1.6.23-0.20231106142408-8445dcaf46d9/go.mod h1:KfO4F7bk+aXPUIvBqYxvPhxLlu2/w4TtSC8Rw/yr5Mg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=