Skip to content

refactor(go/adbc): new utility method StartSpanWithEndSpanHelper #4725

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
birschick-bq wants to merge 5 commits into
base: main
Choose a base branch
from
Open

refactor(go/adbc): new utility method StartSpanWithEndSpanHelper #4725

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c8fa7b4
feat(go/adbc): add overload of StartSpan to return EndSpanHelper
v-birschick-msft Aug 25, 2026
77919cb
Merge branch 'main' into dev/birschick-bq/flight-sql-endspan-helper
v-birschick-msft Aug 27, 2026
41aa382
feat(go/adbc): new utility method StartSpanWithEndSpanHelper
v-birschick-msft Aug 27, 2026
91e0ece
empty: retest
v-birschick-msft Aug 27, 2026
0e60280
attempt fix to c lang build errors on Mac
v-birschick-msft Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 17 additions & 17 deletions c/driver/framework/status.h
View file
Open in desktop

@lidavidm lidavidm Aug 28, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I merged a different fix for this already

@lidavidm lidavidm Aug 28, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Albeit I like this fix better, if you want to rebase and take out the other fix

Original file line number Diff line number Diff line change
Expand Up @@ -129,12 +129,12 @@ class Status {
// Helpers to create statuses with known codes
static Status Ok() { return Status(); }

#define STATUS_CTOR(NAME, CODE) \
template <typename... Args> \
static Status NAME(Args&&... args) { \
std::stringstream ss; \
([&] { ss << args; }(), ...); \
return Status(ADBC_STATUS_##CODE, ss.str()); \
#define STATUS_CTOR(NAME, CODE) \
template <typename... Args> \
[[maybe_unused]] static Status NAME(Args&&... args) { \
std::stringstream ss; \
([&] { ss << args; }(), ...); \
return Status(ADBC_STATUS_##CODE, ss.str()); \
}

STATUS_CTOR(Internal, INTERNAL)
Expand Down Expand Up @@ -305,12 +305,12 @@ namespace adbc::driver::status {

inline driver::Status Ok() { return driver::Status(); }

#define STATUS_CTOR(NAME, CODE) \
template <typename... Args> \
static Status NAME(Args&&... args) { \
std::stringstream ss; \
([&] { ss << args; }(), ...); \
return Status(ADBC_STATUS_##CODE, ss.str()); \
#define STATUS_CTOR(NAME, CODE) \
template <typename... Args> \
[[maybe_unused]] static Status NAME(Args&&... args) { \
std::stringstream ss; \
([&] { ss << args; }(), ...); \
return Status(ADBC_STATUS_##CODE, ss.str()); \
}

// TODO: unit tests for internal utilities
Expand All @@ -329,11 +329,11 @@ STATUS_CTOR(Unknown, UNKNOWN)
#if defined(ADBC_FRAMEWORK_USE_FMT)
namespace adbc::driver::status::fmt {

#define STATUS_CTOR(NAME, CODE) \
template <typename... Args> \
static Status NAME(std::string_view format_string, Args&&... args) { \
auto message = ::fmt::vformat(format_string, ::fmt::make_format_args(args...)); \
return Status(ADBC_STATUS_##CODE, std::move(message)); \
#define STATUS_CTOR(NAME, CODE) \
template <typename... Args> \
[[maybe_unused]] static Status NAME(std::string_view format_string, Args&&... args) { \
auto message = ::fmt::vformat(format_string, ::fmt::make_format_args(args...)); \
return Status(ADBC_STATUS_##CODE, std::move(message)); \
}

// TODO: unit tests for internal utilities
Expand Down
8 changes: 2 additions & 6 deletions go/adbc/driver/flightsql/flightsql_bulk_ingest.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,14 +109,10 @@ func createRecordReaderFromBatch(batch arrow.RecordBatch) (array.RecordReader, e
// This is called from the statement when a target table has been set for bulk ingest.
func (s *statement) executeIngest(ctx context.Context) (nRows int64, err error) {
var startTime = time.Now()
ctx, span := internal.StartSpan(ctx, "FlightSQL.BulkIngest.Execute", s.cnxn)
ctx, span, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, "FlightSQL.BulkIngest.Execute", s.cnxn)
errorRecorded := false
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
WithRecordedError(errorRecorded).
EndSpan()
endSpanHelper.WithError(err).WithRecordedError(errorRecorded).EndSpan()
}()

if s.streamBind == nil && s.bound == nil {
Expand Down
73 changes: 18 additions & 55 deletions go/adbc/driver/flightsql/flightsql_connection.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -247,15 +247,10 @@ func doGetWithResponseMetadata(ctx context.Context, client *flightsql.Client, ti

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)
ctx, span, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, spanName, tracing)
errorRecorded := false
defer func() {
internal.NewEndSpanHelper(span).
WithStartTime(startTime).
WithError(err).
WithRecordedError(errorRecorded).
EndSpan()
endSpanHelper.WithError(err).WithRecordedError(errorRecorded).EndSpan()
}()

streamOpts := make([]grpc.CallOption, 0, len(opts))
Expand Down Expand Up @@ -737,14 +732,10 @@ func (c *connectionImpl) SetOptionDouble(key string, value float64) error {
}

func (c *connectionImpl) PrepareDriverInfo(ctx context.Context, infoCodes []adbc.InfoCode) (err error) {
startTime := time.Now()
const spanName = "FlightSQL.Connection.PrepareDriverInfo"
ctx, span := internal.StartSpan(ctx, spanName, c)
ctx, _, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, spanName, c)
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()

driverInfo := c.DriverInfo
Expand Down Expand Up @@ -867,13 +858,9 @@ func (c *connectionImpl) readInfo(ctx context.Context, expectedSchema *arrow.Sch

func (c *connectionImpl) GetObjectsCatalogs(ctx context.Context, catalog *string) (catalogs []string, err error) {
const spanName = "FlightSQL.Connection.GetObjectsCatalogs"
startTime := time.Now()
ctx, span := internal.StartSpan(ctx, spanName, c)
ctx, _, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, spanName, c)
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()
var (
header, trailer metadata.MD
Expand Down Expand Up @@ -919,13 +906,9 @@ func (c *connectionImpl) GetObjectsCatalogs(ctx context.Context, catalog *string
// Helper function to build up a map of catalogs to DB schemas
func (c *connectionImpl) GetObjectsDbSchemas(ctx context.Context, depth adbc.ObjectDepth, catalog *string, dbSchema *string) (result map[string][]string, err error) {
const spanName = "FlightSQL.Connection.GetObjectsDbSchemas"
startTime := time.Now()
ctx, span := internal.StartSpan(ctx, spanName, c)
ctx, _, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, spanName, c)
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()
if depth == adbc.ObjectDepthCatalogs {
return
Expand Down Expand Up @@ -972,13 +955,9 @@ func (c *connectionImpl) GetObjectsDbSchemas(ctx context.Context, depth adbc.Obj

func (c *connectionImpl) GetObjectsTables(ctx context.Context, depth adbc.ObjectDepth, catalog *string, dbSchema *string, tableName *string, columnName *string, tableType []string) (result internal.SchemaToTableInfo, err error) {
const spanName = "FlightSQL.Connection.GetObjectsTables"
startTime := time.Now()
ctx, span := internal.StartSpan(ctx, spanName, c)
ctx, _, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, spanName, c)
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()

if depth == adbc.ObjectDepthCatalogs || depth == adbc.ObjectDepthDBSchemas {
Expand Down Expand Up @@ -1066,13 +1045,9 @@ func (c *connectionImpl) GetObjectsTables(ctx context.Context, depth adbc.Object

func (c *connectionImpl) GetTableSchema(ctx context.Context, catalog *string, dbSchema *string, tableName string) (schema *arrow.Schema, err error) {
const spanName = "FlightSQL.Connection.GetTableSchema"
startTime := time.Now()
ctx, span := internal.StartSpan(ctx, spanName, c)
ctx, _, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, spanName, c)
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()

opts := &flightsql.GetTablesOpts{
Expand Down Expand Up @@ -1157,13 +1132,9 @@ func (c *connectionImpl) GetTableSchema(ctx context.Context, catalog *string, db
// table_type | utf8 not null
func (c *connectionImpl) GetTableTypes(ctx context.Context) (reader array.RecordReader, err error) {
const spanName = "FlightSQL.Connection.GetTableTypes"
startTime := time.Now()
ctx, span := internal.StartSpan(ctx, spanName, c)
ctx, _, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, spanName, c)
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()

ctx = metadata.NewOutgoingContext(ctx, c.hdrs)
Expand Down Expand Up @@ -1335,13 +1306,9 @@ func (c *connectionImpl) prepareSubstrait(ctx context.Context, plan flightsql.Su
// Close closes this connection and releases any associated resources.
func (c *connectionImpl) Close() (err error) {
const spanName = "FlightSQL.Connection.Close"
startTime := time.Now()
ctx, span := internal.StartSpan(context.Background(), spanName, c)
ctx, span, endSpanHelper := internal.StartSpanWithEndSpanHelper(context.Background(), spanName, c)
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()

if c.cl == nil {
Expand Down Expand Up @@ -1395,13 +1362,9 @@ func (c *connectionImpl) Close() (err error) {
// A partition can be retrieved by using ExecutePartitions on a statement.
func (c *connectionImpl) ReadPartition(ctx context.Context, serializedPartition []byte) (rdr array.RecordReader, err error) {
const spanName = "FlightSQL.Connection.ReadPartition"
startTime := time.Now()
ctx, span := internal.StartSpan(ctx, spanName, c)
ctx, _, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, spanName, c)
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()

var info flight.FlightInfo
Expand Down
31 changes: 8 additions & 23 deletions go/adbc/driver/flightsql/flightsql_database.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -373,16 +373,11 @@ func (d *databaseImpl) SetOptionDouble(key string, value float64) error {

func (d *databaseImpl) Close() (err error) {
const spanName = "FlightSQL.Database.Close"
startTime := time.Now()
var span trace.Span
_, span = internal.StartSpan(context.Background(), spanName, d)
_, span, endSpanHelper := internal.StartSpanWithEndSpanHelper(context.Background(), spanName, d)

span.AddEvent("closing", trace.WithAttributes(attribute.String("target", d.uri.String())))
flushErr := d.ForceFlushTracing(context.Background())
internal.NewEndSpanHelper(span).
WithError(flushErr).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(flushErr).EndSpan()
shutdownErr := d.DatabaseImplBase.Close()
return errors.Join(flushErr, shutdownErr)
}
Expand Down Expand Up @@ -527,34 +522,28 @@ type support struct {
}

func closeCachedFlightClient(d *databaseImpl, location, client interface{}, reason string) {
startTime := time.Now()
var err error
_, span := internal.StartSpan(context.Background(), "FlightSQL.Database.CloseCachedClient", d,
_, _, endSpanHelper := internal.StartSpanWithEndSpanHelper(context.Background(), "FlightSQL.Database.CloseCachedClient", d,
trace.WithAttributes(
attribute.String("flight.location", fmt.Sprint(location)),
attribute.String("flight.cache.reason", reason),
))
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()

err = client.(*flightsql.Client).Close()
}

func (d *databaseImpl) Open(ctx context.Context) (_ adbc.Connection, err error) {
ctx, span := internal.StartSpan(
ctx, span, endSpanHelper := internal.StartSpanWithEndSpanHelper(
ctx,
"FlightSQL.Database.Open",
d,
trace.WithAttributes(traceHeaderAttrsWithPrefix(d.hdrs, traceRequestMetadataPrefix)...),
)
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()

authMiddle := &bearerAuthMiddleware{
Expand All @@ -573,13 +562,9 @@ func (d *databaseImpl) Open(ctx context.Context) (_ adbc.Connection, err error)
cache := gcache.New(20).LRU().
Expiration(5 * time.Minute).
LoaderFunc(func(loc interface{}) (_ interface{}, err error) {
startTime := time.Now()
ctx, cacheSpan := internal.StartSpan(context.Background(), "FlightSQL.Database.LoadCachedClient", d)
ctx, cacheSpan, endCacheSpanHelper := internal.StartSpanWithEndSpanHelper(context.Background(), "FlightSQL.Database.LoadCachedClient", d)
defer func() {
internal.NewEndSpanHelper(cacheSpan).
WithError(err).
WithStartTime(startTime).
EndSpan()
endCacheSpanHelper.WithError(err).EndSpan()
}()

uri, ok := loc.(string)
Expand Down
21 changes: 6 additions & 15 deletions go/adbc/driver/flightsql/flightsql_statement.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -523,12 +523,9 @@ func (s *statement) ExecuteQuery(ctx context.Context) (rdr array.RecordReader, n
spanName = "FlightSQL.Statement." + operationName
)
startTime := time.Now()
ctx, span := internal.StartSpan(ctx, spanName, s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...))
ctx, span, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, spanName, s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...))
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()

if err = s.clearIncrementalQuery(); err != nil {
Expand Down Expand Up @@ -602,12 +599,9 @@ func (s *statement) ExecuteUpdate(ctx context.Context) (n int64, err error) {
spanName = "FlightSQL.Statement." + operationName
)
startTime := time.Now()
ctx, span := internal.StartSpan(ctx, spanName, s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...))
ctx, span, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, spanName, s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...))
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()

if err = s.clearIncrementalQuery(); err != nil {
Expand Down Expand Up @@ -667,12 +661,9 @@ func (s *statement) Prepare(ctx context.Context) (err error) {
spanName = "FlightSQL.Statement." + operationName
)
startTime := time.Now()
ctx, span := internal.StartSpan(ctx, spanName, s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...))
ctx, span, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, spanName, s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...))
defer func() {
internal.NewEndSpanHelper(span).
WithError(err).
WithStartTime(startTime).
EndSpan()
endSpanHelper.WithError(err).EndSpan()
}()

span.AddEvent("starting", trace.WithAttributes(s.queryAttrs()...))
Expand Down
16 changes: 3 additions & 13 deletions go/adbc/driver/flightsql/record_reader.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ import (
"errors"
"fmt"
"sync/atomic"
"time"

"github.com/apache/arrow-adbc/go/adbc"
"github.com/apache/arrow-adbc/go/adbc/driver/internal"
Expand Down Expand Up @@ -78,17 +77,12 @@ type recordReaderConfig struct {
// 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) {
const spanName = "FlightSQL.RecordReader.newRecordReader"
startTime := time.Now()
ctx, span := internal.StartSpan(ctx, spanName, cfg.tracing)
ctx, span, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, spanName, cfg.tracing)
spanOwnedByReader := false
errorRecorded := false
defer func() {
if !spanOwnedByReader {
internal.NewEndSpanHelper(span).
WithStartTime(startTime).
WithError(err).
WithRecordedError(errorRecorded).
EndSpan()
endSpanHelper.WithError(err).WithRecordedError(errorRecorded).EndSpan()
}
}()

Expand Down Expand Up @@ -315,11 +309,7 @@ func newRecordReader(ctx context.Context, cfg recordReaderConfig, opts ...grpc.C
))
}
errorRecorded := reader.err != nil
internal.NewEndSpanHelper(span).
WithStartTime(startTime).
WithError(reader.err).
WithRecordedError(errorRecorded).
EndSpan()
endSpanHelper.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 and tracing have been finalized.
close(chs[lastChannelIndex])
Expand Down
4 changes: 2 additions & 2 deletions go/adbc/driver/flightsql/tracing_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@ func TestTraceHeaderAttrsWithPrefix_AppliedToSpan(t *testing.T) {
},
}

ctx, span := internal.StartSpan(
ctx, _, endSpanHelper := internal.StartSpanWithEndSpanHelper(
context.Background(),
"FlightSQL.Statement.ExecuteQuery",
tracing,
Expand All @@ -155,7 +155,7 @@ func TestTraceHeaderAttrsWithPrefix_AppliedToSpan(t *testing.T) {
}), traceRequestMetadataPrefix)...),
)
_ = ctx
span.End()
endSpanHelper.EndSpan()

spans := recorder.Ended()
if len(spans) != 1 {
Expand Down
Loading
Loading