From 28a87ea15f51f5d4e871a7467dec6c26d415152a Mon Sep 17 00:00:00 2001 From: Marius van Niekerk Date: Thu, 27 Feb 2025 15:00:15 -0500 Subject: [PATCH] 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 ++- .../flightsql/flightsql_adbc_server_test.go | 60 ++++++++++++++++++- .../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 9732bb0e19..1f178e9537 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 e3c1321838..c8a59d72ae 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 a25c839ec4..bbbcbbf061 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 bf4e0b4a80..9e517c7167 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) +}