feat(go/adbc/driver/flightsql): allow passing arbitrary grpc dial opt…

…ions 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

@@ -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)

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 {

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...)

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)
+}