diff --git a/go/adbc/driver/snowflake/driver_test.go b/go/adbc/driver/snowflake/driver_test.go index 2351d9fcc5..496629ebda 100644 --- a/go/adbc/driver/snowflake/driver_test.go +++ b/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,7 +236,7 @@ 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) @@ -244,7 +244,7 @@ func createTempSchema(uri string) string { 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 diff --git a/go/adbc/driver/snowflake/statement.go b/go/adbc/driver/snowflake/statement.go index 755bf3ede7..57ac079d44 100644 --- a/go/adbc/driver/snowflake/statement.go +++ b/go/adbc/driver/snowflake/statement.go @@ -584,6 +584,37 @@ func (st *statement) ExecuteUpdate(ctx context.Context) (int64, error) { return n, nil } +// ExecuteSchema gets the schema of the result set of a query without executing it. +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 { diff --git a/go/adbc/go.mod b/go/adbc/go.mod index 6eea776dbb..45f2ee2daf 100644 --- a/go/adbc/go.mod +++ b/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 diff --git a/go/adbc/go.sum b/go/adbc/go.sum index 04a61c2880..acc7fae809 100644 --- a/go/adbc/go.sum +++ b/go/adbc/go.sum @@ -134,8 +134,8 @@ github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qq github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= 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=