lidavidm commented on code in PR #3325: URL: https://github.com/apache/arrow-adbc/pull/3325#discussion_r2332520238
########## go/adbc/driver/databricks/cloudfetch_e2e_test.go: ########## @@ -0,0 +1,497 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//go:build integration +// +build integration Review Comment: You only need the first directive. ########## go/adbc/driver/databricks/ipc_reader_adapter.go: ########## @@ -0,0 +1,199 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package databricks + +import ( + "bytes" + "context" + "fmt" + "io" + "sync/atomic" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/ipc" + dbsqlrows "github.com/databricks/databricks-sql-go/rows" +) + +// Check if the rows interface supports IPC streams +type rowsWithIPCStream interface { + GetArrowIPCStreams(context.Context) (dbsqlrows.ArrowIPCStreamIterator, error) +} + +// ipcReaderAdapter uses the new IPC stream interface for Arrow access +type ipcReaderAdapter struct { + ipcIterator dbsqlrows.ArrowIPCStreamIterator + currentReader *ipc.Reader + currentRecord arrow.Record + schema *arrow.Schema + closed bool + refCount int64 +} + +// newIPCReaderAdapter creates a RecordReader using direct IPC stream access +func newIPCReaderAdapter(ctx context.Context, rows dbsqlrows.Rows) (array.RecordReader, error) { + // Check if rows supports IPC streams + ipcRows, ok := rows.(rowsWithIPCStream) + if !ok { + return nil, fmt.Errorf("databricks rows do not support IPC stream access") Review Comment: Would this ever happen? ########## go/adbc/driver/databricks/statement.go: ########## @@ -0,0 +1,225 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package databricks + +import ( + "context" + "database/sql" + "database/sql/driver" + "fmt" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + + dbsqlrows "github.com/databricks/databricks-sql-go/rows" +) + +type statementImpl struct { + conn *connectionImpl + query string + parameters []interface{} + prepared *sql.Stmt +} + +func (s *statementImpl) Close() error { + if s.conn == nil { + return adbc.Error{ + Msg: "statement already closed", + Code: adbc.StatusInvalidState, + } + } + if s.prepared != nil { + return s.prepared.Close() + } + return nil +} + +func (s *statementImpl) SetOption(key, val string) error { + // No statement-specific options are supported yet + return adbc.Error{ + Code: adbc.StatusNotImplemented, + Msg: fmt.Sprintf("unsupported statement option: %s", key), + } +} + +func (s *statementImpl) SetSqlQuery(query string) error { + s.query = query + // Reset prepared statement if query changes + if s.prepared != nil { + if err := s.prepared.Close(); err != nil { + return adbc.Error{ + Code: adbc.StatusInvalidState, + Msg: fmt.Sprintf("failed to close previous prepared statement: %w", err), + } + } + s.prepared = nil + } + return nil +} + +func (s *statementImpl) Prepare(ctx context.Context) error { + if s.query == "" { + return adbc.Error{ + Code: adbc.StatusInvalidState, + Msg: "no query set", + } + } + + stmt, err := s.conn.conn.PrepareContext(ctx, s.query) + if err != nil { + return adbc.Error{ + Code: adbc.StatusInvalidState, + Msg: fmt.Sprintf("failed to prepare statement: %v", err), + } + } + + s.prepared = stmt + return nil +} + +func (s *statementImpl) ExecuteQuery(ctx context.Context) (array.RecordReader, int64, error) { + if s.query == "" { + return nil, -1, adbc.Error{ + Code: adbc.StatusInvalidState, + Msg: "no query set", + } + } + + // Execute query using raw driver interface to get Arrow batches + var driverRows driver.Rows + var err error + err = s.conn.conn.Raw(func(driverConn interface{}) error { + // Use raw driver interface for direct Arrow access + if queryerCtx, ok := driverConn.(driver.QueryerContext); ok { + // Convert parameters to driver.NamedValue slice + var driverArgs []driver.NamedValue + for i, param := range s.parameters { + driverArgs = append(driverArgs, driver.NamedValue{ + Ordinal: i + 1, + Value: param, + }) + } + driverRows, err = queryerCtx.QueryContext(ctx, s.query, driverArgs) + return err + } + return fmt.Errorf("driver does not support QueryerContext") Review Comment: Would this ever happen? ########## go/adbc/driver/databricks/ipc_reader_adapter.go: ########## @@ -0,0 +1,199 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package databricks + +import ( + "bytes" + "context" + "fmt" + "io" + "sync/atomic" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/ipc" + dbsqlrows "github.com/databricks/databricks-sql-go/rows" +) + +// Check if the rows interface supports IPC streams +type rowsWithIPCStream interface { + GetArrowIPCStreams(context.Context) (dbsqlrows.ArrowIPCStreamIterator, error) +} + +// ipcReaderAdapter uses the new IPC stream interface for Arrow access +type ipcReaderAdapter struct { + ipcIterator dbsqlrows.ArrowIPCStreamIterator + currentReader *ipc.Reader + currentRecord arrow.Record + schema *arrow.Schema + closed bool + refCount int64 +} + +// newIPCReaderAdapter creates a RecordReader using direct IPC stream access +func newIPCReaderAdapter(ctx context.Context, rows dbsqlrows.Rows) (array.RecordReader, error) { + // Check if rows supports IPC streams + ipcRows, ok := rows.(rowsWithIPCStream) + if !ok { + return nil, fmt.Errorf("databricks rows do not support IPC stream access") + } + + // Get IPC stream iterator + ipcIterator, err := ipcRows.GetArrowIPCStreams(ctx) + if err != nil { + return nil, fmt.Errorf("failed to get IPC streams: %w", err) Review Comment: Ideally we return `adbc.Error`, in general, so that the Python client etc. can provide semantic errors ########## go/adbc/driver/databricks/statement.go: ########## @@ -0,0 +1,225 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package databricks + +import ( + "context" + "database/sql" + "database/sql/driver" + "fmt" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + + dbsqlrows "github.com/databricks/databricks-sql-go/rows" +) + +type statementImpl struct { + conn *connectionImpl + query string + parameters []interface{} + prepared *sql.Stmt +} + +func (s *statementImpl) Close() error { + if s.conn == nil { + return adbc.Error{ + Msg: "statement already closed", + Code: adbc.StatusInvalidState, + } + } + if s.prepared != nil { + return s.prepared.Close() + } + return nil +} + +func (s *statementImpl) SetOption(key, val string) error { + // No statement-specific options are supported yet + return adbc.Error{ + Code: adbc.StatusNotImplemented, + Msg: fmt.Sprintf("unsupported statement option: %s", key), + } +} + +func (s *statementImpl) SetSqlQuery(query string) error { + s.query = query + // Reset prepared statement if query changes + if s.prepared != nil { + if err := s.prepared.Close(); err != nil { + return adbc.Error{ + Code: adbc.StatusInvalidState, + Msg: fmt.Sprintf("failed to close previous prepared statement: %w", err), + } + } + s.prepared = nil + } + return nil +} + +func (s *statementImpl) Prepare(ctx context.Context) error { + if s.query == "" { + return adbc.Error{ + Code: adbc.StatusInvalidState, + Msg: "no query set", + } + } + + stmt, err := s.conn.conn.PrepareContext(ctx, s.query) + if err != nil { + return adbc.Error{ + Code: adbc.StatusInvalidState, + Msg: fmt.Sprintf("failed to prepare statement: %v", err), + } + } + + s.prepared = stmt + return nil +} + +func (s *statementImpl) ExecuteQuery(ctx context.Context) (array.RecordReader, int64, error) { Review Comment: We don't use the prepared statement in here? ########## go/adbc/driver/databricks/statement.go: ########## @@ -0,0 +1,225 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package databricks + +import ( + "context" + "database/sql" + "database/sql/driver" + "fmt" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + + dbsqlrows "github.com/databricks/databricks-sql-go/rows" +) + +type statementImpl struct { + conn *connectionImpl + query string + parameters []interface{} + prepared *sql.Stmt +} + +func (s *statementImpl) Close() error { + if s.conn == nil { + return adbc.Error{ + Msg: "statement already closed", + Code: adbc.StatusInvalidState, + } + } Review Comment: If you're making this check, did you mean to set `s.conn` to nil at the end? ########## go/adbc/driver/databricks/statement_test.go: ########## @@ -0,0 +1,46 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package databricks_test + +import ( + "testing" + + "github.com/apache/arrow-adbc/go/adbc/driver/databricks" + "github.com/apache/arrow-adbc/go/adbc/validation" + "github.com/stretchr/testify/assert" +) + +func TestStatementBasic(t *testing.T) { + // This is a basic test to ensure the code compiles + // Real tests would require a connection to Databricks + + // Create a driver and database + driver := databricks.NewDriver(nil) + db, err := driver.NewDatabase(map[string]string{ + databricks.OptionServerHostname: "mock-host", + databricks.OptionAccessToken: "mock-token", + databricks.OptionHTTPPath: "mock-path", + }) + assert.NoError(t, err) + + defer validation.CheckedClose(t, db) + + // Note: We can't test the actual statement implementation without a real connection + // This test just ensures the public API compiles correctly Review Comment: Also I would encourage you to implement the test suite (again, see the Snowflake or BigQuery tests) which will exercise various bits of the driver. Since they won't run in CI it's OK if they don't all pass but it at least gives something to work towards. ########## go/adbc/driver/databricks/cloudfetch_e2e_test.go: ########## @@ -0,0 +1,497 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//go:build integration +// +build integration Review Comment: But as mentioned I would rather gate on the presence of an env var and skip the tests otherwise. (As is, these tests won't even get compiled.) ########## go/adbc/driver/databricks/statement.go: ########## @@ -0,0 +1,225 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package databricks + +import ( + "context" + "database/sql" + "database/sql/driver" + "fmt" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + + dbsqlrows "github.com/databricks/databricks-sql-go/rows" +) + +type statementImpl struct { + conn *connectionImpl + query string + parameters []interface{} + prepared *sql.Stmt +} + +func (s *statementImpl) Close() error { + if s.conn == nil { + return adbc.Error{ + Msg: "statement already closed", + Code: adbc.StatusInvalidState, + } + } + if s.prepared != nil { + return s.prepared.Close() + } + return nil +} + +func (s *statementImpl) SetOption(key, val string) error { + // No statement-specific options are supported yet + return adbc.Error{ + Code: adbc.StatusNotImplemented, + Msg: fmt.Sprintf("unsupported statement option: %s", key), + } +} + +func (s *statementImpl) SetSqlQuery(query string) error { + s.query = query + // Reset prepared statement if query changes + if s.prepared != nil { + if err := s.prepared.Close(); err != nil { + return adbc.Error{ + Code: adbc.StatusInvalidState, + Msg: fmt.Sprintf("failed to close previous prepared statement: %w", err), + } + } + s.prepared = nil + } + return nil +} + +func (s *statementImpl) Prepare(ctx context.Context) error { + if s.query == "" { + return adbc.Error{ + Code: adbc.StatusInvalidState, + Msg: "no query set", + } + } + + stmt, err := s.conn.conn.PrepareContext(ctx, s.query) + if err != nil { + return adbc.Error{ + Code: adbc.StatusInvalidState, + Msg: fmt.Sprintf("failed to prepare statement: %v", err), + } + } + + s.prepared = stmt + return nil +} + +func (s *statementImpl) ExecuteQuery(ctx context.Context) (array.RecordReader, int64, error) { + if s.query == "" { + return nil, -1, adbc.Error{ + Code: adbc.StatusInvalidState, + Msg: "no query set", + } + } + + // Execute query using raw driver interface to get Arrow batches + var driverRows driver.Rows + var err error + err = s.conn.conn.Raw(func(driverConn interface{}) error { + // Use raw driver interface for direct Arrow access + if queryerCtx, ok := driverConn.(driver.QueryerContext); ok { + // Convert parameters to driver.NamedValue slice + var driverArgs []driver.NamedValue + for i, param := range s.parameters { + driverArgs = append(driverArgs, driver.NamedValue{ + Ordinal: i + 1, + Value: param, + }) + } + driverRows, err = queryerCtx.QueryContext(ctx, s.query, driverArgs) + return err + } + return fmt.Errorf("driver does not support QueryerContext") + }) + + if err != nil { + return nil, -1, adbc.Error{ + Code: adbc.StatusInternal, + Msg: fmt.Sprintf("failed to execute query: %v", err), + } + } + defer func() { _ = driverRows.Close() }() Review Comment: Please don't swallow errors ########## go/adbc/driver/databricks/statement_test.go: ########## @@ -0,0 +1,46 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package databricks_test + +import ( + "testing" + + "github.com/apache/arrow-adbc/go/adbc/driver/databricks" + "github.com/apache/arrow-adbc/go/adbc/validation" + "github.com/stretchr/testify/assert" +) + +func TestStatementBasic(t *testing.T) { + // This is a basic test to ensure the code compiles + // Real tests would require a connection to Databricks + + // Create a driver and database + driver := databricks.NewDriver(nil) + db, err := driver.NewDatabase(map[string]string{ + databricks.OptionServerHostname: "mock-host", + databricks.OptionAccessToken: "mock-token", + databricks.OptionHTTPPath: "mock-path", + }) + assert.NoError(t, err) + + defer validation.CheckedClose(t, db) + + // Note: We can't test the actual statement implementation without a real connection + // This test just ensures the public API compiles correctly Review Comment: If you look at the other tests (e.g. for Snowflake): you can set up real tests and simply gate them on an environment variable, and just skip them if not present. That way we can write nontrivial tests still. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
