refactor(go/adbc): refactor logging instrumentation into OTel tracing - part 1/3 (#4655)

This pull request introduces OpenTelemetry tracing to the FlightSQL
driver, enhancing observability for record reading and endpoint
streaming operations. The main changes involve adding tracing hooks,
attributes, and error recording to critical code paths, as well as
refactoring logging and context management to support tracing.
Additionally, new utility functions for collecting and attaching
response metadata to traces are implemented.

**Tracing and Observability Enhancements:**

* Added OpenTelemetry tracing support to `record_reader.go`, including
span creation, event recording, and error tracking in `newRecordReader`
and endpoint goroutines. This enables detailed tracing of FlightSQL
record reading operations.
* Introduced new tracing utility functions and types in
`flightsql_tracing.go` for collecting response metadata, building trace
attributes for endpoints, and summarizing stream progress as
OpenTelemetry attributes.

**Refactoring for Tracing Integration:**

* Replaced the logger-based endpoint and stream progress attribute
builders in `logging.go` with tracing attribute builders, and removed
now-redundant logging functions.
* Updated context and cancellation handling in `record_reader.go` to use
`context.CancelCauseFunc` for improved error propagation and tracing.

**Internal API and Dependency Updates:**

* Added OpenTelemetry and internal tracing imports to relevant files to
support the new tracing features.
* Extended the `recordReaderConfig` struct to include a tracing
configuration parameter, enabling tracing to be passed through to record
readers.
---
These changes collectively provide fine-grained tracing and error
visibility for FlightSQL operations, making it easier to monitor, debug,
and analyze the driver's behavior in production environments.
Refactors slog instrumentation into OTel tracing.
- update utilities to instrument duration for a span from a given start
time.
- update utilities to separate already recorded error to avoid duplicate
events in the span
- adds a `flightsql_tracing.go` to provide tracing wrappers originally
provided in `logging.go`
- some updates in `driverbase` to handle improved trace handling
- some initial instrumentation in flightsql connection and reader

Part 1 of a multi-part change to refactor logging instrumentation into
OTel tracing.

---------

Co-authored-by: Bruce Irschick (Bit Quill Technologies Inc) <v-birschick@microsoft.com>

Author: birschick-bq

go/adbc/driver/flightsql/flightsql_connection.go

@@ -39,6 +39,8 @@ import (
 	flightproto "github.com/apache/arrow-go/v18/arrow/flight/gen/flight"
 	"github.com/apache/arrow-go/v18/arrow/ipc"
 	"github.com/bluele/gcache"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/grpc"
 	grpccodes "google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -232,6 +234,153 @@ var adbcToFlightSQLInfo = map[adbc.InfoCode]flightsql.SqlInfo{
 	adbc.InfoVendorSubstraitMaxVersion: flightsql.SqlInfoFlightSqlServerSubstraitMaxVersion,
 }
 
+func doGetWithResponseMetadata(ctx context.Context, client *flightsql.Client, ticket *flight.Ticket, opts ...grpc.CallOption) (*flight.Reader, error) {
+	var header, trailer metadata.MD
+	callOpts := append(append([]grpc.CallOption{}, opts...), grpc.Header(&header), grpc.Trailer(&trailer))
+	reader, err := client.DoGet(ctx, ticket, callOpts...)
+	if err != nil {
+		captureResponseMetadata(ctx, metadata.Join(header, trailer))
+	}
+	return reader, err
+}
+
+func doGetWithTracer(ctx context.Context, cl *flightsql.Client, endpoint *flight.FlightEndpoint, clientCache gcache.Cache, tracing adbc.OTelTracing, opts ...grpc.CallOption) (rdr *flight.Reader, err error) {
+	const spanName = "FlightSQL.Connection.DoGet"
+	startTime := time.Now()
+	ctx, span := internal.StartSpan(ctx, spanName, tracing)
+	errorRecorded := false
+	defer func() {
+		internal.NewEndSpanHelper(span).
+			WithStartTime(startTime).
+			WithError(err).
+			WithRecordedError(errorRecorded).
+			EndSpan()
+	}()
+
+	streamOpts := make([]grpc.CallOption, 0, len(opts))
+	for _, opt := range opts {
+		switch opt.(type) {
+		case grpc.HeaderCallOption, *grpc.HeaderCallOption, grpc.TrailerCallOption, *grpc.TrailerCallOption:
+			continue
+		default:
+			streamOpts = append(streamOpts, opt)
+		}
+	}
+
+	if len(endpoint.Location) == 0 {
+		span.AddEvent("flight.location.attempt", trace.WithAttributes(
+			attribute.String("flight.location.source", "default_client"),
+		))
+		start := time.Now()
+		rdr, err = doGetWithResponseMetadata(ctx, cl, endpoint.Ticket, streamOpts...)
+		attrs := []attribute.KeyValue{
+			attribute.Float64("duration_s", time.Since(start).Seconds()),
+			attribute.String("flight.location.source", "default_client"),
+		}
+		if err != nil {
+			attrs = append(attrs, attribute.String("flight.stage", "do_get"))
+			span.RecordError(err, trace.WithAttributes(attrs...), trace.WithStackTrace(true))
+			errorRecorded = true
+		} else {
+			span.AddEvent("flight.location.selected", trace.WithAttributes(attrs...))
+		}
+		return rdr, err
+	}
+
+	var (
+		cc            interface{}
+		hasFallback   bool
+		attemptErrors []string
+	)
+
+	for _, loc := range endpoint.Location {
+		if loc.Uri == flight.LocationReuseConnection {
+			hasFallback = true
+			continue
+		}
+
+		start := time.Now()
+		span.AddEvent("flight.location.attempt", trace.WithAttributes(
+			attribute.String("flight.location", loc.Uri),
+			attribute.String("flight.location.source", "endpoint"),
+		))
+		cc, err = clientCache.Get(loc.Uri)
+		if err != nil {
+			attemptErrors = append(attemptErrors, fmt.Sprintf("clientCache.Get(%q): %s", loc.Uri, err.Error()))
+			span.AddEvent("flight.location.failed", trace.WithAttributes(
+				attribute.String("flight.stage", "client_cache_get"),
+				attribute.String("flight.location", loc.Uri),
+				attribute.Float64("duration_s", time.Since(start).Seconds()),
+				attribute.String("error.message", err.Error()),
+			))
+			continue
+		}
+
+		conn := cc.(*flightsql.Client)
+		rdr, err = doGetWithResponseMetadata(ctx, conn, endpoint.Ticket, streamOpts...)
+		if err != nil {
+			attemptErrors = append(attemptErrors, fmt.Sprintf("DoGet(%q): %s", loc.Uri, err.Error()))
+			span.AddEvent("flight.location.failed", trace.WithAttributes(
+				attribute.String("flight.stage", "do_get"),
+				attribute.String("flight.location", loc.Uri),
+				attribute.Float64("duration_s", time.Since(start).Seconds()),
+				attribute.String("error.message", err.Error()),
+			))
+			continue
+		}
+
+		span.AddEvent("flight.location.selected", trace.WithAttributes(
+			attribute.String("flight.location", loc.Uri),
+			attribute.String("flight.location.source", "endpoint"),
+			attribute.Float64("duration_s", time.Since(start).Seconds()),
+		))
+		return
+	}
+
+	if hasFallback {
+		start := time.Now()
+		span.AddEvent("flight.location.attempt", trace.WithAttributes(
+			attribute.String("flight.location.source", "fallback"),
+		))
+		rdr, err = doGetWithResponseMetadata(ctx, cl, endpoint.Ticket, streamOpts...)
+		if err != nil {
+			attemptErrors = append(attemptErrors, fmt.Sprintf("DoGet(fallback to default client): %s", err.Error()))
+			span.AddEvent("flight.location.failed", trace.WithAttributes(
+				attribute.String("flight.stage", "do_get"),
+				attribute.String("flight.location.source", "fallback"),
+				attribute.Float64("duration_s", time.Since(start).Seconds()),
+				attribute.String("error.message", err.Error()),
+			))
+			err = fmt.Errorf("all DoGet attempts failed: %s; final: %w", strings.Join(attemptErrors, "; "), err)
+			span.RecordError(err, trace.WithAttributes(
+				attribute.String("flight.stage", "all_locations_failed"),
+				attribute.Int("flight.location.attempt_count", len(attemptErrors)),
+			), trace.WithStackTrace(true))
+			errorRecorded = true
+			return nil, err
+		}
+		span.AddEvent("flight.location.selected", trace.WithAttributes(
+			attribute.String("flight.location.source", "fallback"),
+			attribute.Float64("duration_s", time.Since(start).Seconds()),
+		))
+		return rdr, nil
+	}
+
+	if err != nil && len(attemptErrors) > 1 {
+		err = fmt.Errorf("all %d DoGet location(s) failed: %s; final: %w",
+			len(attemptErrors), strings.Join(attemptErrors, "; "), err)
+	}
+	if err != nil {
+		span.RecordError(err, trace.WithAttributes(
+			attribute.String("flight.stage", "all_locations_failed"),
+			attribute.Int("flight.location.attempt_count", len(attemptErrors)),
+		), trace.WithStackTrace(true))
+		errorRecorded = true
+	}
+
+	return nil, err
+}
+
 // doGetWithLogger performs DoGet against an endpoint's locations, logging each
 // attempt and joining all per-location failures into the returned error so the
 // caller can see every location that was tried. logger may be nil.

go/adbc/driver/flightsql/flightsql_database.go

@@ -513,7 +513,11 @@ func (d *databaseImpl) Open(ctx context.Context) (_ adbc.Connection, err error)
 		d,
 		trace.WithAttributes(traceHeaderAttrsWithPrefix(d.hdrs, traceRequestMetadataPrefix)...),
 	)
