diff --git a/go/adbc/driver/flightsql/flightsql_adbc_test.go b/go/adbc/driver/flightsql/flightsql_adbc_test.go index 2be9496e9a..5df3f5a353 100644 --- a/go/adbc/driver/flightsql/flightsql_adbc_test.go +++ b/go/adbc/driver/flightsql/flightsql_adbc_test.go @@ -368,8 +368,41 @@ 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.Database.Close") + require.Contains(t, output, "FlightSQL.Statement.ExecuteQuery") +} + +func TestFlightSQLTracingCleansUpAfterConstructionFailure(t *testing.T) { + alloc := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer alloc.AssertSize(t, 0) + drv := driver.NewDriver(alloc) + + for _, test := range []struct { + name string + uri string + extraOp map[string]string + }{ + {name: "invalid URI", uri: "grpc://%"}, + {name: "invalid option", uri: "grpc://localhost", extraOp: map[string]string{"unknown option": "value"}}, + } { + t.Run(test.name, func(t *testing.T) { + traceDir := t.TempDir() + opts := map[string]string{ + adbc.OptionKeyURI: test.uri, + adbc.OptionKeyTelemetryTracesExporter: string(adbc.TelemetryExporterAdbcFile), + adbc.OptionKeyTelemetryTracesFolderPath: traceDir, + } + for key, value := range test.extraOp { + opts[key] = value + } + + _, err := drv.NewDatabase(opts) + require.Error(t, err) + require.IsType(t, adbc.Error{}, err) + require.NoError(t, os.RemoveAll(traceDir)) + }) + } } // Run the test suite, but validating that a header set on the database is ALWAYS passed diff --git a/go/adbc/driver/flightsql/flightsql_bulk_ingest.go b/go/adbc/driver/flightsql/flightsql_bulk_ingest.go index fea5604d5f..55d484aab5 100644 --- a/go/adbc/driver/flightsql/flightsql_bulk_ingest.go +++ b/go/adbc/driver/flightsql/flightsql_bulk_ingest.go @@ -20,14 +20,16 @@ package flightsql import ( "context" "fmt" - "log/slog" "time" "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-adbc/go/adbc/driver/internal" "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" "github.com/apache/arrow-go/v18/arrow/flight/flightsql" pb "github.com/apache/arrow-go/v18/arrow/flight/gen/flight" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" "google.golang.org/grpc" "google.golang.org/grpc/metadata" ) @@ -105,15 +107,26 @@ func createRecordReaderFromBatch(batch arrow.RecordBatch) (array.RecordReader, e // executeIngest performs bulk ingestion using the FlightSQL client's ExecuteIngest method. // This is called from the statement when a target table has been set for bulk ingest. -func (s *statement) executeIngest(ctx context.Context) (int64, error) { +func (s *statement) executeIngest(ctx context.Context) (nRows int64, err error) { + var startTime = time.Now() + ctx, span := internal.StartSpan(ctx, "FlightSQL.BulkIngest.Execute", s.cnxn) + errorRecorded := false + defer func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + WithRecordedError(errorRecorded). + EndSpan() + }() + if s.streamBind == nil && s.bound == nil { - return -1, adbc.Error{ + err = adbc.Error{ Msg: "[Flight SQL Statement] must call Bind before bulk ingestion", Code: adbc.StatusInvalidState, } + return -1, err } - startTime := time.Now() catalogStr := "" if s.catalog != nil { catalogStr = *s.catalog @@ -122,16 +135,16 @@ func (s *statement) executeIngest(ctx context.Context) (int64, error) { if s.dbSchema != nil { dbSchemaStr = *s.dbSchema } - startAttrs := []any{ - slog.String("target_table", s.targetTable), - slog.String("mode", s.ingestMode), - slog.String("catalog", catalogStr), - slog.String("db_schema", dbSchemaStr), - slog.Bool("temporary", s.temporary), - slog.Bool("streamBind", s.streamBind != nil), - slog.Bool("recordBound", s.bound != nil), + startAttrs := []attribute.KeyValue{ + attribute.String("target_table", s.targetTable), + attribute.String("mode", s.ingestMode), + attribute.String("catalog", catalogStr), + attribute.String("db_schema", dbSchemaStr), + attribute.Bool("temporary", s.temporary), + attribute.Bool("streamBind", s.streamBind != nil), + attribute.Bool("recordBound", s.bound != nil), } - s.log.InfoContext(ctx, "FlightSQL ExecuteIngest start", startAttrs...) + span.AddEvent("flight.ingest.started", trace.WithAttributes(startAttrs...)) opts := ingestOptions{ targetTable: s.targetTable, @@ -145,16 +158,16 @@ func (s *statement) executeIngest(ctx context.Context) (int64, error) { // Get the record reader to ingest var rdr array.RecordReader - var err error if s.streamBind != nil { rdr = s.streamBind } else { rdr, err = createRecordReaderFromBatch(s.bound) if err != nil { - s.log.WarnContext(ctx, "FlightSQL ExecuteIngest finished with error", - slog.Duration("duration", time.Since(startTime)), - "err", err, - ) + span.RecordError(err, trace.WithAttributes( + attribute.String("flight.stage", "create_record_reader"), + attribute.Float64("duration_s", time.Since(startTime).Seconds()), + ), trace.WithStackTrace(true)) + errorRecorded = true return -1, err } } @@ -163,20 +176,20 @@ func (s *statement) executeIngest(ctx context.Context) (int64, error) { var header, trailer metadata.MD callOpts := append([]grpc.CallOption{}, grpc.Header(&header), grpc.Trailer(&trailer), s.timeouts) - nRows, err := s.cnxn.cl.ExecuteIngest(ctx, rdr, ingestOpts, callOpts...) - finishAttrs := []any{ - slog.Duration("duration", time.Since(startTime)), - slog.Int64("rowsIngested", nRows), + nRows, err = s.cnxn.cl.ExecuteIngest(ctx, rdr, ingestOpts, callOpts...) + finishAttrs := []attribute.KeyValue{ + attribute.Float64("duration_s", time.Since(startTime).Seconds()), + attribute.Int64("rowsIngested", nRows), } - finishAttrs = append(finishAttrs, correlationHeaderAttrs(header)...) - finishAttrs = append(finishAttrs, correlationHeaderAttrs(trailer)...) + finishAttrs = append(finishAttrs, correlationHeaderKeyValues(header)...) + finishAttrs = append(finishAttrs, correlationHeaderKeyValues(trailer)...) if err != nil { - wrapped := adbcFromFlightStatusWithDetails(err, header, trailer, "ExecuteIngest") - finishAttrs = append(finishAttrs, "err", wrapped) - s.log.WarnContext(ctx, "FlightSQL ExecuteIngest finished with error", finishAttrs...) - return -1, wrapped + err = adbcFromFlightStatusWithDetails(err, header, trailer, "FlightSQL.BulkIngest.Execute") + span.RecordError(err, trace.WithAttributes(finishAttrs...), trace.WithStackTrace(true)) + errorRecorded = true + return -1, err } - s.log.InfoContext(ctx, "FlightSQL ExecuteIngest finished", finishAttrs...) + span.AddEvent("flight.ingest.completed", trace.WithAttributes(finishAttrs...)) return nRows, nil } diff --git a/go/adbc/driver/flightsql/flightsql_connection.go b/go/adbc/driver/flightsql/flightsql_connection.go index 858b9f40d9..183c30ecec 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)) @@ -279,7 +280,9 @@ func doGetWithTracer(ctx context.Context, cl *flightsql.Client, endpoint *flight } if err != nil { attrs = append(attrs, attribute.String("flight.stage", "do_get")) - span.RecordError(err, trace.WithAttributes(attrs...), trace.WithStackTrace(true)) + if !isRecordReaderSiblingCancellation(ctx) { + span.RecordError(err, trace.WithAttributes(attrs...), trace.WithStackTrace(true)) + } errorRecorded = true } else { span.AddEvent("flight.location.selected", trace.WithAttributes(attrs...)) @@ -352,10 +355,12 @@ func doGetWithTracer(ctx context.Context, cl *flightsql.Client, endpoint *flight 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)) + if !isRecordReaderSiblingCancellation(ctx) { + 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 } @@ -371,100 +376,13 @@ func doGetWithTracer(ctx context.Context, cl *flightsql.Client, endpoint *flight len(attemptErrors), strings.Join(attemptErrors, "; "), err) } if err != nil { - span.RecordError(err, trace.WithAttributes( - attribute.String("flight.stage", "all_locations_failed"), - attribute.Int("flight.location.attempt_count", len(attemptErrors)), - ), trace.WithStackTrace(true)) - errorRecorded = true - } - - return nil, err -} - -// doGetWithLogger performs DoGet against an endpoint's locations, logging each -// attempt and joining all per-location failures into the returned error so the -// caller can see every location that was tried. logger may be nil. -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) + if !isRecordReaderSiblingCancellation(ctx) { + span.RecordError(err, trace.WithAttributes( + attribute.String("flight.stage", "all_locations_failed"), + attribute.Int("flight.location.attempt_count", len(attemptErrors)), + ), trace.WithStackTrace(true)) } - 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) + errorRecorded = true } return nil, err @@ -818,7 +736,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 +767,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 +781,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 +820,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 +851,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 +867,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 +894,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 +912,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 +922,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 +938,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 +971,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 +997,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 +1015,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 +1047,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 +1067,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 +1094,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 +1151,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 +1172,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 +1351,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 +1369,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 +1393,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 +1413,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 +1440,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..42c4593278 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,41 @@ 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, }, + { + Unary: responseMetadataUnaryInterceptor, + 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 +446,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 +559,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 +597,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 +611,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 +625,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 +668,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 +684,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 +694,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..497404f668 100644 --- a/go/adbc/driver/flightsql/flightsql_statement.go +++ b/go/adbc/driver/flightsql/flightsql_statement.go @@ -36,6 +36,7 @@ import ( "github.com/apache/arrow-go/v18/arrow/flight/flightsql" "github.com/apache/arrow-go/v18/arrow/memory" "github.com/bluele/gcache" + "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" "google.golang.org/grpc" "google.golang.org/grpc/metadata" @@ -483,35 +484,43 @@ func (s *statement) SetOptionDouble(key string, value float64) error { // The query can then be executed with any of the Execute methods. // For queries expected to be executed repeatedly, Prepare should be // called before execution. -func (s *statement) SetSqlQuery(query string) error { +func (s *statement) SetSqlQuery(query string) (err error) { + const spanName = "FlightSQL.Statement.SetSqlQuery" + startTime := time.Now() + _, span := internal.StartSpan(context.Background(), spanName, s.cnxn) + defer func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() + if s.prepared != nil { - if err := s.closePreparedStatement(); err != nil { + if err = s.closePreparedStatement(); err != nil { return err } s.prepared = nil } - if err := s.clearIncrementalQuery(); err != nil { + if err = s.clearIncrementalQuery(); err != nil { return err } s.targetTable = "" s.query.setSqlQuery(query) - if s.log != nil { - s.log.Debug("FlightSQL SetSqlQuery", s.queryAttrs()...) - } + span.AddEvent("finished", trace.WithAttributes(s.queryAttrs()...)) return nil } -func (s *statement) queryAttrs() []any { +func (s *statement) queryAttrs() []attribute.KeyValue { if s.query.sqlQuery != "" { - return queryFingerprintAttrs(s.query.sqlQuery) + return queryFingerprintKeyValues(s.query.sqlQuery) } if s.query.substraitPlan != nil { - return substraitFingerprintAttrs(s.query.substraitPlan, s.query.substraitVersion) + return substraitFingerprintKeyValues(s.query.substraitPlan, s.query.substraitVersion) } if s.targetTable != "" { - return []any{slog.String("query_type", "ingest"), slog.String("target_table", s.targetTable)} + return []attribute.KeyValue{attribute.String("query_type", "ingest"), attribute.String("target_table", s.targetTable)} } - return []any{slog.String("query_type", "none")} + return []attribute.KeyValue{attribute.String("query_type", "none")} } // ExecuteQuery executes the current query or prepared statement @@ -520,42 +529,40 @@ func (s *statement) queryAttrs() []any { // // This invalidates any prior result sets on this statement. func (s *statement) ExecuteQuery(ctx context.Context) (rdr array.RecordReader, nrec int64, err error) { - if err := s.clearIncrementalQuery(); err != nil { + const spanName = "FlightSQL.Statement.ExecuteQuery" + startTime := time.Now() + ctx, span := internal.StartSpan(ctx, spanName, s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...)) + defer func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() + + if err = s.clearIncrementalQuery(); err != nil { return nil, -1, err } // Reject staged binds if no ingest target was provided if s.targetTable == "" && s.prepared == nil && (s.bound != nil || s.streamBind != nil) { - return nil, -1, adbc.Error{ + err = adbc.Error{ Msg: "[Flight SQL Statement] must set IngestTargetTable before bulk ingestion", Code: adbc.StatusInvalidState, } + return nil, -1, err } - ctx, span := internal.StartSpan( - ctx, - "FlightSQLStatement.ExecuteQuery", - s.cnxn, - trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...), - ) - defer func() { - internal.NewEndSpanHelper(span). - WithError(err). - EndSpan() - }() - // Handle bulk ingest if s.targetTable != "" { nrec, err = s.executeIngest(ctx) return nil, nrec, err } - startTime := time.Now() - startAttrs := append([]any{ - slog.Bool("prepared", s.prepared != nil), - slog.Bool("hasTxn", s.cnxn.txn != nil), + startAttrs := append([]attribute.KeyValue{ + attribute.Bool("prepared", s.prepared != nil), + attribute.Bool("hasTxn", s.cnxn.txn != nil), }, s.queryAttrs()...) - s.log.InfoContext(ctx, "FlightSQL ExecuteQuery start", startAttrs...) + span.AddEvent("starting", trace.WithAttributes(startAttrs...)) ctx = metadata.NewOutgoingContext(ctx, s.hdrs) var info *flight.FlightInfo @@ -568,25 +575,21 @@ func (s *statement) ExecuteQuery(ctx context.Context) (rdr array.RecordReader, n } defer func() { - finishAttrs := []any{ - slog.Duration("duration", time.Since(startTime)), - slog.String("phase", "GetFlightInfo"), + finishAttrs := []attribute.KeyValue{ + attribute.Float64("duration_s", time.Since(startTime).Seconds()), + attribute.String("flight.stage", "get_flight_info"), } if info != nil { - finishAttrs = append(finishAttrs, flightInfoLogAttrs(info)...) - } - finishAttrs = append(finishAttrs, correlationHeaderAttrs(header)...) - finishAttrs = append(finishAttrs, correlationHeaderAttrs(trailer)...) - if err != nil { - finishAttrs = append(finishAttrs, "err", err) - s.log.WarnContext(ctx, "FlightSQL ExecuteQuery finished with error", finishAttrs...) - } else { - s.log.InfoContext(ctx, "FlightSQL ExecuteQuery finished", finishAttrs...) + finishAttrs = append(finishAttrs, flightInfoTracingKeyValues(info)...) } + finishAttrs = append(finishAttrs, correlationHeaderKeyValues(header)...) + finishAttrs = append(finishAttrs, correlationHeaderKeyValues(trailer)...) + span.AddEvent("finished", trace.WithAttributes(finishAttrs...)) }() if err != nil { - return nil, -1, adbcFromFlightStatusWithDetails(err, header, trailer, "ExecuteQuery") + err = adbcFromFlightStatusWithDetails(err, header, trailer, spanName) + return nil, -1, err } nrec = info.TotalRecords @@ -596,7 +599,7 @@ func (s *statement) ExecuteQuery(ctx context.Context) (rdr array.RecordReader, n info: info, clientCache: s.clientCache, bufferSize: s.queueSize, - logger: s.log, + tracing: s.cnxn, }, s.timeouts) return } @@ -604,21 +607,32 @@ func (s *statement) ExecuteQuery(ctx context.Context) (rdr array.RecordReader, n // ExecuteUpdate executes a statement that does not generate a result // set. It returns the number of rows affected if known, otherwise -1. func (s *statement) ExecuteUpdate(ctx context.Context) (n int64, err error) { - if err := s.clearIncrementalQuery(); err != nil { + const spanName = "FlightSQL.Statement.ExecuteUpdate" + startTime := time.Now() + ctx, span := internal.StartSpan(ctx, spanName, s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...)) + defer func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() + + if err = s.clearIncrementalQuery(); err != nil { return -1, err } // Reject staged binds if no ingest target was provided if s.targetTable == "" && s.prepared == nil && (s.bound != nil || s.streamBind != nil) { - return -1, adbc.Error{ + err = adbc.Error{ Msg: "[Flight SQL Statement] must set IngestTargetTable before bulk ingestion", Code: adbc.StatusInvalidState, } + return -1, err } - ctx, span := internal.StartSpan( + ctx, span = internal.StartSpan( ctx, - "FlightSQLStatement.ExecuteUpdate", + "FlightSQL.Statement.ExecuteUpdate", s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...), ) @@ -633,12 +647,11 @@ func (s *statement) ExecuteUpdate(ctx context.Context) (n int64, err error) { return s.executeIngest(ctx) } - startTime := time.Now() - startAttrs := append([]any{ - slog.Bool("prepared", s.prepared != nil), - slog.Bool("hasTxn", s.cnxn.txn != nil), + startAttrs := append([]attribute.KeyValue{ + attribute.Bool("prepared", s.prepared != nil), + attribute.Bool("hasTxn", s.cnxn.txn != nil), }, s.queryAttrs()...) - s.log.InfoContext(ctx, "FlightSQL ExecuteUpdate start", startAttrs...) + span.AddEvent("starting", trace.WithAttributes(startAttrs...)) ctx = metadata.NewOutgoingContext(ctx, s.hdrs) var header, trailer metadata.MD @@ -650,22 +663,17 @@ func (s *statement) ExecuteUpdate(ctx context.Context) (n int64, err error) { } defer func() { - finishAttrs := []any{ - slog.Duration("duration", time.Since(startTime)), - slog.Int64("rowsAffected", n), - } - finishAttrs = append(finishAttrs, correlationHeaderAttrs(header)...) - finishAttrs = append(finishAttrs, correlationHeaderAttrs(trailer)...) - if err != nil { - finishAttrs = append(finishAttrs, "err", err) - s.log.WarnContext(ctx, "FlightSQL ExecuteUpdate finished with error", finishAttrs...) - } else { - s.log.InfoContext(ctx, "FlightSQL ExecuteUpdate finished", finishAttrs...) + finishAttrs := []attribute.KeyValue{ + attribute.Float64("duration_s", time.Since(startTime).Seconds()), + attribute.Int64("rows_affected", n), } + finishAttrs = append(finishAttrs, correlationHeaderKeyValues(header)...) + finishAttrs = append(finishAttrs, correlationHeaderKeyValues(trailer)...) + span.AddEvent("finished", trace.WithAttributes(finishAttrs...)) }() if err != nil { - err = adbcFromFlightStatusWithDetails(err, header, trailer, "ExecuteQuery") + err = adbcFromFlightStatusWithDetails(err, header, trailer, spanName) } return @@ -676,7 +684,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)...), ) @@ -687,26 +695,32 @@ func (s *statement) Prepare(ctx context.Context) (err error) { }() startTime := time.Now() - s.log.InfoContext(ctx, "FlightSQL Prepare start", s.queryAttrs()...) + const spanName = "FlightSQL.Statement.Prepare" + ctx, span = internal.StartSpan(ctx, spanName, s.cnxn, trace.WithAttributes(traceHeaderAttrsWithPrefix(s.hdrs, traceRequestMetadataPrefix)...)) + defer func() { + internal.NewEndSpanHelper(span). + WithError(err). + WithStartTime(startTime). + EndSpan() + }() + + span.AddEvent("starting", trace.WithAttributes(s.queryAttrs()...)) ctx = metadata.NewOutgoingContext(ctx, s.hdrs) var header, trailer metadata.MD - prep, err := s.query.prepare(ctx, s.cnxn, grpc.Header(&header), grpc.Trailer(&trailer), s.timeouts) + var prep *flightsql.PreparedStatement + prep, err = s.query.prepare(ctx, s.cnxn, grpc.Header(&header), grpc.Trailer(&trailer), s.timeouts) defer func() { - finishAttrs := []any{slog.Duration("duration", time.Since(startTime))} - finishAttrs = append(finishAttrs, correlationHeaderAttrs(header)...) - finishAttrs = append(finishAttrs, correlationHeaderAttrs(trailer)...) - if err != nil { - finishAttrs = append(finishAttrs, "err", err) - s.log.WarnContext(ctx, "FlightSQL Prepare finished with error", finishAttrs...) - } else { - s.log.InfoContext(ctx, "FlightSQL Prepare finished", finishAttrs...) - } + finishAttrs := []attribute.KeyValue{attribute.Float64("duration_s", time.Since(startTime).Seconds())} + finishAttrs = append(finishAttrs, correlationHeaderKeyValues(header)...) + finishAttrs = append(finishAttrs, correlationHeaderKeyValues(trailer)...) + span.AddEvent("finished", trace.WithAttributes(finishAttrs...)) }() if err != nil { - return adbcFromFlightStatusWithDetails(err, header, trailer, "Prepare") + err = adbcFromFlightStatusWithDetails(err, header, trailer, spanName) + return err } s.prepared = prep return nil diff --git a/go/adbc/driver/flightsql/flightsql_tracing.go b/go/adbc/driver/flightsql/flightsql_tracing.go index 168e993de8..18cf0f504d 100644 --- a/go/adbc/driver/flightsql/flightsql_tracing.go +++ b/go/adbc/driver/flightsql/flightsql_tracing.go @@ -19,15 +19,20 @@ package flightsql import ( "context" + "crypto/sha256" "encoding/hex" "fmt" + "slices" "sync" "time" "github.com/apache/arrow-go/v18/arrow/flight" "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" + "golang.org/x/exp/maps" "google.golang.org/grpc" "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" ) type responseMetadataKey struct{} @@ -63,6 +68,41 @@ func (c *responseMetadataCollector) snapshot() metadata.MD { return c.value.Copy() } +func responseMetadataUnaryInterceptor(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) (err error) { + span := trace.SpanFromContext(ctx) + + // Ignore errors + outgoing, _ := metadata.FromOutgoingContext(ctx) + err = invoker(ctx, method, req, reply, cc, opts...) + + if span.IsRecording() { + keys := maps.Keys(outgoing) + slices.Sort(keys) + args := []attribute.KeyValue{ + attribute.String("target", cc.Target()), + attribute.StringSlice("metadata", keys), + } + // Surface curated outbound correlation IDs regardless of level. + args = append(args, outgoingCallHeaderKeyValues(ctx)...) + args = append(args, grpcStatusKeyValues(err)...) + span.AddEvent("Metadata.Unary.Interceptor."+method, trace.WithAttributes(args...)) + } + return err +} + +// outgoingCallHeaderAttrs returns slog attributes for well-known correlation +// headers on ctx's outbound gRPC metadata. Uses the "out_hdr_" prefix. +func outgoingCallHeaderKeyValues(ctx context.Context) []attribute.KeyValue { + if ctx == nil { + return nil + } + md, ok := metadata.FromOutgoingContext(ctx) + if !ok { + return nil + } + return headerKeyValuesWithPrefix(md, "out_hdr_") +} + 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 { @@ -135,6 +175,80 @@ 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()), + } +} + +// queryFingerprintKeyValues builds OpenTelemetry attributes identifying a SQL query +// without exposing it: length and a SHA-256 prefix. The query text itself +// is never recorded because it can embed end-user PII as literals. +func queryFingerprintKeyValues(query string) []attribute.KeyValue { + if query == "" { + return []attribute.KeyValue{attribute.String("query_type", "empty")} + } + h := sha256.Sum256([]byte(query)) + return []attribute.KeyValue{ + attribute.String("query_type", "sql"), + attribute.Int("query_length", len(query)), + attribute.String("query_sha256_prefix", hex.EncodeToString(h[:8])), + } +} + +// substraitFingerprintKeyValues builds OpenTelemetry attributes identifying a Substrait +// plan: length, SHA-256 prefix, and protocol version. Plan bytes are never +// recorded. +func substraitFingerprintKeyValues(plan []byte, version string) []attribute.KeyValue { + if len(plan) == 0 { + return []attribute.KeyValue{attribute.String("query_type", "substrait_empty")} + } + h := sha256.Sum256(plan) + attrs := []attribute.KeyValue{ + attribute.String("query_type", "substrait"), + attribute.Int("substrait_plan_bytes", len(plan)), + attribute.String("substrait_plan_sha256_prefix", hex.EncodeToString(h[:8])), + } + if version != "" { + attrs = append(attrs, attribute.String("substrait_version", version)) + } + return attrs +} + // flightInfoTracingKeyValues returns OpenTelemetry attributes describing a FlightInfo: // descriptor type and command prefix, AppMetadata prefix (some backends // embed a server-side query handle there), and advisory record/byte diff --git a/go/adbc/driver/flightsql/logging.go b/go/adbc/driver/flightsql/logging.go index 48a3427284..34221841f2 100644 --- a/go/adbc/driver/flightsql/logging.go +++ b/go/adbc/driver/flightsql/logging.go @@ -20,18 +20,13 @@ package flightsql import ( "context" "crypto/rand" - "crypto/sha256" "encoding/hex" "io" "log/slog" "strconv" "time" - "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 +90,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 @@ -345,80 +228,3 @@ func newRandomID(prefix string) string { } return prefix + "-" + hex.EncodeToString(b[:]) } - -// queryFingerprintAttrs builds slog attributes identifying a SQL query -// without exposing it: length and a SHA-256 prefix. The query text itself -// is never logged because it can embed end-user PII as literals. -func queryFingerprintAttrs(query string) []any { - if query == "" { - return []any{slog.String("query_type", "empty")} - } - h := sha256.Sum256([]byte(query)) - return []any{ - slog.String("query_type", "sql"), - slog.Int("query_length", len(query)), - slog.String("query_sha256_prefix", hex.EncodeToString(h[:8])), - } -} - -// substraitFingerprintAttrs builds slog attributes identifying a Substrait -// plan: length, SHA-256 prefix, and protocol version. Plan bytes are never -// logged. -func substraitFingerprintAttrs(plan []byte, version string) []any { - if len(plan) == 0 { - return []any{slog.String("query_type", "substrait_empty")} - } - h := sha256.Sum256(plan) - attrs := []any{ - slog.String("query_type", "substrait"), - slog.Int("substrait_plan_bytes", len(plan)), - slog.String("substrait_plan_sha256_prefix", hex.EncodeToString(h[:8])), - } - if version != "" { - attrs = append(attrs, slog.String("substrait_version", version)) - } - return attrs -} - -// flightInfoLogAttrs returns slog 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 flightInfoLogAttrs(info *flight.FlightInfo) []any { - if info == nil { - return nil - } - attrs := []any{ - slog.Int("numEndpoints", len(info.Endpoint)), - slog.Int64("totalRecords", info.TotalRecords), - slog.Int64("totalBytes", info.TotalBytes), - slog.Bool("haveSchemaInFlightInfo", len(info.Schema) > 0), - } - if desc := info.FlightDescriptor; desc != nil { - attrs = append(attrs, slog.String("descriptorType", desc.Type.String())) - if len(desc.Cmd) > 0 { - limit := len(desc.Cmd) - if limit > maxLoggedBlobBytes { - limit = maxLoggedBlobBytes - } - attrs = append(attrs, - slog.Int("descriptorCmdBytes", len(desc.Cmd)), - slog.String("descriptorCmdPrefixHex", hex.EncodeToString(desc.Cmd[:limit])), - ) - } - if len(desc.Path) > 0 { - attrs = append(attrs, slog.Any("descriptorPath", desc.Path)) - } - } - if len(info.AppMetadata) > 0 { - limit := len(info.AppMetadata) - if limit > maxLoggedBlobBytes { - limit = maxLoggedBlobBytes - } - attrs = append(attrs, - slog.Int("appMetadataBytes", len(info.AppMetadata)), - slog.String("appMetadataPrefixHex", hex.EncodeToString(info.AppMetadata[:limit])), - ) - } - return attrs -} diff --git a/go/adbc/driver/flightsql/record_reader.go b/go/adbc/driver/flightsql/record_reader.go index 2cfbeb9d9e..a2003d2e6c 100644 --- a/go/adbc/driver/flightsql/record_reader.go +++ b/go/adbc/driver/flightsql/record_reader.go @@ -21,7 +21,6 @@ import ( "context" "errors" "fmt" - "log/slog" "sync/atomic" "time" @@ -54,6 +53,13 @@ type reader struct { var errReaderReleased = errors.New("record reader released") +type recordReaderCallerContextKey struct{} + +func isRecordReaderSiblingCancellation(ctx context.Context) bool { + callerCtx, ok := ctx.Value(recordReaderCallerContextKey{}).(context.Context) + return ok && ctx.Err() == context.Canceled && callerCtx.Err() == nil +} + // recordReaderConfig bundles the dependencies that newRecordReader // needs to spin up its per-endpoint goroutines. type recordReaderConfig struct { @@ -63,7 +69,6 @@ type recordReaderConfig struct { clientCache gcache.Cache bufferSize int tracing adbc.OTelTracing - logger *slog.Logger } // newRecordReader kicks off a goroutine for each endpoint and returns a @@ -108,6 +113,7 @@ func newRecordReader(ctx context.Context, cfg recordReaderConfig, opts ...grpc.C callerCtx := ctx group, ctx := errgroup.WithContext(ctx) ctx, cancelFn := context.WithCancelCause(ctx) + ctx = context.WithValue(ctx, recordReaderCallerContextKey{}, callerCtx) goEndpoint := func(endpointFn func() error) { group.Go(func() error { err := endpointFn() @@ -233,6 +239,9 @@ func newRecordReader(ctx context.Context, cfg recordReaderConfig, opts ...grpc.C endpointCtx, responseMetadata := withResponseMetadata(ctx) rdr, err := doGetWithTracer(endpointCtx, cfg.cl, endpoint, cfg.clientCache, cfg.tracing, opts...) if err != nil { + if checkRecordReaderContext(err, ctx, callerCtx) == nil { + return nil + } span.RecordError(err, trace.WithAttributes( append( append([]attribute.KeyValue{}, epAttrs...), diff --git a/go/adbc/driver/flightsql/record_reader_test.go b/go/adbc/driver/flightsql/record_reader_test.go index 987f533a73..75713dac86 100644 --- a/go/adbc/driver/flightsql/record_reader_test.go +++ b/go/adbc/driver/flightsql/record_reader_test.go @@ -75,6 +75,10 @@ func (f *testFlightService) DoGet(request *flight.Ticket, stream flight.FlightSe f.failureCount-- return fmt.Errorf("Failed request") } + if request.Ticket[0] == 125 { + <-stream.Context().Done() + return stream.Context().Err() + } schema := orderingSchema() wr := flight.NewRecordWriter(stream, ipc.WithSchema(schema)) @@ -302,7 +306,7 @@ func (suite *RecordReaderTests) TestSiblingCancellationRecordsOneException() { Schema: flight.SerializeSchema(orderingSchema(), suite.alloc), Endpoint: []*flight.FlightEndpoint{ {Ticket: &flight.Ticket{Ticket: []byte{127}}}, - {Ticket: &flight.Ticket{Ticket: []byte{126}}}, + {Ticket: &flight.Ticket{Ticket: []byte{125}}}, }, }, clientCache: suite.clCache, diff --git a/go/adbc/driver/flightsql/tracing_test.go b/go/adbc/driver/flightsql/tracing_test.go index 0bd5a74cf2..be1771582a 100644 --- a/go/adbc/driver/flightsql/tracing_test.go +++ b/go/adbc/driver/flightsql/tracing_test.go @@ -26,9 +26,67 @@ import ( 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 TestResponseMetadataUnaryInterceptorDoesNotEndParentSpan(t *testing.T) { + recorder := tracetest.NewSpanRecorder() + tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder)) + t.Cleanup(func() { + if err := tp.Shutdown(context.Background()); err != nil { + t.Fatalf("TracerProvider shutdown failed: %v", err) + } + }) + + conn, err := grpc.NewClient( + "passthrough:///test", + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + if err != nil { + t.Fatalf("grpc.NewClient() failed: %v", err) + } + t.Cleanup(func() { + if err := conn.Close(); err != nil { + t.Fatalf("ClientConn.Close() failed: %v", err) + } + }) + + ctx, span := tp.Tracer("test.flightsql").Start(context.Background(), "operation") + err = responseMetadataUnaryInterceptor( + ctx, + "/test.Service/Method", + nil, + nil, + conn, + func(context.Context, string, any, any, *grpc.ClientConn, ...grpc.CallOption) error { + return nil + }, + ) + if err != nil { + t.Fatalf("responseMetadataUnaryInterceptor() failed: %v", err) + } + if !span.IsRecording() { + t.Fatal("responseMetadataUnaryInterceptor() ended the parent span") + } + if got := len(recorder.Ended()); got != 0 { + t.Fatalf("ended spans before parent Span.End() = %d, want 0", got) + } + + span.End() + spans := recorder.Ended() + if len(spans) != 1 { + t.Fatalf("ended spans len = %d, want 1", len(spans)) + } + for _, event := range spans[0].Events() { + if event.Name == "Metadata.Unary.Interceptor./test.Service/Method" { + return + } + } + t.Fatal("parent span does not contain the unary metadata event") +} + func TestTraceHeaderAttrsWithPrefix_AllowAndDeny(t *testing.T) { md := metadata.New(map[string]string{ "x-request-id": "req-1", @@ -88,7 +146,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 +172,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=