feat(go/adbc/driver/flightsql): allow passing arbitrary grpc dial options in NewDatabase

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