-	defer internal.EndSpanWithError(span, &err)
+	defer func() {
+		internal.NewEndSpanHelper(span).
+			WithError(err).
+			EndSpan()
+	}()
 
 	authMiddle := &bearerAuthMiddleware{hdrs: d.hdrs.Copy(), logger: safeLogger(d.Logger)}
 	var cookies flight.CookieMiddleware

go/adbc/driver/flightsql/flightsql_statement.go

@@ -538,7 +538,11 @@ func (s *statement) ExecuteQuery(ctx context.Context) (rdr array.RecordReader, n
 		s.cnxn,
 		trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...),
 	)
-	defer internal.EndSpanWithError(span, &err)
+	defer func() {
+		internal.NewEndSpanHelper(span).
+			WithError(err).
+			EndSpan()
+	}()
 
 	// Handle bulk ingest
 	if s.targetTable != "" {
@@ -618,7 +622,11 @@ func (s *statement) ExecuteUpdate(ctx context.Context) (n int64, err error) {
 		s.cnxn,
 		trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...),
 	)
-	defer internal.EndSpanWithError(span, &err)
+	defer func() {
+		internal.NewEndSpanHelper(span).
+			WithError(err).
+			EndSpan()
+	}()
 
 	// Handle bulk ingest
 	if s.targetTable != "" {
@@ -672,7 +680,11 @@ func (s *statement) Prepare(ctx context.Context) (err error) {
 		s.cnxn,
 		trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...),
 	)
