diff --git a/c/driver/framework/status.h b/c/driver/framework/status.h index d42535e5ac..1fac2272ae 100644 --- a/c/driver/framework/status.h +++ b/c/driver/framework/status.h @@ -129,12 +129,12 @@ class Status { // Helpers to create statuses with known codes static Status Ok() { return Status(); } -#define STATUS_CTOR(NAME, CODE) \ - template \ - static Status NAME(Args&&... args) { \ - std::stringstream ss; \ - ([&] { ss << args; }(), ...); \ - return Status(ADBC_STATUS_##CODE, ss.str()); \ +#define STATUS_CTOR(NAME, CODE) \ + template \ + [[maybe_unused]] static Status NAME(Args&&... args) { \ + std::stringstream ss; \ + ([&] { ss << args; }(), ...); \ + return Status(ADBC_STATUS_##CODE, ss.str()); \ } STATUS_CTOR(Internal, INTERNAL) @@ -305,12 +305,12 @@ namespace adbc::driver::status { inline driver::Status Ok() { return driver::Status(); } -#define STATUS_CTOR(NAME, CODE) \ - template \ - static Status NAME(Args&&... args) { \ - std::stringstream ss; \ - ([&] { ss << args; }(), ...); \ - return Status(ADBC_STATUS_##CODE, ss.str()); \ +#define STATUS_CTOR(NAME, CODE) \ + template \ + [[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 @@ -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 \ - 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 \ + [[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 diff --git a/go/adbc/driver/flightsql/flightsql_bulk_ingest.go b/go/adbc/driver/flightsql/flightsql_bulk_ingest.go index b888a3c2a6..b5e4dc185b 100644 --- a/go/adbc/driver/flightsql/flightsql_bulk_ingest.go +++ b/go/adbc/driver/flightsql/flightsql_bulk_ingest.go @@ -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 { diff --git a/go/adbc/driver/flightsql/flightsql_connection.go b/go/adbc/driver/flightsql/flightsql_connection.go index 692f573e10..291228d1a8 100644 --- a/go/adbc/driver/flightsql/flightsql_connection.go +++ b/go/adbc/driver/flightsql/flightsql_connection.go @@ -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)) @@ -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 @@ -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 @@ -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 @@ -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 { @@ -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{ @@ -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) @@ -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 { @@ -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 diff --git a/go/adbc/driver/flightsql/flightsql_database.go b/go/adbc/driver/flightsql/flightsql_database.go index 10b368c325..650fd4c68d 100644 --- a/go/adbc/driver/flightsql/flightsql_database.go +++ b/go/adbc/driver/flightsql/flightsql_database.go @@ -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) } @@ -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{ @@ -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) diff --git a/go/adbc/driver/flightsql/flightsql_statement.go b/go/adbc/driver/flightsql/flightsql_statement.go index 3e6808f3a7..ec11ec9bd4 100644 --- a/go/adbc/driver/flightsql/flightsql_statement.go +++ b/go/adbc/driver/flightsql/flightsql_statement.go @@ -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 { @@ -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 { @@ -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()...)) diff --git a/go/adbc/driver/flightsql/record_reader.go b/go/adbc/driver/flightsql/record_reader.go index 289fc49a26..c3f4c2741a 100644 --- a/go/adbc/driver/flightsql/record_reader.go +++ b/go/adbc/driver/flightsql/record_reader.go @@ -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" @@ -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() } }() @@ -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]) diff --git a/go/adbc/driver/flightsql/tracing_test.go b/go/adbc/driver/flightsql/tracing_test.go index be1771582a..b88f7a19b4 100644 --- a/go/adbc/driver/flightsql/tracing_test.go +++ b/go/adbc/driver/flightsql/tracing_test.go @@ -144,7 +144,7 @@ func TestTraceHeaderAttrsWithPrefix_AppliedToSpan(t *testing.T) { }, } - ctx, span := internal.StartSpan( + ctx, _, endSpanHelper := internal.StartSpanWithEndSpanHelper( context.Background(), "FlightSQL.Statement.ExecuteQuery", tracing, @@ -155,7 +155,7 @@ func TestTraceHeaderAttrsWithPrefix_AppliedToSpan(t *testing.T) { }), traceRequestMetadataPrefix)...), ) _ = ctx - span.End() + endSpanHelper.EndSpan() spans := recorder.Ended() if len(spans) != 1 { diff --git a/go/adbc/driver/internal/driverbase/connection.go b/go/adbc/driver/internal/driverbase/connection.go index 5db09c8350..05f968fde8 100644 --- a/go/adbc/driver/internal/driverbase/connection.go +++ b/go/adbc/driver/internal/driverbase/connection.go @@ -25,7 +25,6 @@ import ( "fmt" "log/slog" "strings" - "time" "github.com/apache/arrow-adbc/go/adbc" "github.com/apache/arrow-adbc/go/adbc/driver/internal" @@ -154,13 +153,9 @@ 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) + _, _, endSpanHelper := internal.StartSpanWithEndSpanHelper(ctx, "ConnectionImplBase.GetInfo", base) defer func() { - internal.NewEndSpanHelper(span). - WithError(err). - WithStartTime(startTime). - EndSpan() + endSpanHelper.WithError(err).EndSpan() }() if len(infoCodes) == 0 { diff --git a/go/adbc/driver/internal/shared_utils.go b/go/adbc/driver/internal/shared_utils.go index 2f5b04b97c..8446d0f6e6 100644 --- a/go/adbc/driver/internal/shared_utils.go +++ b/go/adbc/driver/internal/shared_utils.go @@ -761,6 +761,23 @@ func StartSpan(ctx context.Context, spanName string, tracing adbc.OTelTracing, o return tracing.StartSpan(ctx, spanName, opts...) } +// Starts a trace.Span with the given spanName for the tracing object with +// the given ctx context. Returns an EndSpanHelper to centralize span completion for operations that may fail. +func StartSpanWithEndSpanHelper(ctx context.Context, spanName string, tracing adbc.OTelTracing, opts ...trace.SpanStartOption) (context.Context, trace.Span, *EndSpanHelper) { + startTime := time.Now() + if tracing == nil { + span := trace.SpanFromContext(ctx) + return ctx, span, NewEndSpanHelper(span).WithStartTime(startTime) + } + + attrs := tracing.GetInitialSpanAttributes() + attrs = append(attrs, semconv.DBOperationName(spanName)) + opts = append(opts, trace.WithAttributes(attrs...)) + + newCtx, span := tracing.StartSpan(ctx, spanName, opts...) + return newCtx, span, NewEndSpanHelper(span).WithStartTime(startTime) +} + // Ends the given span. If err is not nil, then the // error is recorded and the status is set appropriately. // Otherwise, the status is set to Ok.