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

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() {
+		if errorRecorded {
+			internal.EndSpanWithStartTimeAndRecordedError(span, &err, &startTime)
+			return
+		}
+		internal.EndSpanWithStartTime(span, &err, &startTime)
+	}()
+
+	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_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()
+	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
+}