-	defer internal.EndSpanWithError(span, &err)
+	defer func() {
+		internal.NewEndSpanHelper(span).
+			WithError(err).
+			EndSpan()
+	}()
 
 	startTime := time.Now()
 	s.log.InfoContext(ctx, "FlightSQL Prepare start", s.queryAttrs()...)

go/adbc/driver/flightsql/flightsql_tracing.go

@@ -0,0 +1,179 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package flightsql
+
+import (
+	"context"
+	"encoding/hex"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/apache/arrow-go/v18/arrow/flight"
+	"go.opentelemetry.io/otel/attribute"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+type responseMetadataKey struct{}
+
+type responseMetadataCollector struct {
+	mutex sync.RWMutex
+	value metadata.MD
+}
+
+func withResponseMetadata(ctx context.Context) (context.Context, *responseMetadataCollector) {
+	collector := &responseMetadataCollector{}
+	return context.WithValue(ctx, responseMetadataKey{}, collector), collector
+}
+
+func captureResponseMetadata(ctx context.Context, value metadata.MD) {
+	collector, ok := responseMetadataFromContext(ctx)
+	if !ok {
+		return
+	}
+	collector.mutex.Lock()
+	collector.value = value.Copy()
+	defer collector.mutex.Unlock()
+}
+
+func responseMetadataFromContext(ctx context.Context) (*responseMetadataCollector, bool) {
+	collector, ok := ctx.Value(responseMetadataKey{}).(*responseMetadataCollector)
+	return collector, ok
+}
+
+func (c *responseMetadataCollector) snapshot() metadata.MD {
+	c.mutex.RLock()
+	defer c.mutex.RUnlock()
+	return c.value.Copy()
+}
+
+func responseMetadataStreamInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+	stream, err := streamer(ctx, desc, cc, method, opts...)
+	if err != nil {
+		return stream, err
+	}
+	if _, ok := responseMetadataFromContext(ctx); !ok {
+		return stream, nil
+	}
+	return &responseMetadataClientStream{ClientStream: stream, ctx: ctx}, nil
+}
+
+type responseMetadataClientStream struct {
+	grpc.ClientStream
+	ctx context.Context
+}
+
+func (s *responseMetadataClientStream) RecvMsg(message interface{}) error {
+	err := s.ClientStream.RecvMsg(message)
+	if err != nil {
+		header, _ := s.Header()
+		captureResponseMetadata(s.ctx, metadata.Join(header, s.Trailer()))
+	}
+	return err
+}
+
+// endpointTraceKeyValues builds OpenTelemetry attributes describing a Flight
+// endpoint. Ticket contents are intentionally never recorded.
+func endpointTraceKeyValues(endpointIndex, numEndpoints int, endpoint *flight.FlightEndpoint) []attribute.KeyValue {
+	attrs := []attribute.KeyValue{
+		attribute.Int("endpointIndex", endpointIndex),
+		attribute.Int("numEndpoints", numEndpoints),
+	}
+	if endpoint == nil {
+		return attrs
+	}
+	if endpoint.Ticket != nil {
+		attrs = append(attrs, attribute.Int("ticketBytes", len(endpoint.Ticket.Ticket)))
+	}
+	if len(endpoint.Location) == 0 {
+		attrs = append(attrs, attribute.String("locations", "<empty: using default client connection>"))
+	} else {
+		uris := make([]string, 0, len(endpoint.Location))
+		for _, loc := range endpoint.Location {
+			uris = append(uris, loc.Uri)
+		}
+		attrs = append(attrs, attribute.StringSlice("locations", uris))
+	}
+	if endpoint.ExpirationTime != nil {
+		attrs = append(attrs, attribute.String("expirationTime", endpoint.ExpirationTime.AsTime().String()))
+	}
+	return attrs
+}
+
+// logKeyValues returns OpenTelemetry attributes summarizing stream progress.
+func (p *streamProgress) logKeyValues() []attribute.KeyValue {
+	attrs := []attribute.KeyValue{
+		attribute.Int64("batchesRead", p.batchesRead),
+		attribute.Int64("recordsRead", p.recordsRead),
+		attribute.Int64("approxBytesRead", p.bytesEstimate),
+		attribute.String("elapsed", time.Since(p.start).String()),
+	}
+	if !p.firstBatchAt.IsZero() {
+		attrs = append(attrs, attribute.String("timeToFirstBatch", p.firstBatchAt.Sub(p.start).String()))
+	} else {
+		attrs = append(attrs, attribute.String("timeToFirstBatch", "never"))
+	}
+	if !p.lastBatchAt.IsZero() {
+		attrs = append(attrs, attribute.String("timeSinceLastBatch", time.Since(p.lastBatchAt).String()))
+	}
+	return attrs
+}
+
+// flightInfoTracingKeyValues returns OpenTelemetry attributes describing a FlightInfo:
+// descriptor type and command prefix, AppMetadata prefix (some backends
+// embed a server-side query handle there), and advisory record/byte
+// counts. Returns nil for a nil info.
+func flightInfoTracingKeyValues(info *flight.FlightInfo) []attribute.KeyValue {
+	if info == nil {
+		return nil
+	}
+	attrs := []attribute.KeyValue{
+		attribute.Int("numEndpoints", len(info.Endpoint)),
+		attribute.Int64("totalRecords", info.TotalRecords),
+		attribute.Int64("totalBytes", info.TotalBytes),
+		attribute.Bool("haveSchemaInFlightInfo", len(info.Schema) > 0),
+	}
+	if desc := info.FlightDescriptor; desc != nil {
+		attrs = append(attrs, attribute.String("descriptorType", desc.Type.String()))
+		if len(desc.Cmd) > 0 {
+			limit := len(desc.Cmd)
+			if limit > maxLoggedBlobBytes {
+				limit = maxLoggedBlobBytes
+			}
+			attrs = append(attrs,
+				attribute.Int("descriptorCmdBytes", len(desc.Cmd)),
+				attribute.String("descriptorCmdPrefixHex", hex.EncodeToString(desc.Cmd[:limit])),
+			)
+		}
+		if len(desc.Path) > 0 {
+			attrs = append(attrs, attribute.String("descriptorPath", fmt.Sprint(desc.Path)))
+		}
+	}
+	if len(info.AppMetadata) > 0 {
+		limit := len(info.AppMetadata)
+		if limit > maxLoggedBlobBytes {
+			limit = maxLoggedBlobBytes
+		}
+		attrs = append(attrs,
+			attribute.Int("appMetadataBytes", len(info.AppMetadata)),
+			attribute.String("appMetadataPrefixHex", hex.EncodeToString(info.AppMetadata[:limit])),
+		)
+	}
+	return attrs
+}