diff --git a/go/adbc/driver/flightsql/flightsql_adbc_test.go b/go/adbc/driver/flightsql/flightsql_adbc_test.go index 2be9496e9a..edc6c804f4 100644 --- a/go/adbc/driver/flightsql/flightsql_adbc_test.go +++ b/go/adbc/driver/flightsql/flightsql_adbc_test.go @@ -368,7 +368,7 @@ func TestFlightSQLTracingProducesTraceFiles(t *testing.T) { } output := traceOutput.String() - require.Contains(t, output, "FlightSQLDatabase.Open") + require.Contains(t, output, "FlightSQL.Database.Open") require.Contains(t, output, "FlightSQLStatement.ExecuteQuery") } diff --git a/go/adbc/driver/flightsql/flightsql_connection.go b/go/adbc/driver/flightsql/flightsql_connection.go index 11d9a6473d..95e9a9b61a 100644 --- a/go/adbc/driver/flightsql/flightsql_connection.go +++ b/go/adbc/driver/flightsql/flightsql_connection.go @@ -23,7 +23,6 @@ import ( "encoding/json" "fmt" "io" - "log/slog" "math" "strings" "time" @@ -39,6 +38,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,22 +233,59 @@ var adbcToFlightSQLInfo = map[adbc.InfoCode]flightsql.SqlInfo{ adbc.InfoVendorSubstraitMaxVersion: flightsql.SqlInfoFlightSqlServerSubstraitMaxVersion, } -// doGetWithLogger performs DoGet against an endpoint's locations, logging each +// doGetWithTracer 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. -func doGetWithLogger(ctx context.Context, cl *flightsql.Client, endpoint *flight.FlightEndpoint, clientCache gcache.Cache, logger *slog.Logger, opts ...grpc.CallOption) (rdr *flight.Reader, err error) { - log := safeLogger(logger) +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" + var 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 { - log.DebugContext(ctx, "FlightSQL doGet", - "phase", "noLocations", - ) + span.AddEvent("flight.location.attempt", trace.WithAttributes( + attribute.String("flight.location.source", "default_client"), + )) start := time.Now() - rdr, err = cl.DoGet(ctx, endpoint.Ticket, opts...) - log.DebugContext(ctx, "FlightSQL doGet", - "phase", "defaultClientResult", - "duration", time.Since(start), - "err", err, - ) + 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 } @@ -264,52 +302,69 @@ func doGetWithLogger(ctx context.Context, cl *flightsql.Client, endpoint *flight } 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())) - log.WarnContext(ctx, "FlightSQL doGet location attempt failed", - "phase", "clientCacheGet", - "location", loc.Uri, - "duration", time.Since(start), - "err", err, - ) + 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 = conn.DoGet(ctx, endpoint.Ticket, opts...) + rdr, err = doGetWithResponseMetadata(ctx, conn, endpoint.Ticket, streamOpts...) if err != nil { attemptErrors = append(attemptErrors, fmt.Sprintf("DoGet(%q): %s", loc.Uri, err.Error())) - log.WarnContext(ctx, "FlightSQL doGet location attempt failed", - "phase", "doGet", - "location", loc.Uri, - "duration", time.Since(start), - "err", err, - ) + 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 } - log.DebugContext(ctx, "FlightSQL doGet succeeded", - "location", loc.Uri, - "duration", time.Since(start), - ) + 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() - rdr, err = cl.DoGet(ctx, endpoint.Ticket, opts...) + 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())) - log.WarnContext(ctx, "FlightSQL doGet fallback to default client failed", - "duration", time.Since(start), - "err", err, - ) - return nil, fmt.Errorf("all DoGet attempts failed: %s; final: %w", strings.Join(attemptErrors, "; "), err) + 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 } - log.DebugContext(ctx, "FlightSQL doGet succeeded via default client fallback", - "duration", time.Since(start), - ) + span.AddEvent("flight.location.selected", trace.WithAttributes( + attribute.String("flight.location.source", "fallback"), + attribute.Float64("duration_s", time.Since(start).Seconds()), + )) return rdr, nil } @@ -317,6 +372,13 @@ func doGetWithLogger(ctx context.Context, cl *flightsql.Client, endpoint *flight 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 } @@ -669,7 +731,12 @@ func (c *connectionImpl) SetOptionDouble(key string, value float64) error { return c.ConnectionImplBase.SetOptionDouble(key, value) } -func (c *connectionImpl) PrepareDriverInfo(ctx context.Context, infoCodes []adbc.InfoCode) 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) + defer internal.EndSpanWithStartTime(span, &err, &startTime) + driverInfo := c.DriverInfo if len(infoCodes) == 0 { @@ -690,7 +757,8 @@ func (c *connectionImpl) PrepareDriverInfo(ctx context.Context, infoCodes []adbc ctx = metadata.NewOutgoingContext(ctx, c.hdrs) var header, trailer metadata.MD - info, err := c.cl.GetSqlInfo(ctx, translated, grpc.Header(&header), grpc.Trailer(&trailer), c.timeouts) + var info *flight.FlightInfo + info, err = c.cl.GetSqlInfo(ctx, translated, grpc.Header(&header), grpc.Trailer(&trailer), c.timeouts) // Just return local driver info if GetSqlInfo hasn't been implemented on the server if grpcstatus.Code(err) == grpccodes.Unimplemented { @@ -703,10 +771,12 @@ func (c *connectionImpl) PrepareDriverInfo(ctx context.Context, infoCodes []adbc // No error, go get the SqlInfo from the server for i, endpoint := range info.Endpoint { - var header, trailer metadata.MD - rdr, err := doGetWithLogger(ctx, c.cl, endpoint, c.clientCache, c.Logger, grpc.Header(&header), grpc.Trailer(&trailer), c.timeouts) + var responseMetadata *responseMetadataCollector + ctx, responseMetadata = withResponseMetadata(ctx) + var rdr *flight.Reader + rdr, err = doGetWithTracer(ctx, c.cl, endpoint, c.clientCache, c, c.timeouts) if err != nil { - return adbcFromFlightStatusWithDetails(err, header, trailer, "GetInfo(DoGet): endpoint %d: %s", i, endpoint.Location) + return adbcFromFlightStatusWithDetails(err, responseMetadata.snapshot(), nil, "GetInfo(DoGet): endpoint %d: %s", i, endpoint.Location) } for rdr.Next() { @@ -740,20 +810,22 @@ func (c *connectionImpl) PrepareDriverInfo(ctx context.Context, infoCodes []adbc case *array.Boolean: v = arr.Value(idx) default: - return adbc.Error{ + err = adbc.Error{ Msg: fmt.Sprintf("unsupported field_type %T for info_value", arr), Code: adbc.StatusInvalidArgument, } + return err } - if err := driverInfo.RegisterInfoCode(adbcInfoCode, v); err != nil { + if err = driverInfo.RegisterInfoCode(adbcInfoCode, v); err != nil { return err } } } - if err := checkContext(rdr.Err(), ctx); err != nil { - return adbcFromFlightStatusWithDetails(err, header, trailer, "GetInfo(DoGet): endpoint %d: %s", i, endpoint.Location) + if err = checkContext(rdr.Err(), ctx); err != nil { + err = adbcFromFlightStatusWithDetails(err, responseMetadata.snapshot(), nil, "GetInfo(DoGet): endpoint %d: %s", i, endpoint.Location) + return err } } @@ -769,7 +841,7 @@ func (c *connectionImpl) readInfo(ctx context.Context, expectedSchema *arrow.Sch info: info, clientCache: c.clientCache, bufferSize: 5, - logger: c.Logger, + tracing: c, }, opts...) if err != nil { return nil, adbcFromFlightStatus(err, "DoGet") @@ -785,7 +857,11 @@ func (c *connectionImpl) readInfo(ctx context.Context, expectedSchema *arrow.Sch return rdr, nil } -func (c *connectionImpl) GetObjectsCatalogs(ctx context.Context, catalog *string) ([]string, error) { +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) + defer internal.EndSpanWithStartTime(span, &err, &startTime) var ( header, trailer metadata.MD numCatalogs int64 @@ -803,13 +879,15 @@ func (c *connectionImpl) GetObjectsCatalogs(ctx context.Context, catalog *string header = metadata.MD{} trailer = metadata.MD{} - rdr, err := c.readInfo(ctx, schema_ref.Catalogs, info, c.timeouts, grpc.Header(&header), grpc.Trailer(&trailer)) + var rdr array.RecordReader + rdr, err = c.readInfo(ctx, schema_ref.Catalogs, info, c.timeouts, grpc.Header(&header), grpc.Trailer(&trailer)) if err != nil { - return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetCatalogs)") + err = adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetCatalogs)") + return nil, err } defer rdr.Release() - catalogs := make([]string, 0, numCatalogs) + catalogs = make([]string, 0, numCatalogs) for rdr.Next() { arr := rdr.RecordBatch().Column(0).(*array.String) for i := 0; i < arr.Len(); i++ { @@ -819,8 +897,9 @@ func (c *connectionImpl) GetObjectsCatalogs(ctx context.Context, catalog *string } } - if err := checkContext(rdr.Err(), ctx); err != nil { - return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetCatalogs)") + if err = checkContext(rdr.Err(), ctx); err != nil { + err = adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetCatalogs)") + return nil, err } return catalogs, nil @@ -828,6 +907,10 @@ 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) + defer internal.EndSpanWithStartTime(span, &err, &startTime) if depth == adbc.ObjectDepthCatalogs { return } @@ -835,16 +918,20 @@ func (c *connectionImpl) GetObjectsDbSchemas(ctx context.Context, depth adbc.Obj result = make(map[string][]string) var header, trailer metadata.MD // Pre-populate the map of which schemas are in which catalogs - info, err := c.cl.GetDBSchemas(ctx, &flightsql.GetDBSchemasOpts{DbSchemaFilterPattern: dbSchema}, grpc.Header(&header), grpc.Trailer(&trailer), c.timeouts) + var info *flight.FlightInfo + info, err = c.cl.GetDBSchemas(ctx, &flightsql.GetDBSchemasOpts{DbSchemaFilterPattern: dbSchema}, grpc.Header(&header), grpc.Trailer(&trailer), c.timeouts) if err != nil { - return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetDBSchemas)") + err = adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetDBSchemas)") + return nil, err } header = metadata.MD{} trailer = metadata.MD{} - rdr, err := c.readInfo(ctx, schema_ref.DBSchemas, info, c.timeouts, grpc.Header(&header), grpc.Trailer(&trailer)) + var rdr array.RecordReader + rdr, err = c.readInfo(ctx, schema_ref.DBSchemas, info, c.timeouts, grpc.Header(&header), grpc.Trailer(&trailer)) if err != nil { - return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetDBSchemas)") + err = adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetDBSchemas)") + return nil, err } defer rdr.Release() @@ -864,12 +951,18 @@ func (c *connectionImpl) GetObjectsDbSchemas(ctx context.Context, depth adbc.Obj } if err := checkContext(rdr.Err(), ctx); err != nil { - return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetCatalogs)") + err = adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetCatalogs)") + return nil, err } return } 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) + defer internal.EndSpanWithStartTime(span, &err, &startTime) + if depth == adbc.ObjectDepthCatalogs || depth == adbc.ObjectDepthDBSchemas { return } @@ -879,14 +972,16 @@ func (c *connectionImpl) GetObjectsTables(ctx context.Context, depth adbc.Object // Pre-populate the map of which schemas are in which catalogs includeSchema := depth == adbc.ObjectDepthAll || depth == adbc.ObjectDepthColumns var header, trailer metadata.MD - info, err := c.cl.GetTables(ctx, &flightsql.GetTablesOpts{ + var info *flight.FlightInfo + info, err = c.cl.GetTables(ctx, &flightsql.GetTablesOpts{ DbSchemaFilterPattern: dbSchema, TableNameFilterPattern: tableName, TableTypes: tableType, IncludeSchema: includeSchema, }, grpc.Header(&header), grpc.Trailer(&trailer), c.timeouts) if err != nil { - return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetTables)") + err = adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetTables)") + return nil, err } expectedSchema := schema_ref.Tables @@ -895,9 +990,11 @@ func (c *connectionImpl) GetObjectsTables(ctx context.Context, depth adbc.Object } header = metadata.MD{} trailer = metadata.MD{} - rdr, err := c.readInfo(ctx, expectedSchema, info, c.timeouts, grpc.Header(&header), grpc.Trailer(&trailer)) + var rdr array.RecordReader + rdr, err = c.readInfo(ctx, expectedSchema, info, c.timeouts, grpc.Header(&header), grpc.Trailer(&trailer)) if err != nil { - return nil, adbcFromFlightStatus(err, "GetObjects(GetTables)") + err = adbcFromFlightStatus(err, "GetObjects(GetTables)") + return nil, err } defer rdr.Release() @@ -925,7 +1022,8 @@ func (c *connectionImpl) GetObjectsTables(ctx context.Context, depth adbc.Object var schema *arrow.Schema if includeSchema { - reader, err := ipc.NewReader(bytes.NewReader(rdr.RecordBatch().Column(4).(*array.Binary).Value(i))) + var reader *ipc.Reader + reader, err = ipc.NewReader(bytes.NewReader(rdr.RecordBatch().Column(4).(*array.Binary).Value(i))) if err != nil { return nil, adbc.Error{ Msg: err.Error(), @@ -944,13 +1042,19 @@ func (c *connectionImpl) GetObjectsTables(ctx context.Context, depth adbc.Object } } - if err := checkContext(rdr.Err(), ctx); err != nil { - return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetTables)") + if err = checkContext(rdr.Err(), ctx); err != nil { + err = adbcFromFlightStatusWithDetails(err, header, trailer, "GetObjects(GetTables)") + return nil, err } return } -func (c *connectionImpl) GetTableSchema(ctx context.Context, catalog *string, dbSchema *string, tableName string) (*arrow.Schema, error) { +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) + defer internal.EndSpanWithStartTime(span, &err, &startTime) + opts := &flightsql.GetTablesOpts{ Catalog: catalog, DbSchemaFilterPattern: dbSchema, @@ -960,41 +1064,48 @@ func (c *connectionImpl) GetTableSchema(ctx context.Context, catalog *string, db ctx = metadata.NewOutgoingContext(ctx, c.hdrs) var header, trailer metadata.MD - info, err := c.cl.GetTables(ctx, opts, c.timeouts, grpc.Header(&header), grpc.Trailer(&trailer)) + var info *flight.FlightInfo + info, err = c.cl.GetTables(ctx, opts, c.timeouts, grpc.Header(&header), grpc.Trailer(&trailer)) if err != nil { return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "GetTableSchema(GetTables)") } - header = metadata.MD{} - trailer = metadata.MD{} - rdr, err := doGetWithLogger(ctx, c.cl, info.Endpoint[0], c.clientCache, c.Logger, c.timeouts, grpc.Header(&header), grpc.Trailer(&trailer)) + ctx, responseMetadata := withResponseMetadata(ctx) + var rdr *flight.Reader + rdr, err = doGetWithTracer(ctx, c.cl, info.Endpoint[0], c.clientCache, c, c.timeouts) if err != nil { - return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "GetTableSchema(DoGet)") + err = adbcFromFlightStatusWithDetails(err, responseMetadata.snapshot(), nil, "GetTableSchema(DoGet)") + return nil, err } defer rdr.Release() - rec, err := rdr.Read() + var rec arrow.RecordBatch + rec, err = rdr.Read() if err != nil { if err == io.EOF { - return nil, adbc.Error{ + err = adbc.Error{ Msg: "No table found", Code: adbc.StatusNotFound, } + return nil, err } - return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "GetTableSchema(DoGet)") + err = adbcFromFlightStatusWithDetails(err, responseMetadata.snapshot(), nil, "GetTableSchema(DoGet)") + return nil, err } numRows := rec.NumRows() switch { case numRows == 0: - return nil, adbc.Error{ + err = adbc.Error{ Code: adbc.StatusNotFound, } + return nil, err case numRows > math.MaxInt32: - return nil, adbc.Error{ + err = adbc.Error{ Msg: "[Flight SQL] GetTableSchema cannot handle tables with number of rows > 2^31 - 1", Code: adbc.StatusNotImplemented, } + return nil, err } var s *arrow.Schema @@ -1010,16 +1121,18 @@ func (c *connectionImpl) GetTableSchema(ctx context.Context, catalog *string, db schemaBytes := rec.Column(4).(*array.Binary).Value(i) s, err = flight.DeserializeSchema(schemaBytes, c.db.Alloc) if err != nil { - return nil, adbcFromFlightStatus(err, "GetTableSchema") + err = adbcFromFlightStatus(err, "GetTableSchema") + return nil, err } return s, nil } } - return s, adbc.Error{ + err = adbc.Error{ Msg: "[Flight SQL] GetTableSchema could not find a table with a matching schema", Code: adbc.StatusNotFound, } + return s, err } // GetTableTypes returns a list of the table types in the database. @@ -1029,22 +1142,30 @@ func (c *connectionImpl) GetTableSchema(ctx context.Context, catalog *string, db // Field Name | Field Type // ----------------|-------------- // table_type | utf8 not null -func (c *connectionImpl) GetTableTypes(ctx context.Context) (array.RecordReader, error) { +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) + defer internal.EndSpanWithStartTime(span, &err, &startTime) + ctx = metadata.NewOutgoingContext(ctx, c.hdrs) var header, trailer metadata.MD - info, err := c.cl.GetTableTypes(ctx, c.timeouts, grpc.Header(&header), grpc.Trailer(&trailer)) + var info *flight.FlightInfo + info, err = c.cl.GetTableTypes(ctx, c.timeouts, grpc.Header(&header), grpc.Trailer(&trailer)) if err != nil { - return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "GetTableTypes") + err = adbcFromFlightStatusWithDetails(err, header, trailer, "GetTableTypes") + return nil, err } - return newRecordReader(ctx, recordReaderConfig{ + reader, err = newRecordReader(ctx, recordReaderConfig{ alloc: c.db.Alloc, cl: c.cl, info: info, clientCache: c.clientCache, bufferSize: 5, - logger: c.Logger, + tracing: c, }) + return reader, err } // Commit commits any pending transactions on this connection, it should @@ -1195,35 +1316,37 @@ func (c *connectionImpl) prepareSubstrait(ctx context.Context, plan flightsql.Su } // Close closes this connection and releases any associated resources. -func (c *connectionImpl) Close() error { +func (c *connectionImpl) Close() (err error) { + const spanName = "FlightSQL.Connection.Close" + startTime := time.Now() + ctx, span := internal.StartSpan(context.Background(), spanName, c) + defer internal.EndSpanWithStartTime(span, &err, &startTime) + if c.cl == nil { - return adbc.Error{ + err = adbc.Error{ Msg: "[Flight SQL Connection] trying to close already closed connection", Code: adbc.StatusInvalidState, } + return err } - closeStart := time.Now() // Snapshot fields before tearing down c.cl; log "closing" and // "closed" separately so a hung CloseSession is still visible. - logger := safeLogger(c.Logger) connID := c.id openedAt := c.openedAt + span.AddEvent("closing", trace.WithAttributes(attribute.String("connection_id", connID))) - logger.Info("FlightSQL connection closing", - "connection_id", connID, - ) - - ctx := metadata.NewOutgoingContext(context.Background(), c.hdrs) + ctx = metadata.NewOutgoingContext(ctx, c.hdrs) var header, trailer metadata.MD - _, err := c.cl.CloseSession(ctx, &flight.CloseSessionRequest{}, grpc.Header(&header), grpc.Trailer(&trailer), c.timeouts) + _, err = c.cl.CloseSession(ctx, &flight.CloseSessionRequest{}, grpc.Header(&header), grpc.Trailer(&trailer), c.timeouts) if err != nil { grpcStatus := grpcstatus.Convert(err) // Ignore unimplemented if grpcStatus.Code() != grpccodes.Unimplemented { // Ignore the error since server may not support it and may not properly return UNIMPLEMENTED // TODO(https://github.com/apache/arrow-adbc/issues/1243): log a proper warning - c.db.Logger.Debug("failed to close session", "error", err.Error()) + // Note: this does not set the status to error. It just records the error as an event in the span. + span.RecordError(err) } } @@ -1231,22 +1354,19 @@ func (c *connectionImpl) Close() error { err = c.cl.Close() c.cl = nil - args := []any{ - "connection_id", connID, - "close_duration", time.Since(closeStart), + args := []attribute.KeyValue{ + attribute.String("connection_id", connID), } if !openedAt.IsZero() { - args = append(args, "lifetime", time.Since(openedAt)) + args = append(args, attribute.Float64("lifetime_s", time.Since(openedAt).Seconds())) } if err != nil { - args = append(args, "err", err) - args = append(args, grpcStatusAttrs(err)...) - logger.Info("FlightSQL connection closed with error", args...) - } else { - logger.Info("FlightSQL connection closed", args...) + args = append(args, grpcStatusKeyValues(err)...) } + span.AddEvent("closed", trace.WithAttributes(args...)) - return adbcFromFlightStatus(err, "Close") + err = adbcFromFlightStatus(err, spanName) + return err } // ReadPartition constructs a statement for a partition of a query. The @@ -1254,6 +1374,11 @@ func (c *connectionImpl) Close() 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) + defer internal.EndSpanWithStartTime(span, &err, &startTime) + var info flight.FlightInfo if err := proto.Unmarshal(serializedPartition, &info); err != nil { return nil, adbc.Error{ @@ -1271,7 +1396,7 @@ func (c *connectionImpl) ReadPartition(ctx context.Context, serializedPartition } ctx = metadata.NewOutgoingContext(ctx, c.hdrs) - rdr, err = doGetWithLogger(ctx, c.cl, info.Endpoint[0], c.clientCache, c.Logger, c.timeouts) + rdr, err = doGetWithTracer(ctx, c.cl, info.Endpoint[0], c.clientCache, c, c.timeouts) if err != nil { return nil, adbcFromFlightStatus(err, "ReadPartition(DoGet)") } diff --git a/go/adbc/driver/flightsql/flightsql_database.go b/go/adbc/driver/flightsql/flightsql_database.go index 28e73399d3..a2f3e887ca 100644 --- a/go/adbc/driver/flightsql/flightsql_database.go +++ b/go/adbc/driver/flightsql/flightsql_database.go @@ -21,8 +21,8 @@ import ( "context" "crypto/tls" "crypto/x509" + "errors" "fmt" - "log/slog" "net/url" "strconv" "strings" @@ -36,6 +36,8 @@ import ( "github.com/apache/arrow-go/v18/arrow/flight" "github.com/apache/arrow-go/v18/arrow/flight/flightsql" "github.com/bluele/gcache" + "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" + "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" "google.golang.org/grpc" "google.golang.org/grpc/credentials" @@ -369,35 +371,49 @@ func (d *databaseImpl) SetOptionDouble(key string, value float64) error { return d.DatabaseImplBase.SetOptionDouble(key, value) } -func (d *databaseImpl) Close() error { - if d.Logger != nil { - d.Logger.Info("FlightSQL database closed", - "target", d.uri.String(), - ) - } - return d.DatabaseImplBase.Close() +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.AddEvent("closing", trace.WithAttributes(attribute.String("target", d.uri.String()))) + return closeTracing(context.Background(), &d.DatabaseImplBase, func(flushErr error) { + internal.EndSpanWithStartTime(span, &flushErr, &startTime) + }) +} + +type tracingLifecycle interface { + ForceFlushTracing(context.Context) error + Close() error } -func getFlightClient(ctx context.Context, loc string, d *databaseImpl, authMiddle *bearerAuthMiddleware, cookies flight.CookieMiddleware) (*flightsql.Client, error) { +func closeTracing(ctx context.Context, lifecycle tracingLifecycle, finishSpan func(error)) error { + flushErr := lifecycle.ForceFlushTracing(ctx) + finishSpan(flushErr) + shutdownErr := lifecycle.Close() + return errors.Join(flushErr, shutdownErr) +} + +func getFlightClient(ctx context.Context, loc string, d *databaseImpl, authMiddle *bearerAuthMiddleware, cookies flight.CookieMiddleware, span trace.Span) (client *flightsql.Client, err error) { middleware := []flight.ClientMiddleware{ - { - Unary: makeUnaryLoggingInterceptor(d.Logger), - Stream: makeStreamLoggingInterceptor(d.Logger), - }, flight.CreateClientMiddleware(authMiddle), { Unary: unaryTimeoutInterceptor, Stream: streamTimeoutInterceptor, }, + {Stream: responseMetadataStreamInterceptor}, } if d.enableCookies { middleware = append(middleware, flight.CreateClientMiddleware(cookies)) } - uri, err := url.Parse(loc) + var uri *url.URL + uri, err = url.Parse(loc) if err != nil { - return nil, adbc.Error{Msg: fmt.Sprintf("Invalid URI '%s': %s", loc, err), Code: adbc.StatusInvalidArgument} + err = adbc.Error{Msg: fmt.Sprintf("Invalid URI '%s': %s", loc, err), Code: adbc.StatusInvalidArgument} + return nil, err } creds := d.creds @@ -436,104 +452,132 @@ func getFlightClient(ctx context.Context, loc string, d *databaseImpl, authMiddl dv, _ := d.DriverInfo.GetInfoForInfoCode(adbc.InfoDriverVersion) driverVersion := dv.(string) - dialOpts := append(d.dialOpts.opts, grpc.WithConnectParams(d.timeout.connectParams()), grpc.WithTransportCredentials(creds), grpc.WithUserAgent("ADBC Flight SQL Driver "+driverVersion)) + dialOpts := append(d.dialOpts.opts, + grpc.WithStatsHandler(otelgrpc.NewClientHandler( + otelgrpc.WithTracerProvider(d.GetTracerProvider()), + )), + grpc.WithConnectParams(d.timeout.connectParams()), + grpc.WithTransportCredentials(creds), + grpc.WithUserAgent("ADBC Flight SQL Driver "+driverVersion), + ) dialOpts = append(dialOpts, d.userDialOpts...) if d.oauthToken != nil { dialOpts = append(dialOpts, grpc.WithPerRPCCredentials(d.oauthToken)) } - d.Logger.DebugContext(ctx, "new client", "location", loc) - cl, err := flightsql.NewClient(target, nil, middleware, dialOpts...) + span.AddEvent("flight.client.connecting", trace.WithAttributes( + attribute.String("flight.location", loc), + )) + client, err = flightsql.NewClient(target, nil, middleware, dialOpts...) if err != nil { - return nil, adbc.Error{ + err = adbc.Error{ Msg: err.Error(), Code: adbc.StatusIO, } + return nil, err } - cl.Alloc = d.Alloc + client.Alloc = d.Alloc // Authorization header is already set, continue if len(authMiddle.hdrs.Get("authorization")) > 0 { - d.Logger.DebugContext(ctx, "reusing auth token", "location", loc) - return cl, nil + span.AddEvent("flight.auth.token.reused", trace.WithAttributes( + attribute.String("flight.location", loc), + )) + return client, nil } var authValue string if d.user != "" || d.pass != "" { authStart := time.Now() - d.Logger.InfoContext(ctx, "FlightSQL basic auth started", - "target", loc, - "user", d.user, - ) + span.AddEvent("flight.auth.basic.started", trace.WithAttributes( + attribute.String("target", loc), + attribute.String("user", d.user), + )) var header, trailer metadata.MD - ctx, err = cl.Client.AuthenticateBasicToken(ctx, d.user, d.pass, grpc.Header(&header), grpc.Trailer(&trailer), d.timeout) + ctx, err = client.Client.AuthenticateBasicToken(ctx, d.user, d.pass, grpc.Header(&header), grpc.Trailer(&trailer), d.timeout) if err != nil { - args := []any{ - "target", loc, - "user", d.user, - "duration", time.Since(authStart), - "err", err, + args := []attribute.KeyValue{ + attribute.String("target", loc), + attribute.String("user", d.user), + attribute.Float64("duration_s", time.Since(authStart).Seconds()), } - args = append(args, correlationHeaderAttrs(header)...) - args = append(args, correlationHeaderAttrs(trailer)...) - args = append(args, grpcStatusAttrs(err)...) - d.Logger.InfoContext(ctx, "FlightSQL basic auth failed", args...) - return nil, adbcFromFlightStatusWithDetails(err, header, trailer, "AuthenticateBasicToken") + args = append(args, correlationHeaderKeyValues(header)...) + args = append(args, correlationHeaderKeyValues(trailer)...) + args = append(args, grpcStatusKeyValues(err)...) + span.SetAttributes(args...) + err = adbcFromFlightStatusWithDetails(err, header, trailer, "AuthenticateBasicToken") + return nil, err } if md, ok := metadata.FromOutgoingContext(ctx); ok { authValue = md.Get("Authorization")[0] } - d.Logger.InfoContext(ctx, "FlightSQL basic auth succeeded", - "target", loc, - "user", d.user, - "duration", time.Since(authStart), - "token_length", len(authValue), - ) + span.AddEvent("flight.auth.basic.completed", trace.WithAttributes( + attribute.String("target", loc), + attribute.String("user", d.user), + attribute.String("duration_s", time.Since(authStart).String()), + attribute.Int("token_length", len(authValue)), + )) } if authValue != "" { - authMiddle.SetHeader(authValue) + authMiddle.SetHeader(authValue, span) } - return cl, nil + return client, nil } type support struct { transactions bool } +func closeCachedFlightClient(d *databaseImpl, location, client interface{}, reason string) { + startTime := time.Now() + var err error + _, span := internal.StartSpan(context.Background(), "FlightSQL.Database.CloseCachedClient", d, + trace.WithAttributes( + attribute.String("flight.location", fmt.Sprint(location)), + attribute.String("flight.cache.reason", reason), + )) + defer internal.EndSpanWithStartTime(span, &err, &startTime) + + err = client.(*flightsql.Client).Close() +} + func (d *databaseImpl) Open(ctx context.Context) (_ adbc.Connection, err error) { - ctx, span := internal.StartSpan( - ctx, - "FlightSQLDatabase.Open", - d, - trace.WithAttributes(traceHeaderAttrsWithPrefix(d.hdrs, traceRequestMetadataPrefix)...), - ) - defer internal.EndSpanWithError(span, &err) + startTime := time.Now() + ctx, span := internal.StartSpan(ctx, "FlightSQL.Database.Open", d, trace.WithAttributes(traceHeaderAttrsWithPrefix(d.hdrs, traceRequestMetadataPrefix)...)) + defer internal.EndSpanWithStartTime(span, &err, &startTime) - authMiddle := &bearerAuthMiddleware{hdrs: d.hdrs.Copy(), logger: safeLogger(d.Logger)} + authMiddle := &bearerAuthMiddleware{ + hdrs: d.hdrs.Copy(), + } var cookies flight.CookieMiddleware if d.enableCookies { cookies = flight.NewCookieMiddleware() } - cl, err := getFlightClient(ctx, d.uri.String(), d, authMiddle, cookies) + cl, err := getFlightClient(ctx, d.uri.String(), d, authMiddle, cookies, span) if err != nil { return nil, err } cache := gcache.New(20).LRU(). Expiration(5 * time.Minute). - LoaderFunc(func(loc interface{}) (interface{}, error) { + LoaderFunc(func(loc interface{}) (_ interface{}, err error) { + startTime := time.Now() + ctx, cacheSpan := internal.StartSpan(context.Background(), "FlightSQL.Database.LoadCachedClient", d) + defer internal.EndSpanWithStartTime(cacheSpan, &err, &startTime) + uri, ok := loc.(string) if !ok { return nil, adbc.Error{Msg: fmt.Sprintf("Location must be a string, got %#v", uri), Code: adbc.StatusInternal} } + cacheSpan.SetAttributes(attribute.String("flight.location", uri)) var cookieMiddleware flight.CookieMiddleware // if cookies are enabled, start by cloning the existing cookies @@ -541,8 +585,13 @@ func (d *databaseImpl) Open(ctx context.Context) (_ adbc.Connection, err error) cookieMiddleware = cookies.Clone() } // use the existing auth token if there is one - cl, err := getFlightClient(context.Background(), uri, d, - &bearerAuthMiddleware{hdrs: authMiddle.hdrs.Copy(), logger: safeLogger(d.Logger)}, cookieMiddleware) + cl, err := getFlightClient( + ctx, + uri, + d, + &bearerAuthMiddleware{hdrs: authMiddle.hdrs.Copy()}, + cookieMiddleware, + cacheSpan) if err != nil { return nil, err } @@ -550,18 +599,10 @@ func (d *databaseImpl) Open(ctx context.Context) (_ adbc.Connection, err error) cl.Alloc = d.Alloc return cl, nil }). - EvictedFunc(func(_, client interface{}) { - conn := client.(*flightsql.Client) - err := conn.Close() - if err != nil { - d.Logger.Debug("failed to close client", "error", err.Error()) - } - }).PurgeVisitorFunc(func(_ interface{}, client interface{}) { - conn := client.(*flightsql.Client) - err := conn.Close() - if err != nil { - d.Logger.Debug("failed to close client", "error", err.Error()) - } + EvictedFunc(func(location, client interface{}) { + closeCachedFlightClient(d, location, client, "evicted") + }).PurgeVisitorFunc(func(location, client interface{}) { + closeCachedFlightClient(d, location, client, "purged") }).Build() var cnxnSupport support @@ -572,7 +613,7 @@ func (d *databaseImpl) Open(ctx context.Context) (_ adbc.Connection, err error) const int32code = 3 for _, endpoint := range info.Endpoint { - rdr, err := doGetWithLogger(ctx, cl, endpoint, cache, d.Logger, d.timeout) + rdr, err := doGetWithTracer(ctx, cl, endpoint, cache, d, d.timeout) if err != nil { continue } @@ -615,12 +656,11 @@ func (d *databaseImpl) Open(ctx context.Context) (_ adbc.Connection, err error) // this connection (and any statements derived from it). conn.id = newRandomID("conn") conn.openedAt = time.Now() - conn.Logger = safeLogger(conn.Logger).With("connection_id", conn.id) - conn.Logger.InfoContext(ctx, "FlightSQL connection opened", - "target", d.uri.String(), - "transactionsSupported", cnxnSupport.transactions, - "driver", infoDriverName, - ) + span.AddEvent("finished", trace.WithAttributes( + attribute.String("target", d.uri.String()), + attribute.Bool("transactionsSupported", cnxnSupport.transactions), + attribute.String("driver", infoDriverName), + )) return driverbase.NewConnectionBuilder(conn). WithDriverInfoPreparer(conn). @@ -632,9 +672,6 @@ func (d *databaseImpl) Open(ctx context.Context) (_ adbc.Connection, err error) type bearerAuthMiddleware struct { mutex sync.RWMutex hdrs metadata.MD - // logger, when non-nil, receives an Info event each time the bearer - // token is rotated. Only token lengths are logged, never values. - logger *slog.Logger } func (b *bearerAuthMiddleware) StartCall(ctx context.Context) context.Context { @@ -645,50 +682,46 @@ func (b *bearerAuthMiddleware) StartCall(ctx context.Context) context.Context { } // rotateAuth atomically replaces the stored Authorization metadata and -// returns the previous value plus the current logger. Callers invoke -// the logger outside the critical section. -func (b *bearerAuthMiddleware) rotateAuth(headers ...string) (previous []string, logger *slog.Logger) { +// returns the previous value. +func (b *bearerAuthMiddleware) rotateAuth(headers ...string) (previous []string) { b.mutex.Lock() defer b.mutex.Unlock() previous = b.hdrs.Get("authorization") b.hdrs.Set("authorization", headers...) - return previous, b.logger + return previous } func (b *bearerAuthMiddleware) HeadersReceived(ctx context.Context, md metadata.MD) { + captureResponseMetadata(ctx, md) // apache/arrow-adbc#584 headers := md.Get("authorization") if len(headers) == 0 { return } - previous, logger := b.rotateAuth(headers...) - if logger == nil { - return - } + previous := b.rotateAuth(headers...) // Log lengths, never values, so credentials never reach the log path. var prevLen int if len(previous) > 0 { prevLen = len(previous[0]) } - logger.InfoContext(ctx, "FlightSQL bearer token rotated by server", - "previous_token_length", prevLen, - "new_token_length", len(headers[0]), - "source", "HeadersReceived", - ) + if span := trace.SpanFromContext(ctx); span != nil && span.IsRecording() { + span.AddEvent("auth_token.rotated_by_server", trace.WithAttributes( + attribute.Int("previous_token_length", prevLen), + attribute.Int("new_token_length", len(headers[0])), + attribute.String("source", "HeadersReceived"), + )) + } } -func (b *bearerAuthMiddleware) SetHeader(authValue string) { - previous, logger := b.rotateAuth(authValue) - if logger == nil { - return - } +func (b *bearerAuthMiddleware) SetHeader(authValue string, span trace.Span) { + previous := b.rotateAuth(authValue) var prevLen int if len(previous) > 0 { prevLen = len(previous[0]) } - logger.Info("FlightSQL bearer token rotated by client", - "previous_token_length", prevLen, - "new_token_length", len(authValue), - "source", "SetHeader", - ) + span.AddEvent("auth_token.rotated_by_client", trace.WithAttributes( + attribute.Int("previous_token_length", prevLen), + attribute.Int("new_token_length", len(authValue)), + attribute.String("source", "SetHeader"), + )) } diff --git a/go/adbc/driver/flightsql/flightsql_driver.go b/go/adbc/driver/flightsql/flightsql_driver.go index 169b60a084..090d7c1661 100644 --- a/go/adbc/driver/flightsql/flightsql_driver.go +++ b/go/adbc/driver/flightsql/flightsql_driver.go @@ -40,6 +40,7 @@ package flightsql import ( "context" + "errors" "net/url" "time" @@ -127,7 +128,7 @@ func (d *driverImpl) NewDatabaseWithOptions(opts map[string]string, userDialOpts return d.NewDatabaseWithOptionsContext(context.Background(), opts, userDialOpts...) } -func (d *driverImpl) NewDatabaseWithOptionsContext(ctx context.Context, opts map[string]string, userDialOpts ...grpc.DialOption) (adbc.Database, error) { +func (d *driverImpl) NewDatabaseWithOptionsContext(ctx context.Context, opts map[string]string, userDialOpts ...grpc.DialOption) (_ adbc.Database, err error) { opts = maps.Clone(opts) uri, ok := opts[adbc.OptionKeyURI] if !ok { @@ -151,6 +152,14 @@ func (d *driverImpl) NewDatabaseWithOptionsContext(ctx context.Context, opts map if err != nil { return nil, err } + constructionComplete := false + defer func() { + if !constructionComplete { + if closeErr := dbBase.Close(); closeErr != nil { + err = errors.Join(err, closeErr) + } + } + }() db := &databaseImpl{ DatabaseImplBase: dbBase, timeout: timeoutOption{ @@ -174,6 +183,7 @@ func (d *driverImpl) NewDatabaseWithOptionsContext(ctx context.Context, opts map return nil, err } + constructionComplete = true return driverbase.NewDatabase(db), nil } diff --git a/go/adbc/driver/flightsql/flightsql_tracing.go b/go/adbc/driver/flightsql/flightsql_tracing.go new file mode 100644 index 0000000000..0d63196e73 --- /dev/null +++ b/go/adbc/driver/flightsql/flightsql_tracing.go @@ -0,0 +1,220 @@ +// 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" + "google.golang.org/grpc/status" +) + +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", "")) + } 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 +} + +// headerKeyValuesWithPrefix is the shared implementation behind +// correlationHeaderAttrs (incoming) and outgoingCallHeaderAttrs +// (outbound). Only headers in wellKnownCorrelationHeaders are emitted; +// returns nil when none are present. +func headerKeyValuesWithPrefix(md metadata.MD, prefix string) []attribute.KeyValue { + if len(md) == 0 { + return nil + } + out := make([]attribute.KeyValue, 0, 4) + for _, k := range wellKnownCorrelationHeaders { + if vals := md.Get(k); len(vals) > 0 { + out = append(out, attribute.StringSlice(prefix+k, vals)) + } + } + return out +} + +// correlationHeaderKeyValues returns OpenTelemetry attributes for well-known +// correlation headers present in md (typically incoming headers/trailers). Uses the +// "hdr_" prefix; only allow-listed headers are emitted. +func correlationHeaderKeyValues(md metadata.MD) []attribute.KeyValue { + return headerKeyValuesWithPrefix(md, "hdr_") +} + +// grpcStatusKeyValues returns OpenTelemetry attributes for the gRPC status +// embedded in err, or nil if err has no status. +func grpcStatusKeyValues(err error) []attribute.KeyValue { + if err == nil { + return nil + } + st, ok := status.FromError(err) + if !ok { + return nil + } + return []attribute.KeyValue{ + attribute.String("grpc_code", st.Code().String()), + attribute.String("grpc_message", st.Message()), + } +} + +// 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 9dae400b8d..e60627c33d 100644 --- a/go/adbc/driver/flightsql/logging.go +++ b/go/adbc/driver/flightsql/logging.go @@ -29,9 +29,6 @@ import ( "github.com/apache/arrow-go/v18/arrow/flight" "go.opentelemetry.io/otel/trace" - "golang.org/x/exp/maps" - "golang.org/x/exp/slices" - "google.golang.org/grpc" "google.golang.org/grpc/metadata" "google.golang.org/grpc/status" ) @@ -51,35 +48,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", "")) - } 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 +76,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 { @@ -143,118 +92,6 @@ func formatInt(n int64) string { return strconv.FormatInt(n, 10) } -func makeUnaryLoggingInterceptor(logger *slog.Logger) grpc.UnaryClientInterceptor { - interceptor := func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error { - start := time.Now() - // Ignore errors - outgoing, _ := metadata.FromOutgoingContext(ctx) - err := invoker(ctx, method, req, reply, cc, opts...) - if logger.Enabled(ctx, slog.LevelDebug) { - args := []any{"target", cc.Target(), "duration", time.Since(start), "err", err, "metadata", outgoing} - args = append(args, outgoingCallHeaderAttrs(ctx)...) - args = append(args, grpcStatusAttrs(err)...) - logger.DebugContext(ctx, method, args...) - } else { - keys := maps.Keys(outgoing) - slices.Sort(keys) - args := []any{"target", cc.Target(), "duration", time.Since(start), "err", err, "metadata", keys} - // Surface curated outbound correlation IDs regardless of level. - args = append(args, outgoingCallHeaderAttrs(ctx)...) - args = append(args, grpcStatusAttrs(err)...) - logger.InfoContext(ctx, method, args...) - } - return err - } - return interceptor -} - -func makeStreamLoggingInterceptor(logger *slog.Logger) grpc.StreamClientInterceptor { - interceptor := func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) { - start := time.Now() - // Ignore errors - outgoing, _ := metadata.FromOutgoingContext(ctx) - stream, err := streamer(ctx, desc, cc, method, opts...) - if err != nil { - args := []any{"target", cc.Target(), "duration", time.Since(start), "err", err} - args = append(args, outgoingCallHeaderAttrs(ctx)...) - args = append(args, grpcStatusAttrs(err)...) - logger.InfoContext(ctx, method, args...) - return stream, err - } - - return &loggedStream{ClientStream: stream, logger: logger, ctx: ctx, method: method, start: start, target: cc.Target(), outgoing: outgoing}, err - } - return interceptor -} - -type loggedStream struct { - grpc.ClientStream - - logger *slog.Logger - ctx context.Context - method string - start time.Time - target string - outgoing metadata.MD - - // recvCount tracks how many messages were received before the stream - // ended; logged on termination so EOFs on empty streams are distinguishable - // from mid-stream failures. - recvCount int64 -} - -func (stream *loggedStream) RecvMsg(m any) error { - err := stream.ClientStream.RecvMsg(m) - if err == nil { - stream.recvCount++ - return nil - } - - loggedErr := err - if loggedErr == io.EOF { - loggedErr = nil - } - - // Capture trailers from the terminated stream; they often carry - // server-side diagnostic information for failure triage. - trailer := stream.Trailer() - - if stream.logger.Enabled(stream.ctx, slog.LevelDebug) { - stream.logger.DebugContext(stream.ctx, stream.method, - "target", stream.target, - "duration", time.Since(stream.start), - "err", loggedErr, - "recvMessages", stream.recvCount, - "metadata", stream.outgoing, - "trailer", trailer, - ) - } else { - keys := maps.Keys(stream.outgoing) - slices.Sort(keys) - trailerKeys := maps.Keys(trailer) - slices.Sort(trailerKeys) - args := []any{ - "target", stream.target, - "duration", time.Since(stream.start), - "err", loggedErr, - "recvMessages", stream.recvCount, - "metadata", keys, - "trailer", trailerKeys, - } - // Promote curated correlation headers from the trailer. - args = append(args, correlationHeaderAttrs(trailer)...) - // Promote the outbound correlation IDs the caller supplied. - args = append(args, outgoingCallHeaderAttrs(stream.ctx)...) - // EOF is a clean close in Flight, so loggedErr was nil-ed above; - // only attach status attrs for real errors. - if loggedErr != nil { - args = append(args, grpcStatusAttrs(loggedErr)...) - } - stream.logger.InfoContext(stream.ctx, stream.method, args...) - } - return err -} - // wellKnownCorrelationHeaders is the curated allow-list of inbound gRPC // header/trailer keys that are surfaced verbatim into log records, for // cross-referencing client-side logs with server-side traces. Includes diff --git a/go/adbc/driver/flightsql/record_reader.go b/go/adbc/driver/flightsql/record_reader.go index 071cd1880b..ed9316e068 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 { + if errorRecorded { + internal.EndSpanWithStartTimeAndRecordedError(span, &err, &startTime) + return + } + internal.EndSpanWithStartTime(span, &err, &startTime) + } + }() + 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,59 @@ 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), + )) } + internal.EndSpanWithStartTimeAndRecordedError(span, &reader.err, &startTime) // 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 +330,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 ab7b5f1794..987f533a73 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 fbd88bc618..64d6d9f8a3 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,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) - defer internal.EndSpanWithError(span, &err) + defer internal.EndSpanWithStartTime(span, &err, &startTime) 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 991e7270eb..fb664f7199 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 d44578d40d..24999f1a18 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,52 @@ 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...) +} + +// EndSpanWithRecordedError ends a span whose error event has already been +// recorded by the caller, setting only the final status and error type. +func EndSpanWithRecordedError(span trace.Span, err *error, options ...trace.SpanEndOption) { + if err != nil { + setSpanStatus(span, *err) + } else { + setSpanStatus(span, nil) + } + span.End(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...) +} + +// Ends the given span. +// If startTime is not nil, then the duration of the span is recorded as an attribute. +// If err is not nil, then the +// error is recorded and the status is set appropriately. +// Otherwise, the status is set to Ok. +func EndSpanWithStartTime(span trace.Span, err *error, startTime *time.Time, options ...trace.SpanEndOption) { + if startTime != nil { + span.SetAttributes(attribute.Float64("span.duration_s", time.Since(*startTime).Seconds())) + } + EndSpanWithError(span, err, options...) +} + +// EndSpanWithStartTimeAndRecordedError records duration and ends a span whose +// error event has already been recorded by the caller. +func EndSpanWithStartTimeAndRecordedError(span trace.Span, err *error, startTime *time.Time, options ...trace.SpanEndOption) { + if startTime != nil { + span.SetAttributes(attribute.Float64("span.duration_s", time.Since(*startTime).Seconds())) + } + EndSpanWithRecordedError(span, err, options...) } diff --git a/go/adbc/go.mod b/go/adbc/go.mod index 0861f07285..1fa209eade 100644 --- a/go/adbc/go.mod +++ b/go/adbc/go.mod @@ -29,6 +29,7 @@ require ( github.com/golang/protobuf v1.5.4 github.com/google/uuid v1.6.0 github.com/stretchr/testify v1.11.1 + go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.69.0 go.opentelemetry.io/otel v1.44.0 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.44.0 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.44.0 diff --git a/go/adbc/go.sum b/go/adbc/go.sum index 5966e0c083..a9147b84cb 100644 --- a/go/adbc/go.sum +++ b/go/adbc/go.sum @@ -78,6 +78,8 @@ github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.69.0 h1:2yEATaop1/a1I4psnSLgWVPLWwCzkqWakgJy7xTDVy0= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.69.0/go.mod h1:D7J12YRapIekYyPWgGPlA/23pRmpSEZC5xJC/TTLI9U= go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU= go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc= go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.44.0 h1:4YsVu3B8+3qtWYYrsUYgn0OG78pN0rnNPRGX4SbokQI=