diff --git a/go/adbc/driver/flightsql/flightsql_adbc_test.go b/go/adbc/driver/flightsql/flightsql_adbc_test.go index 2be9496e9a..79d318b107 100644 --- a/go/adbc/driver/flightsql/flightsql_adbc_test.go +++ b/go/adbc/driver/flightsql/flightsql_adbc_test.go @@ -368,8 +368,8 @@ func TestFlightSQLTracingProducesTraceFiles(t *testing.T) { } output := traceOutput.String() - require.Contains(t, output, "FlightSQLDatabase.Open") - require.Contains(t, output, "FlightSQLStatement.ExecuteQuery") + require.Contains(t, output, "FlightSQL.Database.Open") + require.Contains(t, output, "FlightSQL.Statement.ExecuteQuery") } // Run the test suite, but validating that a header set on the database is ALWAYS passed diff --git a/go/adbc/driver/flightsql/flightsql_connection.go b/go/adbc/driver/flightsql/flightsql_connection.go index 858b9f40d9..6db86c1b3b 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" @@ -234,6 +233,8 @@ var adbcToFlightSQLInfo = map[adbc.InfoCode]flightsql.SqlInfo{ adbc.InfoVendorSubstraitMaxVersion: flightsql.SqlInfoFlightSqlServerSubstraitMaxVersion, } +// doGetWithTracer performs DoGet against an endpoint's locations, tracing each +// attempt and joining all per-location failures into the returned error. 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)) @@ -381,95 +382,6 @@ func doGetWithTracer(ctx context.Context, cl *flightsql.Client, endpoint *flight return nil, err } -// doGetWithLogger performs DoGet against an endpoint's locations, logging each -// attempt and joining all per-location failures into the returned error so the -// caller can see every location that was tried. logger may be nil. -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) - if len(endpoint.Location) == 0 { - log.DebugContext(ctx, "FlightSQL doGet", - "phase", "noLocations", - ) - start := time.Now() - rdr, err = cl.DoGet(ctx, endpoint.Ticket, opts...) - log.DebugContext(ctx, "FlightSQL doGet", - "phase", "defaultClientResult", - "duration", time.Since(start), - "err", err, - ) - return rdr, err - } - - var ( - cc interface{} - hasFallback bool - attemptErrors []string - ) - - for _, loc := range endpoint.Location { - if loc.Uri == flight.LocationReuseConnection { - hasFallback = true - continue - } - - start := time.Now() - 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, - ) - continue - } - - conn := cc.(*flightsql.Client) - rdr, err = conn.DoGet(ctx, endpoint.Ticket, opts...) - 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, - ) - continue - } - - log.DebugContext(ctx, "FlightSQL doGet succeeded", - "location", loc.Uri, - "duration", time.Since(start), - ) - return - } - - if hasFallback { - start := time.Now() - rdr, err = cl.DoGet(ctx, endpoint.Ticket, opts...) - 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) - } - log.DebugContext(ctx, "FlightSQL doGet succeeded via default client fallback", - "duration", time.Since(start), - ) - return rdr, nil - } - - if err != nil && len(attemptErrors) > 1 { - err = fmt.Errorf("all %d DoGet location(s) failed: %s; final: %w", - len(attemptErrors), strings.Join(attemptErrors, "; "), err) - } - - return nil, err -} - func (c *connectionImpl) getSessionOptions(ctx context.Context) (map[string]interface{}, error) { ctx = metadata.NewOutgoingContext(ctx, c.hdrs) var header, trailer metadata.MD @@ -818,7 +730,17 @@ 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 func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() + driverInfo := c.DriverInfo if len(infoCodes) == 0 { @@ -839,7 +761,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 { @@ -852,10 +775,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() { @@ -889,20 +814,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 } } @@ -918,7 +845,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") @@ -934,7 +861,16 @@ 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 func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() var ( header, trailer metadata.MD numCatalogs int64 @@ -952,13 +888,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++ { @@ -968,8 +906,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 @@ -977,6 +916,15 @@ 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 func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() if depth == adbc.ObjectDepthCatalogs { return } @@ -984,16 +932,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() @@ -1013,12 +965,23 @@ 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 func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() + if depth == adbc.ObjectDepthCatalogs || depth == adbc.ObjectDepthDBSchemas { return } @@ -1028,14 +991,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 @@ -1044,9 +1009,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() @@ -1074,7 +1041,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(), @@ -1093,13 +1061,24 @@ 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 func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() + opts := &flightsql.GetTablesOpts{ Catalog: catalog, DbSchemaFilterPattern: dbSchema, @@ -1109,41 +1088,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 @@ -1159,16 +1145,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. @@ -1178,22 +1166,35 @@ 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 func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() + 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 @@ -1344,7 +1345,17 @@ 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 func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() + if c.cl == nil { return adbc.Error{ Msg: "[Flight SQL Connection] trying to close already closed connection", @@ -1352,27 +1363,23 @@ func (c *connectionImpl) Close() error { } } - 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) } } @@ -1380,22 +1387,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 @@ -1403,6 +1407,16 @@ 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 func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() + var info flight.FlightInfo if err := proto.Unmarshal(serializedPartition, &info); err != nil { return nil, adbc.Error{ @@ -1420,7 +1434,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 0ead6422ce..0931de3b9c 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,33 +371,38 @@ 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()))) + flushErr := d.ForceFlushTracing(context.Background()) + internal.NewEndSpanHelper(span). + WithError(flushErr). + WithStartTime(startTime). + EndSpan() + shutdownErr := d.DatabaseImplBase.Close() + return errors.Join(flushErr, shutdownErr) } -func getFlightClient(ctx context.Context, loc string, d *databaseImpl, authMiddle *bearerAuthMiddleware, cookies flight.CookieMiddleware) (*flightsql.Client, error) { +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} } @@ -436,80 +443,110 @@ 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 func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() + + err = client.(*flightsql.Client).Close() +} + func (d *databaseImpl) Open(ctx context.Context) (_ adbc.Connection, err error) { ctx, span := internal.StartSpan( ctx, - "FlightSQLDatabase.Open", + "FlightSQL.Database.Open", d, trace.WithAttributes(traceHeaderAttrsWithPrefix(d.hdrs, traceRequestMetadataPrefix)...), ) @@ -519,25 +556,37 @@ func (d *databaseImpl) Open(ctx context.Context) (_ adbc.Connection, err error) EndSpan() }() - 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 func() { + internal.NewEndSpanHelper(cacheSpan). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() + 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 @@ -545,8 +594,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 } @@ -554,18 +608,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 @@ -576,7 +622,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 } @@ -619,12 +665,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). @@ -636,9 +681,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 { @@ -649,50 +691,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_statement.go b/go/adbc/driver/flightsql/flightsql_statement.go index 61911e9e33..57aca4f94b 100644 --- a/go/adbc/driver/flightsql/flightsql_statement.go +++ b/go/adbc/driver/flightsql/flightsql_statement.go @@ -534,7 +534,7 @@ func (s *statement) ExecuteQuery(ctx context.Context) (rdr array.RecordReader, n ctx, span := internal.StartSpan( ctx, - "FlightSQLStatement.ExecuteQuery", + "FlightSQL.Statement.ExecuteQuery", s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...), ) @@ -618,7 +618,7 @@ func (s *statement) ExecuteUpdate(ctx context.Context) (n int64, err error) { ctx, span := internal.StartSpan( ctx, - "FlightSQLStatement.ExecuteUpdate", + "FlightSQL.Statement.ExecuteUpdate", s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...), ) @@ -676,7 +676,7 @@ func (s *statement) ExecuteUpdate(ctx context.Context) (n int64, err error) { func (s *statement) Prepare(ctx context.Context) (err error) { ctx, span := internal.StartSpan( ctx, - "FlightSQLStatement.Prepare", + "FlightSQL.Statement.Prepare", s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...), ) diff --git a/go/adbc/driver/flightsql/flightsql_tracing.go b/go/adbc/driver/flightsql/flightsql_tracing.go index 168e993de8..853be6897a 100644 --- a/go/adbc/driver/flightsql/flightsql_tracing.go +++ b/go/adbc/driver/flightsql/flightsql_tracing.go @@ -28,6 +28,7 @@ import ( "go.opentelemetry.io/otel/attribute" "google.golang.org/grpc" "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" ) type responseMetadataKey struct{} @@ -135,6 +136,46 @@ func (p *streamProgress) logKeyValues() []attribute.KeyValue { 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 diff --git a/go/adbc/driver/flightsql/logging.go b/go/adbc/driver/flightsql/logging.go index 48a3427284..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" ) @@ -95,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/tracing_test.go b/go/adbc/driver/flightsql/tracing_test.go index 0bd5a74cf2..8dfca8c1c3 100644 --- a/go/adbc/driver/flightsql/tracing_test.go +++ b/go/adbc/driver/flightsql/tracing_test.go @@ -88,7 +88,7 @@ func TestTraceHeaderAttrsWithPrefix_AppliedToSpan(t *testing.T) { ctx, span := internal.StartSpan( context.Background(), - "FlightSQLStatement.ExecuteQuery", + "FlightSQL.Statement.ExecuteQuery", tracing, trace.WithAttributes(traceHeaderAttrsWithPrefix(metadata.New(map[string]string{ "x-request-id": "req-123", @@ -114,7 +114,7 @@ func TestTraceHeaderAttrsWithPrefix_AppliedToSpan(t *testing.T) { if _, ok := got["rpc.request.metadata.x-random-header"]; ok { t.Fatalf("x-random-header leaked into span attrs: %v", got) } - if v := got["db.operation.name"]; len(v) != 1 || v[0] != "FlightSQLStatement.ExecuteQuery" { + if v := got["db.operation.name"]; len(v) != 1 || v[0] != "FlightSQL.Statement.ExecuteQuery" { t.Fatalf("db.operation.name = %v, want [FlightSQLStatement.ExecuteQuery]", v) } } diff --git a/go/adbc/go.mod b/go/adbc/go.mod index 09bcc0c7e9..c54b320c6c 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.12.0 + go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.70.0 go.opentelemetry.io/otel v1.45.0 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.45.0 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.45.0 diff --git a/go/adbc/go.sum b/go/adbc/go.sum index fca6e93b5c..fdb5875b49 100644 --- a/go/adbc/go.sum +++ b/go/adbc/go.sum @@ -74,6 +74,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.70.0 h1:oECp5f+hN7nkwjU/8BxQ/q23bGPb8FIrD839owX222E= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.70.0/go.mod h1:DqEFwLumhzMBDQv9PcWbyoDxHI/4lAk6CM4nJBH39sc= go.opentelemetry.io/otel v1.45.0 h1:pdrWmLHofpubmArBv1LgFSv1Z0Ie/ppdZzu+kUN5EeU= go.opentelemetry.io/otel v1.45.0/go.mod h1:XZxIqPapzEYnhNSScF5DIqXhm/rYi0FzCe2XddAwZfQ= go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.45.0 h1:QRefszxJmfPdjXUUm3j6iDzY03mTPXMjqErFqQ67vUg=