This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new f3dbc6b31 refactor(go/adbc): refactor logging instrumentation into
OTel tracing - part 1/3 (#4655)
f3dbc6b31 is described below
commit f3dbc6b314f93ad06a0039ed2c1dc4a730ebcc2b
Author: Bruce Irschick <[email protected]>
AuthorDate: Mon Aug 17 17:51:56 2026 -0700
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)
<[email protected]>
---
go/adbc/driver/flightsql/flightsql_connection.go | 149 +++++++++++++++++
go/adbc/driver/flightsql/flightsql_database.go | 6 +-
go/adbc/driver/flightsql/flightsql_statement.go | 18 +-
go/adbc/driver/flightsql/flightsql_tracing.go | 179 ++++++++++++++++++++
go/adbc/driver/flightsql/logging.go | 48 ------
go/adbc/driver/flightsql/record_reader.go | 199 +++++++++++++++--------
go/adbc/driver/flightsql/record_reader_test.go | 159 +++++++++++++++++-
go/adbc/driver/internal/driverbase/connection.go | 9 +-
go/adbc/driver/internal/driverbase/database.go | 34 +++-
go/adbc/driver/internal/shared_utils.go | 84 +++++++++-
10 files changed, 750 insertions(+), 135 deletions(-)
diff --git a/go/adbc/driver/flightsql/flightsql_connection.go
b/go/adbc/driver/flightsql/flightsql_connection.go
index 11d9a6473..858b9f40d 100644
--- a/go/adbc/driver/flightsql/flightsql_connection.go
+++ b/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.
diff --git a/go/adbc/driver/flightsql/flightsql_database.go
b/go/adbc/driver/flightsql/flightsql_database.go
index 28e73399d..0ead6422c 100644
--- a/go/adbc/driver/flightsql/flightsql_database.go
+++ b/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
diff --git a/go/adbc/driver/flightsql/flightsql_statement.go
b/go/adbc/driver/flightsql/flightsql_statement.go
index 5a5b175aa..61911e9e3 100644
--- a/go/adbc/driver/flightsql/flightsql_statement.go
+++ b/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()...)
diff --git a/go/adbc/driver/flightsql/flightsql_tracing.go
b/go/adbc/driver/flightsql/flightsql_tracing.go
new file mode 100644
index 000000000..168e993de
--- /dev/null
+++ b/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
+}
diff --git a/go/adbc/driver/flightsql/logging.go
b/go/adbc/driver/flightsql/logging.go
index 9dae400b8..48a342728 100644
--- a/go/adbc/driver/flightsql/logging.go
+++ b/go/adbc/driver/flightsql/logging.go
@@ -51,35 +51,6 @@ func safeLogger(logger *slog.Logger) *slog.Logger {
// tickets are not logged at all because they may carry sensitive data.
const maxLoggedBlobBytes = 32
-// endpointLogAttrs builds slog attributes describing a Flight endpoint
-// (index, ticket length, locations) for per-endpoint log records. Ticket
-// contents are intentionally never logged.
-func endpointLogAttrs(endpointIndex, numEndpoints int, endpoint
*flight.FlightEndpoint) []any {
- attrs := []any{
- slog.Int("endpointIndex", endpointIndex),
- slog.Int("numEndpoints", numEndpoints),
- }
- if endpoint == nil {
- return attrs
- }
- if endpoint.Ticket != nil {
- attrs = append(attrs, slog.Int("ticketBytes",
len(endpoint.Ticket.Ticket)))
- }
- if len(endpoint.Location) == 0 {
- attrs = append(attrs, slog.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, slog.Any("locations", uris))
- }
- if endpoint.ExpirationTime != nil {
- attrs = append(attrs, slog.Time("expirationTime",
endpoint.ExpirationTime.AsTime()))
- }
- return attrs
-}
-
// streamProgress tracks per-endpoint streaming statistics for log records
// and error messages emitted when a stream ends. Not safe for concurrent
// use; intended to be owned by the goroutine driving one endpoint.
@@ -108,25 +79,6 @@ func (p *streamProgress) recordBatch(rows int64, bytes
int64) {
p.bytesEstimate += bytes
}
-// logAttrs returns slog attributes summarizing this stream's progress.
-func (p *streamProgress) logAttrs() []any {
- attrs := []any{
- slog.Int64("batchesRead", p.batchesRead),
- slog.Int64("recordsRead", p.recordsRead),
- slog.Int64("approxBytesRead", p.bytesEstimate),
- slog.Duration("elapsed", time.Since(p.start)),
- }
- if !p.firstBatchAt.IsZero() {
- attrs = append(attrs, slog.Duration("timeToFirstBatch",
p.firstBatchAt.Sub(p.start)))
- } else {
- attrs = append(attrs, slog.String("timeToFirstBatch", "never"))
- }
- if !p.lastBatchAt.IsZero() {
- attrs = append(attrs, slog.Duration("timeSinceLastBatch",
time.Since(p.lastBatchAt)))
- }
- return attrs
-}
-
// summary returns a compact human-readable summary of the stream's progress
// suitable for embedding into wrapped error messages.
func (p *streamProgress) summary() string {
diff --git a/go/adbc/driver/flightsql/record_reader.go
b/go/adbc/driver/flightsql/record_reader.go
index 071cd1880..2cfbeb9d9 100644
--- a/go/adbc/driver/flightsql/record_reader.go
+++ b/go/adbc/driver/flightsql/record_reader.go
@@ -19,11 +19,14 @@ package flightsql
import (
"context"
+ "errors"
"fmt"
"log/slog"
"sync/atomic"
+ "time"
"github.com/apache/arrow-adbc/go/adbc"
+ "github.com/apache/arrow-adbc/go/adbc/driver/internal"
"github.com/apache/arrow-adbc/go/adbc/utils"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
@@ -32,9 +35,10 @@ import (
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/apache/arrow-go/v18/arrow/util"
"github.com/bluele/gcache"
+ "go.opentelemetry.io/otel/attribute"
+ "go.opentelemetry.io/otel/trace"
"golang.org/x/sync/errgroup"
"google.golang.org/grpc"
- "google.golang.org/grpc/metadata"
)
type reader struct {
@@ -45,9 +49,11 @@ type reader struct {
rec arrow.RecordBatch
err error
- cancelFn context.CancelFunc
+ cancelFn context.CancelCauseFunc
}
+var errReaderReleased = errors.New("record reader released")
+
// recordReaderConfig bundles the dependencies that newRecordReader
// needs to spin up its per-endpoint goroutines.
type recordReaderConfig struct {
@@ -56,18 +62,30 @@ type recordReaderConfig struct {
info *flight.FlightInfo
clientCache gcache.Cache
bufferSize int
+ tracing adbc.OTelTracing
logger *slog.Logger
}
// newRecordReader kicks off a goroutine for each endpoint and returns a
-// reader which gathers all of the records as they come in. cfg.logger
-// may be nil.
+// reader which gathers all of the records as they come in.
func newRecordReader(ctx context.Context, cfg recordReaderConfig, opts
...grpc.CallOption) (rdr array.RecordReader, err error) {
- log := safeLogger(cfg.logger)
+ const spanName = "FlightSQL.RecordReader.newRecordReader"
+ startTime := time.Now()
+ ctx, span := internal.StartSpan(ctx, spanName, cfg.tracing)
+ spanOwnedByReader := false
+ errorRecorded := false
+ defer func() {
+ if !spanOwnedByReader {
+ internal.NewEndSpanHelper(span).
+ WithStartTime(startTime).
+ WithError(err).
+ WithRecordedError(errorRecorded).
+ EndSpan()
+ }
+ }()
+
info := cfg.info
endpoints := info.Endpoint
- var header, trailer metadata.MD
- opts = append(append([]grpc.CallOption{}, opts...),
grpc.Header(&header), grpc.Trailer(&trailer))
var schema *arrow.Schema
if len(endpoints) == 0 {
if info.Schema == nil {
@@ -87,21 +105,31 @@ func newRecordReader(ctx context.Context, cfg
recordReaderConfig, opts ...grpc.C
}
ch := make(chan arrow.RecordBatch, cfg.bufferSize)
+ callerCtx := ctx
group, ctx := errgroup.WithContext(ctx)
- ctx, cancelFn := context.WithCancel(ctx)
+ ctx, cancelFn := context.WithCancelCause(ctx)
+ goEndpoint := func(endpointFn func() error) {
+ group.Go(func() error {
+ err := endpointFn()
+ if err != nil {
+ cancelFn(err)
+ }
+ return err
+ })
+ }
// We may mutate endpoints below
numEndpoints := len(endpoints)
- log.DebugContext(ctx, "FlightSQL newRecordReader start",
- append([]any{
- slog.Int("bufferSize", cfg.bufferSize),
- }, flightInfoLogAttrs(info)...)...,
- )
+ span.AddEvent("endpoint_stream.starting", trace.WithAttributes(
+ append([]attribute.KeyValue{
+ attribute.Int("bufferSize", cfg.bufferSize),
+ }, flightInfoTracingKeyValues(info)...)...,
+ ))
defer func() {
if err != nil {
close(ch)
- cancelFn()
+ cancelFn(err)
}
}()
@@ -114,22 +142,26 @@ func newRecordReader(ctx context.Context, cfg
recordReaderConfig, opts ...grpc.C
}
} else {
firstEndpoint := endpoints[0]
- epAttrs := endpointLogAttrs(0, numEndpoints, firstEndpoint)
- log.DebugContext(ctx, "FlightSQL endpoint stream opening
(schema discovery)", epAttrs...)
+ epAttrs := endpointTraceKeyValues(0, numEndpoints,
firstEndpoint)
+ span.AddEvent("endpoint_stream.opening_schema_discovery",
trace.WithAttributes(epAttrs...))
startSchemaFetch := newStreamProgress()
- rdr, err := doGetWithLogger(ctx, cfg.cl, firstEndpoint,
cfg.clientCache, log, opts...)
+ endpointCtx, responseMetadata := withResponseMetadata(ctx)
+ var rdr array.RecordReader
+ rdr, err = doGetWithTracer(endpointCtx, cfg.cl, firstEndpoint,
cfg.clientCache, cfg.tracing, opts...)
if err != nil {
- log.ErrorContext(ctx, "FlightSQL endpoint DoGet failed
(schema discovery)",
- append(append([]any{}, epAttrs...),
- "err", err,
- "elapsed", startSchemaFetch.summary(),
+ span.RecordError(err, trace.WithAttributes(
+ append(append([]attribute.KeyValue{},
epAttrs...),
+ attribute.String("elapsed",
startSchemaFetch.summary()),
+ attribute.String("flight.stage",
"schema_discovery"),
)...,
- )
- return nil, adbcFromFlightStatusWithDetails(err,
header, trailer,
+ ))
+ errorRecorded = true
+ return nil, adbcFromFlightStatusWithDetails(err,
responseMetadata.snapshot(), nil,
"DoGet: endpoint 0: remote: %s",
firstEndpoint.Location)
}
schema = rdr.Schema()
- group.Go(func() error {
+ goEndpoint(func() error {
+ span := trace.SpanFromContext(ctx)
defer rdr.Release()
if numEndpoints > 1 {
defer close(ch)
@@ -142,20 +174,24 @@ func newRecordReader(ctx context.Context, cfg
recordReaderConfig, opts ...grpc.C
rec.Retain()
ch <- rec
}
- if err := checkContext(rdr.Err(), ctx); err != nil {
- log.ErrorContext(ctx, "FlightSQL endpoint
stream ended with error",
- append(append([]any{},
endpointLogAttrs(0, numEndpoints, firstEndpoint)...),
- append([]any{"err", err},
progress.logAttrs()...)...,
- )...,
+ if err := checkRecordReaderContext(rdr.Err(), ctx,
callerCtx); err != nil {
+ attrs := endpointTraceKeyValues(0,
numEndpoints, firstEndpoint)
+ attrs = append(attrs,
progress.logKeyValues()...)
+ span.RecordError(err,
+ /*"FlightSQL endpoint stream ended with
error",*/
+ trace.WithAttributes(attrs...),
)
- return adbcFromFlightStatusWithDetails(err,
header, trailer,
+ return adbcFromFlightStatusWithDetails(err,
responseMetadata.snapshot(), nil,
"DoGet: endpoint 0: remote: %s",
firstEndpoint.Location)
}
- log.DebugContext(ctx, "FlightSQL endpoint stream
completed",
- append(append([]any{}, endpointLogAttrs(0,
numEndpoints, firstEndpoint)...),
- progress.logAttrs()...,
+ span.AddEvent("endpoint_stream.completed",
trace.WithAttributes(
+ append(
+ append(
+ []attribute.KeyValue{},
+ endpointTraceKeyValues(0,
numEndpoints, firstEndpoint)...),
+ progress.logKeyValues()...,
)...,
- )
+ ))
return nil
})
@@ -185,37 +221,43 @@ func newRecordReader(ctx context.Context, cfg
recordReaderConfig, opts ...grpc.C
logEndpointIndex = endpointIndex + 1
}
chs[endpointIndex] = make(chan arrow.RecordBatch,
cfg.bufferSize)
- group.Go(func() error {
+ goEndpoint(func() error {
// Close channels (except the last) so that Next can
move on to the next channel properly
if endpointIndex != lastChannelIndex {
defer close(chs[endpointIndex])
}
- epAttrs := endpointLogAttrs(logEndpointIndex,
numEndpoints, endpoint)
- log.DebugContext(ctx, "FlightSQL endpoint stream
opening", epAttrs...)
+ epAttrs := endpointTraceKeyValues(logEndpointIndex,
numEndpoints, endpoint)
+ span.AddEvent("endpoint_stream.opening",
trace.WithAttributes(epAttrs...))
doGetStart := newStreamProgress()
- rdr, err := doGetWithLogger(ctx, cfg.cl, endpoint,
cfg.clientCache, log, opts...)
+ endpointCtx, responseMetadata :=
withResponseMetadata(ctx)
+ rdr, err := doGetWithTracer(endpointCtx, cfg.cl,
endpoint, cfg.clientCache, cfg.tracing, opts...)
if err != nil {
- log.ErrorContext(ctx, "FlightSQL endpoint DoGet
failed",
- append(append([]any{}, epAttrs...),
- "err", err,
- "elapsed", doGetStart.summary(),
+ span.RecordError(err, trace.WithAttributes(
+ append(
+ append([]attribute.KeyValue{},
epAttrs...),
+ attribute.String("err",
err.Error()),
+ attribute.String("elapsed",
doGetStart.summary()),
+
attribute.String("flight.stage", "do_get"),
)...,
- )
- return adbcFromFlightStatusWithDetails(err,
header, trailer,
+ ))
+ return adbcFromFlightStatusWithDetails(err,
responseMetadata.snapshot(), nil,
"DoGet: endpoint %d: %s",
logEndpointIndex, endpoint.Location)
}
defer rdr.Release()
streamSchema := utils.RemoveSchemaMetadata(rdr.Schema())
if !streamSchema.Equal(referenceSchema) {
- log.ErrorContext(ctx, "FlightSQL endpoint
returned inconsistent schema",
- append(append([]any{}, epAttrs...),
- "expectedSchema",
referenceSchema.String(),
- "actualSchema",
streamSchema.String(),
+ err = fmt.Errorf("endpoint %d returned
inconsistent schema: expected %s but got %s", logEndpointIndex,
referenceSchema.String(), streamSchema.String())
+ span.RecordError(err, trace.WithAttributes(
+ append(
+ append([]attribute.KeyValue{},
epAttrs...),
+
attribute.String("expectedSchema", referenceSchema.String()),
+
attribute.String("actualSchema", streamSchema.String()),
+ attribute.String("stage",
"FlightSQL endpoint returned inconsistent schema"),
)...,
- )
- return fmt.Errorf("endpoint %d returned
inconsistent schema: expected %s but got %s", logEndpointIndex,
referenceSchema.String(), streamSchema.String())
+ ))
+ return err
}
progress := newStreamProgress()
@@ -226,45 +268,64 @@ func newRecordReader(ctx context.Context, cfg
recordReaderConfig, opts ...grpc.C
chs[endpointIndex] <- rec
}
- if err := checkContext(rdr.Err(), ctx); err != nil {
- log.ErrorContext(ctx, "FlightSQL endpoint
stream ended with error",
- append(append([]any{}, epAttrs...),
- append([]any{"err", err},
progress.logAttrs()...)...,
+ if err := checkRecordReaderContext(rdr.Err(), ctx,
callerCtx); err != nil {
+ span.RecordError(err, trace.WithAttributes(
+ append(append([]attribute.KeyValue{},
epAttrs...),
+ append([]attribute.KeyValue{
+ attribute.String("err",
err.Error()),
+
attribute.String("stage", "FlightSQL endpoint stream ended with error"),
+ },
progress.logKeyValues()...)...,
)...,
- )
- return adbcFromFlightStatusWithDetails(err,
header, trailer,
+ ))
+ return adbcFromFlightStatusWithDetails(err,
responseMetadata.snapshot(), nil,
"DoGet: endpoint %d: %s",
logEndpointIndex, endpoint.Location)
}
- log.DebugContext(ctx, "FlightSQL endpoint stream
completed",
- append(append([]any{}, epAttrs...),
- progress.logAttrs()...,
+ span.AddEvent("endpoint_stream.completed",
trace.WithAttributes(
+ append(append([]attribute.KeyValue{},
epAttrs...),
+ progress.logKeyValues()...,
)...,
- )
+ ))
return nil
})
}
+ spanOwnedByReader = true
go func() {
err := group.Wait()
reader.err = err
if reader.err != nil {
- log.WarnContext(ctx, "FlightSQL record reader finished
with error",
- "err", reader.err,
- "numEndpoints", numEndpoints,
- )
+ span.AddEvent("record_reader.failed",
trace.WithAttributes(
+ attribute.Int("numEndpoints", numEndpoints),
+ ))
} else {
- log.DebugContext(ctx, "FlightSQL record reader finished
successfully",
- "numEndpoints", numEndpoints,
- )
+ span.AddEvent("record_reader.completed",
trace.WithAttributes(
+ attribute.Int("numEndpoints", numEndpoints),
+ ))
}
+ errorRecorded := reader.err != nil
+ internal.NewEndSpanHelper(span).
+ WithStartTime(startTime).
+ WithError(reader.err).
+ WithRecordedError(errorRecorded).
+ EndSpan()
// Don't close the last channel until after the group is
finished, so that
- // Next() can only return after reader.err may have been set
+ // Next() can only return after reader.err and tracing have
been finalized.
close(chs[lastChannelIndex])
}()
return reader, nil
}
+func checkRecordReaderContext(maybeErr error, ctx, callerCtx context.Context)
error {
+ if errors.Is(context.Cause(ctx), errReaderReleased) {
+ return nil
+ }
+ if ctx.Err() == context.Canceled && callerCtx.Err() == nil {
+ return nil
+ }
+ return checkContext(maybeErr, ctx)
+}
+
func (r *reader) Retain() {
atomic.AddInt64(&r.refCount, 1)
}
@@ -274,7 +335,7 @@ func (r *reader) Release() {
if r.rec != nil {
r.rec.Release()
}
- r.cancelFn()
+ r.cancelFn(errReaderReleased)
for _, ch := range r.chs {
for rec := range ch {
rec.Release()
diff --git a/go/adbc/driver/flightsql/record_reader_test.go
b/go/adbc/driver/flightsql/record_reader_test.go
index ab7b5f179..987f533a7 100644
--- a/go/adbc/driver/flightsql/record_reader_test.go
+++ b/go/adbc/driver/flightsql/record_reader_test.go
@@ -33,8 +33,13 @@ import (
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/bluele/gcache"
"github.com/stretchr/testify/suite"
+ "go.opentelemetry.io/otel/attribute"
+ sdktrace "go.opentelemetry.io/otel/sdk/trace"
+ "go.opentelemetry.io/otel/sdk/trace/tracetest"
+ "go.opentelemetry.io/otel/trace"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
+ "google.golang.org/grpc/metadata"
)
func orderingSchema() *arrow.Schema {
@@ -50,6 +55,20 @@ type testFlightService struct {
failureCount int
}
+type recorderTracing struct {
+ tracer trace.Tracer
+}
+
+func (*recorderTracing) SetTraceParent(string) {}
+
+func (*recorderTracing) GetTraceParent() string { return "" }
+
+func (t *recorderTracing) StartSpan(ctx context.Context, name string, opts
...trace.SpanStartOption) (context.Context, trace.Span) {
+ return t.tracer.Start(ctx, name, opts...)
+}
+
+func (*recorderTracing) GetInitialSpanAttributes() []attribute.KeyValue {
return nil }
+
func (f *testFlightService) DoGet(request *flight.Ticket, stream
flight.FlightService_DoGetServer) (err error) {
// Crude way to make requests fail until retried enough times
if f.failureCount > 0 {
@@ -78,12 +97,20 @@ func (f *testFlightService) DoGet(request *flight.Ticket,
stream flight.FlightSe
if err := wr.Write(rec); err != nil {
return err
}
+ if request.Ticket[0] == 126 {
+ <-stream.Context().Done()
+ return stream.Context().Err()
+ }
+ if request.Ticket[0] == 127 {
+ stream.SetTrailer(metadata.Pairs("x-request-id",
"late-stream-error"))
+ return fmt.Errorf("late stream failure")
+ }
}
return nil
}
-func getFlightClientTest(ctx context.Context, loc string) (*flightsql.Client,
error) {
+func getFlightClientTest(_ context.Context, loc string) (*flightsql.Client,
error) {
uri, err := url.Parse(loc)
if err != nil {
return nil, err
@@ -179,6 +206,131 @@ func (suite *RecordReaderTests)
TestFallbackFailedConnection() {
suite.NoError(reader.Err())
}
+func (suite *RecordReaderTests) TestFallbackTracing() {
+ goodLocation := "grpc://" + suite.server.Addr().String()
+ badLocation := "grpc://127.0.0.2:1234"
+ recorder := tracetest.NewSpanRecorder()
+ provider :=
sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
+ defer func() {
+ suite.NoError(provider.Shutdown(context.Background()))
+ }()
+ tracing := &recorderTracing{tracer: provider.Tracer("test")}
+
+ endpoint := &flight.FlightEndpoint{
+ Ticket: &flight.Ticket{Ticket: []byte{0}},
+ Location: []*flight.Location{{Uri: badLocation}, {Uri:
goodLocation}},
+ }
+ reader, err := doGetWithTracer(context.Background(), suite.cl,
endpoint, suite.clCache, tracing)
+ suite.NoError(err)
+ reader.Release()
+ suite.Equal(1, countSpanEvents(recorder.Ended(),
"flight.location.failed"))
+ suite.Zero(countSpanEvents(recorder.Ended(), "exception"))
+
+ recorder.Reset()
+ endpoint.Location = []*flight.Location{{Uri: badLocation}, {Uri:
badLocation}}
+ reader, err = doGetWithTracer(context.Background(), suite.cl, endpoint,
suite.clCache, tracing)
+ suite.Nil(reader)
+ suite.Error(err)
+ suite.Equal(2, countSpanEvents(recorder.Ended(),
"flight.location.failed"))
+ suite.Equal(1, countSpanEvents(recorder.Ended(), "exception"))
+}
+
+func (suite *RecordReaderTests) TestLateStreamErrorMetadata() {
+ middleware := []flight.ClientMiddleware{
+ flight.CreateClientMiddleware(&bearerAuthMiddleware{hdrs:
make(metadata.MD)}),
+ {Stream: responseMetadataStreamInterceptor},
+ }
+ client, err := flightsql.NewClient(suite.server.Addr().String(), nil,
middleware, grpc.WithTransportCredentials(insecure.NewCredentials()))
+ suite.Require().NoError(err)
+ defer func() {
+ suite.NoError(client.Close())
+ }()
+
+ ctx, responseMetadata := withResponseMetadata(context.Background())
+ reader, err := doGetWithTracer(ctx, client, &flight.FlightEndpoint{
+ Ticket: &flight.Ticket{Ticket: []byte{127}},
+ }, suite.clCache, nil)
+ suite.Require().NoError(err)
+ defer reader.Release()
+
+ for reader.Next() {
+ }
+ suite.Error(reader.Err())
+ suite.Equal([]string{"late-stream-error"},
responseMetadata.snapshot().Get("x-request-id"))
+}
+
+func (suite *RecordReaderTests) TestEarlyReleaseTracing() {
+ recorder := tracetest.NewSpanRecorder()
+ provider :=
sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
+ defer func() {
+ suite.NoError(provider.Shutdown(context.Background()))
+ }()
+
+ reader, err := newRecordReader(context.Background(), recordReaderConfig{
+ alloc: suite.alloc,
+ cl: suite.cl,
+ info: &flight.FlightInfo{
+ Schema: flight.SerializeSchema(orderingSchema(),
suite.alloc),
+ Endpoint: []*flight.FlightEndpoint{{
+ Ticket: &flight.Ticket{Ticket: []byte{126}},
+ }},
+ },
+ clientCache: suite.clCache,
+ bufferSize: 1,
+ tracing: &recorderTracing{tracer: provider.Tracer("test")},
+ })
+ suite.Require().NoError(err)
+ suite.True(reader.Next())
+ reader.Release()
+
+ suite.Zero(countSpanEvents(recorder.Ended(), "exception"))
+ suite.Zero(countSpanEvents(recorder.Ended(), "record_reader.failed"))
+ suite.Equal(1, countSpanEvents(recorder.Ended(),
"record_reader.completed"))
+}
+
+func (suite *RecordReaderTests) TestSiblingCancellationRecordsOneException() {
+ recorder := tracetest.NewSpanRecorder()
+ provider :=
sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
+ defer func() {
+ suite.NoError(provider.Shutdown(context.Background()))
+ }()
+
+ reader, err := newRecordReader(context.Background(), recordReaderConfig{
+ alloc: suite.alloc,
+ cl: suite.cl,
+ info: &flight.FlightInfo{
+ Schema: flight.SerializeSchema(orderingSchema(),
suite.alloc),
+ Endpoint: []*flight.FlightEndpoint{
+ {Ticket: &flight.Ticket{Ticket: []byte{127}}},
+ {Ticket: &flight.Ticket{Ticket: []byte{126}}},
+ },
+ },
+ clientCache: suite.clCache,
+ bufferSize: 1,
+ tracing: &recorderTracing{tracer: provider.Tracer("test")},
+ })
+ suite.Require().NoError(err)
+ defer reader.Release()
+
+ for reader.Next() {
+ }
+ suite.Error(reader.Err())
+ suite.Equal(1, countSpanEvents(recorder.Ended(), "exception"))
+ suite.Equal(1, countSpanEvents(recorder.Ended(),
"record_reader.failed"))
+}
+
+func countSpanEvents(spans []sdktrace.ReadOnlySpan, name string) int {
+ count := 0
+ for _, span := range spans {
+ for _, event := range span.Events() {
+ if event.Name == name {
+ count++
+ }
+ }
+ }
+ return count
+}
+
func (suite *RecordReaderTests) TestFallbackFailedDoGet() {
defer func() {
suite.service.failureCount = 0
@@ -393,13 +545,14 @@ func (suite *RecordReaderTests) TestOrdering() {
},
}
+ var header, trailer metadata.MD
reader, err := newRecordReader(context.Background(), recordReaderConfig{
alloc: suite.alloc,
cl: suite.cl,
info: &info,
clientCache: suite.clCache,
bufferSize: 3,
- })
+ }, grpc.Header(&header), grpc.Trailer(&trailer))
suite.NoError(err)
defer reader.Release()
@@ -423,6 +576,8 @@ func (suite *RecordReaderTests) TestOrdering() {
}
suite.False(reader.Next())
suite.NoError(reader.Err())
+ suite.Nil(header)
+ suite.Nil(trailer)
}
func TestRecordReader(t *testing.T) {
diff --git a/go/adbc/driver/internal/driverbase/connection.go
b/go/adbc/driver/internal/driverbase/connection.go
index fbd88bc61..5db09c835 100644
--- a/go/adbc/driver/internal/driverbase/connection.go
+++ b/go/adbc/driver/internal/driverbase/connection.go
@@ -25,6 +25,7 @@ import (
"fmt"
"log/slog"
"strings"
+ "time"
"github.com/apache/arrow-adbc/go/adbc"
"github.com/apache/arrow-adbc/go/adbc/driver/internal"
@@ -153,8 +154,14 @@ func (base *ConnectionImplBase) Rollback(context.Context)
error {
}
func (base *ConnectionImplBase) GetInfo(ctx context.Context, infoCodes
[]adbc.InfoCode) (reader array.RecordReader, err error) {
+ startTime := time.Now()
_, span := internal.StartSpan(ctx, "ConnectionImplBase.GetInfo", base)
- defer internal.EndSpanWithError(span, &err)
+ defer func() {
+ internal.NewEndSpanHelper(span).
+ WithError(err).
+ WithStartTime(startTime).
+ EndSpan()
+ }()
if len(infoCodes) == 0 {
infoCodes = base.DriverInfo.InfoSupportedCodes()
diff --git a/go/adbc/driver/internal/driverbase/database.go
b/go/adbc/driver/internal/driverbase/database.go
index 991e7270e..fb664f719 100644
--- a/go/adbc/driver/internal/driverbase/database.go
+++ b/go/adbc/driver/internal/driverbase/database.go
@@ -106,8 +106,10 @@ type DatabaseImplBase struct {
Logger *slog.Logger
Tracer trace.Tracer
- tracerShutdownFunc func(context.Context) error
- traceParent string
+ tracerForceFlushFunc func(context.Context) error
+ tracerShutdownFunc func(context.Context) error
+ tracerProvider trace.TracerProvider
+ traceParent string
}
type TracingOptions struct {
@@ -128,11 +130,12 @@ type TracingOptions struct {
// driver, allowing the Arrow allocator and error handler to be reused.
func NewDatabaseImplBase(ctx context.Context, driver *DriverImplBase, opts
TracingOptions) (DatabaseImplBase, error) {
database := DatabaseImplBase{
- Alloc: driver.Alloc,
- ErrorHelper: driver.ErrorHelper,
- DriverInfo: driver.DriverInfo,
- Logger: nilLogger(),
- Tracer: nilTracer(),
+ Alloc: driver.Alloc,
+ ErrorHelper: driver.ErrorHelper,
+ DriverInfo: driver.DriverInfo,
+ Logger: nilLogger(),
+ Tracer: nilTracer(),
+ tracerProvider: otel.GetTracerProvider(),
}
err := database.InitTracing(
ctx,
@@ -180,17 +183,25 @@ func (base *DatabaseImplBase) SetOptionInt(key string,
val int64) error {
}
func (base *database) Close() error {
- return base.Base().Close()
+ return base.DatabaseImpl.Close()
}
func (base *DatabaseImplBase) Close() (err error) {
if base.Base().tracerShutdownFunc != nil {
err = base.Base().tracerShutdownFunc(context.Background())
base.Base().tracerShutdownFunc = nil
+ base.Base().tracerForceFlushFunc = nil
}
return
}
+func (base *DatabaseImplBase) ForceFlushTracing(ctx context.Context) error {
+ if base.Base().tracerForceFlushFunc == nil {
+ return nil
+ }
+ return base.Base().tracerForceFlushFunc(ctx)
+}
+
func (base *DatabaseImplBase) Open(ctx context.Context) (adbc.Connection,
error) {
return nil, base.ErrorHelper.Errorf(adbc.StatusNotImplemented, "Open")
}
@@ -225,6 +236,10 @@ func (d *DatabaseImplBase) StartSpan(
return d.Tracer.Start(ctx, spanName, opts...)
}
+func (d *DatabaseImplBase) GetTracerProvider() trace.TracerProvider {
+ return d.tracerProvider
+}
+
// database is the implementation of adbc.Database.
type database struct {
DatabaseImpl
@@ -264,6 +279,7 @@ func (base *DatabaseImplBase) InitTracing(
// Empty exporter
if exporterName == "" {
+ base.tracerProvider = otel.GetTracerProvider()
base.Tracer = otel.Tracer(fullyQualifiedDriverName)
return
}
@@ -361,7 +377,9 @@ func newTracer(
if err != nil {
return
}
+ base.Base().tracerForceFlushFunc = tracerProvider.ForceFlush
base.Base().tracerShutdownFunc = tracerProvider.Shutdown
+ base.Base().tracerProvider = tracerProvider
tracer = tracerProvider.Tracer(
fullyQualifiedDriverName,
trace.WithInstrumentationVersion(driverVersion),
diff --git a/go/adbc/driver/internal/shared_utils.go
b/go/adbc/driver/internal/shared_utils.go
index d44578d40..2f5b04b97 100644
--- a/go/adbc/driver/internal/shared_utils.go
+++ b/go/adbc/driver/internal/shared_utils.go
@@ -22,11 +22,13 @@ import (
"regexp"
"strconv"
"strings"
+ "time"
"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/attribute"
"go.opentelemetry.io/otel/codes"
semconv "go.opentelemetry.io/otel/semconv/v1.30.0"
"go.opentelemetry.io/otel/trace"
@@ -781,12 +783,88 @@ func EndSpan(span trace.Span, err error, options
...trace.SpanEndOption) {
func EndSpanWithError(span trace.Span, err *error, options
...trace.SpanEndOption) {
if err != nil && *err != nil {
span.RecordError(*err)
- if adbcError, ok := (*err).(adbc.Error); ok {
+ setSpanStatus(span, *err)
+ } else {
+ setSpanStatus(span, nil)
+ }
+ span.End(options...)
+}
+
+// EndSpanHelper centralizes span completion for operations that may fail.
+// It records errors when present, records the elapsed duration if a start time
+// was provided, and ensures the span status and end options are applied
uniformly.
+type EndSpanHelper struct {
+ span trace.Span
+ err *error
+ errorRecorded bool
+ startTime *time.Time
+ options []trace.SpanEndOption
+}
+
+// NewEndSpanHelper creates a helper for finalizing a span after an operation.
+func NewEndSpanHelper(span trace.Span) *EndSpanHelper {
+ return &EndSpanHelper{
+ span: span,
+ errorRecorded: false,
+ startTime: nil,
+ err: nil,
+ }
+}
+
+// WithError configures the error to be recorded and reported on the span.
+// It returns the helper itself to allow fluent chaining.
+func (h *EndSpanHelper) WithError(err error) *EndSpanHelper {
+ h.err = &err
+ return h
+}
+
+// WithRecordedError marks whether an error has already been recorded
elsewhere.
+// It returns the helper itself to allow fluent chaining.
+func (h *EndSpanHelper) WithRecordedError(errorRecorded bool) *EndSpanHelper {
+ h.errorRecorded = errorRecorded
+ return h
+}
+
+// WithOptions sets the span end options to apply when the span is closed.
+// It returns the helper itself to allow fluent chaining.
+func (h *EndSpanHelper) WithOptions(options ...trace.SpanEndOption)
*EndSpanHelper {
+ h.options = options
+ return h
+}
+
+// WithStartTime sets the operation start time so the elapsed duration can be
recorded.
+// It returns the helper itself to allow fluent chaining.
+func (h *EndSpanHelper) WithStartTime(startTime time.Time) *EndSpanHelper {
+ h.startTime = &startTime
+ return h
+}
+
+// EndSpan completes the span, records any error, sets status, and emits
duration metadata if available.
+func (h *EndSpanHelper) EndSpan() {
+ if h.span == nil {
+ return
+ }
+ if h.startTime != nil {
+ h.span.SetAttributes(attribute.Float64("span.duration_s",
time.Since(*h.startTime).Seconds()))
+ }
+ if !h.errorRecorded && h.err != nil && *h.err != nil {
+ h.span.RecordError(*h.err)
+ }
+ if h.err != nil && *h.err != nil {
+ setSpanStatus(h.span, *h.err)
+ } else {
+ setSpanStatus(h.span, nil)
+ }
+ h.span.End(h.options...)
+}
+
+func setSpanStatus(span trace.Span, err error) {
+ if err != nil {
+ if adbcError, ok := err.(adbc.Error); ok {
span.SetAttributes(semconv.ErrorTypeKey.String(adbcError.Code.String()))
}
- span.SetStatus(codes.Error, (*err).Error())
+ span.SetStatus(codes.Error, err.Error())
} else {
span.SetStatus(codes.Ok, "")
}
- span.End(options...)
}
