feat(go/adbc): refactor logging instrumentation into OTel tracing - part 3/3

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

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() {
+		if errorRecorded {
+			internal.EndSpanWithStartTimeAndRecordedError(span, &err, &startTime)
+			return
+		}
+		internal.EndSpanWithStartTime(span, &err, &startTime)
+	}()
+
 	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
 }