feat(go/adbc): prototype OpenTelemetry trace file exporter in go driver

go/adbc/adbc.go

@@ -250,6 +250,19 @@ const (
 	OptionKeyURI                      = "uri"
 	OptionKeyUsername                 = "username"
 	OptionKeyPassword                 = "password"
+	// EXPERIMENTAL. Sets/Gets the trace parent on OpenTelemetry traces
+	OptionKeyTelemetryTraceParent = "adbc.telemetry.trace_parent"
+)
+
+// EXPERIMENTAL. Traces Telemetry exporter option type
+type OptionTelemetryExporter string
+
+// EXPERIMENTAL. Traces Telemetry exporter options
+const (
+	TelemetryExporterNone     OptionTelemetryExporter = "none"
+	TelemetryExporterOtlp     OptionTelemetryExporter = "otlp"
+	TelemetryExporterConsole  OptionTelemetryExporter = "console"
+	TelemetryExporterAdbcFile OptionTelemetryExporter = "adbcfile"
 )
 
 type OptionIsolationLevel string

go/adbc/driver/bigquery/driver.go

@@ -18,6 +18,7 @@
 package bigquery
 
 import (
+	"context"
 	"fmt"
 	"runtime/debug"
 	"strings"
@@ -111,8 +112,16 @@ func NewDriver(alloc memory.Allocator) adbc.Driver {
 }
 
 func (d *driverImpl) NewDatabase(opts map[string]string) (adbc.Database, error) {
+	return d.NewDatabaseWithContext(context.Background(), opts)
+}
+
+func (d *driverImpl) NewDatabaseWithContext(ctx context.Context, opts map[string]string) (adbc.Database, error) {
+	dbBase, err := driverbase.NewDatabaseImplBase(ctx, &d.DriverImplBase)
+	if err != nil {
+		return nil, err
+	}
 	db := &databaseImpl{
-		DatabaseImplBase: driverbase.NewDatabaseImplBase(&d.DriverImplBase),
+		DatabaseImplBase: dbBase,
 		authType:         OptionValueAuthTypeDefault,
 	}
 	if err := db.SetOptions(opts); err != nil {

go/adbc/driver/flightsql/flightsql_driver.go

@@ -32,6 +32,7 @@
 package flightsql
 
 import (
+	"context"
 	"net/url"
 	"time"
 
@@ -101,6 +102,7 @@ type driverImpl struct {
 type Driver interface {
 	adbc.Driver
 	NewDatabaseWithOptions(map[string]string, ...grpc.DialOption) (adbc.Database, error)
+	NewDatabaseWithOptionsContext(context.Context, map[string]string, ...grpc.DialOption) (adbc.Database, error)
 }
 
 // NewDriver creates a new Flight SQL driver using the given Arrow allocator.
@@ -115,6 +117,10 @@ func NewDriver(alloc memory.Allocator) Driver {
 // 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) {
+	return d.NewDatabaseWithOptionsContext(context.Background(), opts, userDialOpts...)
+}
+
+func (d *driverImpl) NewDatabaseWithOptionsContext(ctx context.Context, opts map[string]string, userDialOpts ...grpc.DialOption) (adbc.Database, error) {
 	opts = maps.Clone(opts)
 	uri, ok := opts[adbc.OptionKeyURI]
 	if !ok {
@@ -125,8 +131,12 @@ func (d *driverImpl) NewDatabaseWithOptions(opts map[string]string, userDialOpts
 	}
 	delete(opts, adbc.OptionKeyURI)
 
+	dbBase, err := driverbase.NewDatabaseImplBase(ctx, &d.DriverImplBase)
+	if err != nil {
+		return nil, err
+	}
 	db := &databaseImpl{
-		DatabaseImplBase: driverbase.NewDatabaseImplBase(&d.DriverImplBase),
+		DatabaseImplBase: dbBase,
 		timeout: timeoutOption{
 			// Match gRPC default
 			connectTimeout: time.Second * 20,
@@ -135,7 +145,6 @@ func (d *driverImpl) NewDatabaseWithOptions(opts map[string]string, userDialOpts
 		userDialOpts: userDialOpts,
 	}
 
-	var err error
 	if db.uri, err = url.Parse(uri); err != nil {
 		return nil, adbc.Error{Msg: err.Error(), Code: adbc.StatusInvalidArgument}
 	}
@@ -154,5 +163,9 @@ func (d *driverImpl) NewDatabaseWithOptions(opts map[string]string, userDialOpts
 
 // 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)
+	return d.NewDatabaseWithContext(context.Background(), opts)
+}
+
+func (d *driverImpl) NewDatabaseWithContext(ctx context.Context, opts map[string]string) (adbc.Database, error) {
+	return d.NewDatabaseWithOptionsContext(ctx, opts)
 }

go/adbc/driver/internal/driverbase/connection.go

@@ -24,26 +24,31 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"strings"
 
 	"github.com/apache/arrow-adbc/go/adbc"
 	"github.com/apache/arrow-go/v18/arrow"
 	"github.com/apache/arrow-go/v18/arrow/array"
 	"github.com/apache/arrow-go/v18/arrow/memory"
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/sync/errgroup"
 )
 
 const (
-	ConnectionMessageOptionUnknown     = "Unknown connection option"
-	ConnectionMessageOptionUnsupported = "Unsupported connection option"
-	ConnectionMessageCannotCommit      = "Cannot commit when autocommit is enabled"
-	ConnectionMessageCannotRollback    = "Cannot rollback when autocommit is enabled"
+	ConnectionMessageOptionUnknown              = "Unknown connection option"
+	ConnectionMessageOptionUnsupported          = "Unsupported connection option"
+	ConnectionMessageCannotCommit               = "Cannot commit when autocommit is enabled"
+	ConnectionMessageCannotRollback             = "Cannot rollback when autocommit is enabled"
+	ConnectionMessageTraceParentIncorrectFormat = "Incorrect or unsupported trace parent format"
 )
 
 // ConnectionImpl is an interface that drivers implement to provide
 // vendor-specific functionality.
 type ConnectionImpl interface {
 	adbc.Connection
 	adbc.GetSetOptions
+	adbc.OTelTracing
 	Base() *ConnectionImplBase
 }
 
@@ -107,9 +112,12 @@ type ConnectionImplBase struct {
 	ErrorHelper ErrorHelper
 	DriverInfo  *DriverInfo
 	Logger      *slog.Logger
+	Tracer      trace.Tracer
 
 	Autocommit bool
 	Closed     bool
+
+	traceParent string
 }
 
 // NewConnectionImplBase instantiates ConnectionImplBase.
@@ -122,6 +130,7 @@ func NewConnectionImplBase(database *DatabaseImplBase) ConnectionImplBase {
 		ErrorHelper: database.ErrorHelper,
 		DriverInfo:  database.DriverInfo,
 		Logger:      database.Logger,
+		Tracer:      database.Tracer,
 		Autocommit:  true,
 		Closed:      false,
 	}
@@ -223,6 +232,10 @@ func (base *ConnectionImplBase) ReadPartition(ctx context.Context, serializedPar
 }
 
 func (base *ConnectionImplBase) GetOption(key string) (string, error) {
+	switch strings.ToLower(key) {
+	case adbc.OptionKeyTelemetryTraceParent:
+		return base.GetTraceParent(), nil
+	}
 	return "", base.ErrorHelper.Errorf(adbc.StatusNotFound, "%s '%s'", ConnectionMessageOptionUnknown, key)
 }
 
@@ -239,9 +252,12 @@ func (base *ConnectionImplBase) GetOptionInt(key string) (int64, error) {
 }
 
 func (base *ConnectionImplBase) SetOption(key string, val string) error {
-	switch key {
+	switch strings.ToLower(key) {
 	case adbc.OptionKeyAutoCommit:
 		return base.ErrorHelper.Errorf(adbc.StatusNotImplemented, "%s '%s'", ConnectionMessageOptionUnsupported, key)
+	case adbc.OptionKeyTelemetryTraceParent:
+		base.SetTraceParent(strings.TrimSpace(val))
+		return nil
 	}
 	return base.ErrorHelper.Errorf(adbc.StatusNotImplemented, "%s '%s'", ConnectionMessageOptionUnknown, key)
 }
@@ -332,6 +348,23 @@ func (b *ConnectionBuilder) Connection() Connection {
 	return conn
 }
 
+func (cnxn *ConnectionImplBase) GetTraceParent() string {
+	return cnxn.traceParent
+}
+
+func (cnxn *ConnectionImplBase) SetTraceParent(traceParent string) {
+	cnxn.traceParent = traceParent
+}
+
+func (cnxn *ConnectionImplBase) StartSpan(
+	ctx context.Context,
+	spanName string,
+	opts ...trace.SpanStartOption,
+) (context.Context, trace.Span) {
+	ctx, _ = maybeAddTraceParent(ctx, cnxn, nil)
+	return cnxn.Tracer.Start(ctx, spanName, opts...)
+}
+
 // GetObjects implements Connection.
 func (cnxn *connection) GetObjects(ctx context.Context, depth adbc.ObjectDepth, catalog *string, dbSchema *string, tableName *string, columnName *string, tableType []string) (array.RecordReader, error) {
 	helper := cnxn.dbObjectsEnumerator
@@ -714,4 +747,37 @@ func ValueOrZero[T any](val *T) T {
 	return *val
 }
 
+func maybeAddTraceParent(ctx context.Context, cnxn adbc.OTelTracing, st adbc.OTelTracing) (context.Context, error) {
+	var traceParentStr string
+	if st != nil && st.GetTraceParent() != "" {
+		traceParentStr = st.GetTraceParent()
+	} else if cnxn != nil && cnxn.GetTraceParent() != "" {
+		traceParentStr = cnxn.GetTraceParent()
+	}
+	if traceParentStr != "" {
+		spanContext, err := propagateTraceParent(ctx, traceParentStr)
+		if err != nil {
+			return ctx, err
+		}
+		ctx = trace.ContextWithRemoteSpanContext(ctx, spanContext)
+	}
+	return ctx, nil
+}
+
+func propagateTraceParent(ctx context.Context, traceParentStr string) (trace.SpanContext, error) {
+	if strings.TrimSpace(traceParentStr) == "" {
+		return trace.SpanContext{}, errors.New("traceparent string is empty")
+	}
+
+	propagator := propagation.TraceContext{}
+	carrier := propagation.MapCarrier{"traceparent": traceParentStr}
+	extractedContext := propagator.Extract(ctx, carrier)
+
+	spanContext := trace.SpanContextFromContext(extractedContext)
+	if !spanContext.IsValid() {
+		return trace.SpanContext{}, errors.New("invalid traceparent string")
+	}
+	return spanContext, nil
+}
+
 var _ ConnectionImpl = (*ConnectionImplBase)(nil)