This is an automated email from the ASF dual-hosted git repository.
zeroshade 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 28a87ea15 feat(go/adbc/driver/flightsql): allow passing arbitrary grpc
dial options in NewDatabase (#2563)
28a87ea15 is described below
commit 28a87ea15f51f5d4e871a7467dec6c26d415152a
Author: Marius van Niekerk <[email protected]>
AuthorDate: Thu Feb 27 15:00:15 2025 -0500
feat(go/adbc/driver/flightsql): allow passing arbitrary grpc dial options
in NewDatabase (#2563)
Add a new database constructor NewDatabaseWithOptions to allow passing
arbitrary user-specified grpc dial options.
This is useful for constructs like
```go
driver := flightsql.NewDriver(memory.DefaultAllocator)
driver.NewDatabaseWithOptions(map[string]string{
"uri": uri,
},
grpc.WithStatsHandler(otelgrpc.NewClientHandler())
)
```
which allows usage of the opentelemetry grpc instrumentation for
example.
This also provides an escape valve for users that need the ability to
use less commonly used grpc client features that we have not exposed via
the string map args.
---
go/adbc/adbc.go | 9 ++--
.../driver/flightsql/flightsql_adbc_server_test.go | 60 +++++++++++++++++++++-
go/adbc/driver/flightsql/flightsql_database.go | 2 +
go/adbc/driver/flightsql/flightsql_driver.go | 29 +++++++++--
4 files changed, 89 insertions(+), 11 deletions(-)
diff --git a/go/adbc/adbc.go b/go/adbc/adbc.go
index 9732bb0e1..1f178e953 100644
--- a/go/adbc/adbc.go
+++ b/go/adbc/adbc.go
@@ -310,20 +310,19 @@ const (
)
// Driver is the entry point for the interface. It is similar to
-// database/sql.Driver taking a map of keys and values as options
-// to initialize a Connection to the database. Any common connection
+// [database/sql.Driver] taking a map of keys and values as options
+// to initialize a [Connection] to the database. Any common connection
// state can live in the Driver itself, for example an in-memory database
// can place ownership of the actual database in this driver.
//
// Any connection specific options should be set using SetOptions before
// calling Open.
//
-// The provided context.Context is for dialing purposes only
-// (see net.DialContext) and should not be stored or used for other purposes.
+// The provided [context.Context] is for dialing purposes only.
// A default timeout should still be used when dialing as a connection
// pool may call Connect asynchronously to any query.
//
-// A driver can also optionally implement io.Closer if there is a need
+// A driver can also optionally implement [io.Closer] if there is a need
// or desire for it.
type Driver interface {
NewDatabase(opts map[string]string) (Database, error)
diff --git a/go/adbc/driver/flightsql/flightsql_adbc_server_test.go
b/go/adbc/driver/flightsql/flightsql_adbc_server_test.go
index e3c132183..c8a59d72a 100644
--- a/go/adbc/driver/flightsql/flightsql_adbc_server_test.go
+++ b/go/adbc/driver/flightsql/flightsql_adbc_server_test.go
@@ -51,6 +51,7 @@ import (
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
+ "google.golang.org/grpc/stats"
"google.golang.org/grpc/status"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/anypb"
@@ -67,7 +68,7 @@ type ServerBasedTests struct {
cnxn adbc.Connection
}
-func (suite *ServerBasedTests) DoSetupSuite(srv flightsql.Server,
srvMiddleware []flight.ServerMiddleware, dbArgs map[string]string) {
+func (suite *ServerBasedTests) DoSetupSuite(srv flightsql.Server,
srvMiddleware []flight.ServerMiddleware, dbArgs map[string]string, dialOpts
...grpc.DialOption) {
suite.s = flight.NewServerWithMiddleware(srvMiddleware)
suite.s.RegisterFlightService(flightsql.NewFlightServer(srv))
suite.Require().NoError(suite.s.Init("localhost:0"))
@@ -83,7 +84,7 @@ func (suite *ServerBasedTests) DoSetupSuite(srv
flightsql.Server, srvMiddleware
"uri": uri,
}
maps.Copy(args, dbArgs)
- suite.db, err =
(driver.NewDriver(memory.DefaultAllocator)).NewDatabase(args)
+ suite.db, err =
(driver.NewDriver(memory.DefaultAllocator)).NewDatabaseWithOptions(args,
dialOpts...)
suite.Require().NoError(err)
}
@@ -109,6 +110,10 @@ func TestAuthn(t *testing.T) {
suite.Run(t, &AuthnTests{})
}
+func TestGrpcDialerOptions(t *testing.T) {
+ suite.Run(t, &DialerOptionsTests{})
+}
+
func TestErrorDetails(t *testing.T) {
suite.Run(t, &ErrorDetailsTests{})
}
@@ -244,6 +249,57 @@ func (suite *AuthnTests) TestBearerTokenUpdated() {
defer reader.Release()
}
+// ---- Grpc Dialer Options Tests --------------
+
+type DialerOptionsTests struct {
+ ServerBasedTests
+ statsHandler *dialerOptionsGrpcStatsHandler
+}
+
+type dialerOptionsGrpcStatsHandler struct {
+ connectionsHandled int
+ rpcsHandled int
+ connectionsTagged int
+ rpcsTagged int
+}
+
+func (d *dialerOptionsGrpcStatsHandler) HandleConn(ctx context.Context, stat
stats.ConnStats) {
+ d.connectionsHandled++
+}
+func (d *dialerOptionsGrpcStatsHandler) HandleRPC(ctx context.Context, stat
stats.RPCStats) {
+ d.rpcsHandled++
+}
+func (d *dialerOptionsGrpcStatsHandler) TagConn(ctx context.Context, stat
*stats.ConnTagInfo) context.Context {
+ d.connectionsTagged++
+ return ctx
+}
+func (d *dialerOptionsGrpcStatsHandler) TagRPC(ctx context.Context, stat
*stats.RPCTagInfo) context.Context {
+ d.rpcsTagged++
+ return ctx
+}
+
+func (suite *DialerOptionsTests) SetupSuite() {
+ suite.statsHandler = &dialerOptionsGrpcStatsHandler{}
+ suite.DoSetupSuite(&AuthnTestServer{}, nil, nil,
grpc.WithStatsHandler(suite.statsHandler))
+}
+
+// TestGrpcObserved validates that the grpc stats handler that was passed
through correctly to the underlying grpc client.
+func (suite *DialerOptionsTests) TestGrpcObserved() {
+ stmt, err := suite.cnxn.NewStatement()
+ suite.Require().NoError(err)
+ defer stmt.Close()
+
+ suite.Require().NoError(stmt.SetSqlQuery("timeout"))
+ reader, _, err := stmt.ExecuteQuery(context.Background())
+ suite.NoError(err)
+ defer reader.Release()
+
+ suite.Less(0, suite.statsHandler.connectionsTagged)
+ suite.Less(0, suite.statsHandler.connectionsHandled)
+ suite.Less(0, suite.statsHandler.rpcsTagged)
+ suite.Less(0, suite.statsHandler.rpcsHandled)
+}
+
// ---- Error Details Tests --------------------
type ErrorDetailsTestServer struct {
diff --git a/go/adbc/driver/flightsql/flightsql_database.go
b/go/adbc/driver/flightsql/flightsql_database.go
index a25c839ec..bbbcbbf06 100644
--- a/go/adbc/driver/flightsql/flightsql_database.go
+++ b/go/adbc/driver/flightsql/flightsql_database.go
@@ -67,6 +67,7 @@ type databaseImpl struct {
dialOpts dbDialOpts
enableCookies bool
options map[string]string
+ userDialOpts []grpc.DialOption
}
func (d *databaseImpl) SetOptions(cnOptions map[string]string) error {
@@ -371,6 +372,7 @@ func getFlightClient(ctx context.Context, loc string, d
*databaseImpl, authMiddl
dv, _ :=
d.DatabaseImplBase.DriverInfo.GetInfoForInfoCode(adbc.InfoDriverVersion)
driverVersion := dv.(string)
dialOpts := append(d.dialOpts.opts,
grpc.WithConnectParams(d.timeout.connectParams()),
grpc.WithTransportCredentials(creds), grpc.WithUserAgent("ADBC Flight SQL
Driver "+driverVersion))
+ dialOpts = append(dialOpts, d.userDialOpts...)
d.Logger.DebugContext(ctx, "new client", "location", loc)
cl, err := flightsql.NewClient(target, nil, middleware, dialOpts...)
diff --git a/go/adbc/driver/flightsql/flightsql_driver.go
b/go/adbc/driver/flightsql/flightsql_driver.go
index bf4e0b4a8..9e517c716 100644
--- a/go/adbc/driver/flightsql/flightsql_driver.go
+++ b/go/adbc/driver/flightsql/flightsql_driver.go
@@ -39,6 +39,7 @@ import (
"github.com/apache/arrow-adbc/go/adbc/driver/internal/driverbase"
"github.com/apache/arrow-go/v18/arrow/memory"
"golang.org/x/exp/maps"
+ "google.golang.org/grpc"
"google.golang.org/grpc/metadata"
)
@@ -76,13 +77,27 @@ type driverImpl struct {
driverbase.DriverImplBase
}
+// Driver is the extended [adbc.Driver] interface for Flight SQL.
+//
+// It adds an additional method to create a database with grpc specific
options that cannot be
+// passed through the options map.
+type Driver interface {
+ adbc.Driver
+ NewDatabaseWithOptions(map[string]string, ...grpc.DialOption)
(adbc.Database, error)
+}
+
// NewDriver creates a new Flight SQL driver using the given Arrow allocator.
-func NewDriver(alloc memory.Allocator) adbc.Driver {
+func NewDriver(alloc memory.Allocator) Driver {
info := driverbase.DefaultDriverInfo("Flight SQL")
- return driverbase.NewDriver(&driverImpl{DriverImplBase:
driverbase.NewDriverImplBase(info, alloc)})
+ return &driverImpl{DriverImplBase: driverbase.NewDriverImplBase(info,
alloc)}
}
-func (d *driverImpl) NewDatabase(opts map[string]string) (adbc.Database,
error) {
+// NewDatabase creates a new Flight SQL database using the given options.
+//
+// Additional grpc client options can can be passed as grpc.DialOption.
+// This enables the use of additional grpc client options not directly exposed
by the options map.
+// such as grpc.WithStatsHandler() for enabling various telemetry handlers.
+func (d *driverImpl) NewDatabaseWithOptions(opts map[string]string,
userDialOpts ...grpc.DialOption) (adbc.Database, error) {
opts = maps.Clone(opts)
uri, ok := opts[adbc.OptionKeyURI]
if !ok {
@@ -99,7 +114,8 @@ func (d *driverImpl) NewDatabase(opts map[string]string)
(adbc.Database, error)
// Match gRPC default
connectTimeout: time.Second * 20,
},
- hdrs: make(metadata.MD),
+ hdrs: make(metadata.MD),
+ userDialOpts: userDialOpts,
}
var err error
@@ -118,3 +134,8 @@ func (d *driverImpl) NewDatabase(opts map[string]string)
(adbc.Database, error)
return driverbase.NewDatabase(db), nil
}
+
+// NewDatabase creates a new Flight SQL database using the given options.
+func (d *driverImpl) NewDatabase(opts map[string]string) (adbc.Database,
error) {
+ return d.NewDatabaseWithOptions(opts)
+}
