This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new c7c3b0395 feat(go/adbc/driver/flightsql): implement tracing for core
Flight SQL operations (#4486)
c7c3b0395 is described below
commit c7c3b03954127ac1bacf091d141b8099eedfda83
Author: Artur Rakhmatulin <[email protected]>
AuthorDate: Fri Jul 10 10:00:40 2026 +0100
feat(go/adbc/driver/flightsql): implement tracing for core Flight SQL
operations (#4486)
## Summary
This change wires OpenTelemetry tracing into the core Flight SQL
execution path and ensures configured trace exporters are flushed when
the database is closed.
Before this change, the Flight SQL driver accepted tracing configuration
but did not emit trace data for the main open/query path. The Flight SQL
database close path also did not delegate to the shared database close
implementation, so tracer shutdown and exporter flush did not run.
Closes #4487
---
go/adbc/driver/flightsql/flightsql_adbc_test.go | 56 ++++++++++++++++++++++++
go/adbc/driver/flightsql/flightsql_database.go | 9 +++-
go/adbc/driver/flightsql/flightsql_statement.go | 15 ++++++-
go/adbc/driver/internal/driverbase/connection.go | 3 +-
go/adbc/driver/internal/driverbase/database.go | 27 +++++++++++-
5 files changed, 104 insertions(+), 6 deletions(-)
diff --git a/go/adbc/driver/flightsql/flightsql_adbc_test.go
b/go/adbc/driver/flightsql/flightsql_adbc_test.go
index 93d8b2b68..92c00e5b4 100644
--- a/go/adbc/driver/flightsql/flightsql_adbc_test.go
+++ b/go/adbc/driver/flightsql/flightsql_adbc_test.go
@@ -58,6 +58,7 @@ import (
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/metadata"
"google.golang.org/protobuf/proto"
+ _ "modernc.org/sqlite"
)
type HeaderServerMiddleware struct {
@@ -316,6 +317,61 @@ func TestADBCFlightSQL(t *testing.T) {
suite.Run(t, &DomainSocketTests{db: db})
}
+func TestFlightSQLTracingProducesTraceFiles(t *testing.T) {
+ // TODO: Reuse example.CreateDB() here after apache/arrow-go#916
isolates
+ // the shared in-memory SQLite database per call.
+ db, err := sql.Open("sqlite",
fmt.Sprintf("file:%s-%d?mode=memory&cache=private", strings.ToLower(t.Name()),
time.Now().UnixNano()))
+ require.NoError(t, err)
+ defer validation.CheckedClose(t, db)
+
+ q := &FlightSQLQuirks{db: db}
+ drv := q.SetupDriver(t)
+ defer q.TearDownDriver(t, drv)
+
+ traceDir := t.TempDir()
+ opts := q.DatabaseOptions()
+ opts[adbc.OptionKeyTelemetryTracesExporter] =
string(adbc.TelemetryExporterAdbcFile)
+ opts[adbc.OptionKeyTelemetryTracesFolderPath] = traceDir
+
+ func() {
+ adbcDB, err := drv.NewDatabase(opts)
+ require.NoError(t, err)
+ defer validation.CheckedClose(t, adbcDB)
+
+ cnxn, err := adbcDB.Open(context.Background())
+ require.NoError(t, err)
+ defer validation.CheckedClose(t, cnxn)
+
+ stmt, err := cnxn.NewStatement()
+ require.NoError(t, err)
+ defer validation.CheckedClose(t, stmt)
+
+ require.NoError(t, stmt.SetSqlQuery("SELECT 1"))
+ reader, _, err := stmt.ExecuteQuery(context.Background())
+ require.NoError(t, err)
+ defer reader.Release()
+
+ for reader.Next() {
+ }
+ require.NoError(t, reader.Err())
+ }()
+
+ files, err := filepath.Glob(filepath.Join(traceDir, "*.jsonl"))
+ require.NoError(t, err)
+ require.NotEmpty(t, files)
+
+ var traceOutput strings.Builder
+ for _, file := range files {
+ data, err := os.ReadFile(file)
+ require.NoError(t, err)
+ traceOutput.Write(data)
+ }
+
+ output := traceOutput.String()
+ require.Contains(t, output, "FlightSQLDatabase.Open")
+ require.Contains(t, output, "FlightSQLStatement.ExecuteQuery")
+}
+
// Run the test suite, but validating that a header set on the database is
ALWAYS passed
type FlightSQLWithHeaderQuirks struct {
diff --git a/go/adbc/driver/flightsql/flightsql_database.go
b/go/adbc/driver/flightsql/flightsql_database.go
index 259e534ed..15b2f90c2 100644
--- a/go/adbc/driver/flightsql/flightsql_database.go
+++ b/go/adbc/driver/flightsql/flightsql_database.go
@@ -30,6 +30,7 @@ import (
"time"
"github.com/apache/arrow-adbc/go/adbc"
+ "github.com/apache/arrow-adbc/go/adbc/driver/internal"
"github.com/apache/arrow-adbc/go/adbc/driver/internal/driverbase"
"github.com/apache/arrow-go/v18/arrow/array"
"github.com/apache/arrow-go/v18/arrow/flight"
@@ -373,7 +374,7 @@ func (d *databaseImpl) Close() error {
"target", d.uri.String(),
)
}
- return nil
+ return d.DatabaseImplBase.Close()
}
func getFlightClient(ctx context.Context, loc string, d *databaseImpl,
authMiddle *bearerAuthMiddleware, cookies flight.CookieMiddleware)
(*flightsql.Client, error) {
@@ -480,7 +481,11 @@ type support struct {
transactions bool
}
-func (d *databaseImpl) Open(ctx context.Context) (adbc.Connection, error) {
+func (d *databaseImpl) Open(ctx context.Context) (_ adbc.Connection, err
error) {
+ ctx, span := internal.StartSpan(ctx, "FlightSQLDatabase.Open", d)
+ // TODO(apache/arrow-adbc#4494): replace with a shared telemetry helper.
+ defer func() { internal.EndSpan(span, err) }()
+
authMiddle := &bearerAuthMiddleware{hdrs: d.hdrs.Copy(), logger:
safeLogger(d.Logger)}
var cookies flight.CookieMiddleware
if d.enableCookies {
diff --git a/go/adbc/driver/flightsql/flightsql_statement.go
b/go/adbc/driver/flightsql/flightsql_statement.go
index 61705ae6f..14283380e 100644
--- a/go/adbc/driver/flightsql/flightsql_statement.go
+++ b/go/adbc/driver/flightsql/flightsql_statement.go
@@ -29,6 +29,7 @@ import (
"unsafe"
"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"
@@ -530,6 +531,10 @@ func (s *statement) ExecuteQuery(ctx context.Context) (rdr
array.RecordReader, n
}
}
+ ctx, span := internal.StartSpan(ctx, "FlightSQLStatement.ExecuteQuery",
s.cnxn)
+ // TODO(apache/arrow-adbc#4494): replace with a shared telemetry helper.
+ defer func() { internal.EndSpan(span, err) }()
+
// Handle bulk ingest
if s.targetTable != "" {
nrec, err = s.executeIngest(ctx)
@@ -602,6 +607,10 @@ func (s *statement) ExecuteUpdate(ctx context.Context) (n
int64, err error) {
}
}
+ ctx, span := internal.StartSpan(ctx,
"FlightSQLStatement.ExecuteUpdate", s.cnxn)
+ // TODO(apache/arrow-adbc#4494): replace with a shared telemetry helper.
+ defer func() { internal.EndSpan(span, err) }()
+
// Handle bulk ingest
if s.targetTable != "" {
return s.executeIngest(ctx)
@@ -647,7 +656,11 @@ func (s *statement) ExecuteUpdate(ctx context.Context) (n
int64, err error) {
// Prepare turns this statement into a prepared statement to be executed
// multiple times. This invalidates any prior result sets.
-func (s *statement) Prepare(ctx context.Context) error {
+func (s *statement) Prepare(ctx context.Context) (err error) {
+ ctx, span := internal.StartSpan(ctx, "FlightSQLStatement.Prepare",
s.cnxn)
+ // TODO(apache/arrow-adbc#4494): replace with a shared telemetry helper.
+ defer func() { internal.EndSpan(span, err) }()
+
startTime := time.Now()
s.log.InfoContext(ctx, "FlightSQL Prepare start", s.queryAttrs()...)
diff --git a/go/adbc/driver/internal/driverbase/connection.go
b/go/adbc/driver/internal/driverbase/connection.go
index 272504c10..d716edc5c 100644
--- a/go/adbc/driver/internal/driverbase/connection.go
+++ b/go/adbc/driver/internal/driverbase/connection.go
@@ -154,7 +154,8 @@ func (base *ConnectionImplBase) Rollback(context.Context)
error {
func (base *ConnectionImplBase) GetInfo(ctx context.Context, infoCodes
[]adbc.InfoCode) (reader array.RecordReader, err error) {
_, span := internal.StartSpan(ctx, "ConnectionImplBase.GetInfo", base)
- defer internal.EndSpan(span, err)
+ // TODO(apache/arrow-adbc#4494): replace with a shared telemetry helper.
+ defer func() { internal.EndSpan(span, err) }()
if len(infoCodes) == 0 {
infoCodes = base.DriverInfo.InfoSupportedCodes()
diff --git a/go/adbc/driver/internal/driverbase/database.go
b/go/adbc/driver/internal/driverbase/database.go
index 5cff37a84..991e7270e 100644
--- a/go/adbc/driver/internal/driverbase/database.go
+++ b/go/adbc/driver/internal/driverbase/database.go
@@ -20,6 +20,7 @@ package driverbase
import (
"context"
"errors"
+ "io"
"log/slog"
"os"
"strings"
@@ -423,7 +424,24 @@ func newOtlpTraceExporters(ctx context.Context)
([]sdktrace.SpanExporter, error)
return []sdktrace.SpanExporter{grpcExporter, httpExporter}, nil
}
-func newAdbcFileExporter(driverName, folderPath string)
(*stdouttrace.Exporter, error) {
+type closableSpanExporter struct {
+ exporter sdktrace.SpanExporter
+ closer io.Closer
+}
+
+func (e *closableSpanExporter) ExportSpans(ctx context.Context, spans
[]sdktrace.ReadOnlySpan) error {
+ return e.exporter.ExportSpans(ctx, spans)
+}
+
+func (e *closableSpanExporter) Shutdown(ctx context.Context) error {
+ err := e.exporter.Shutdown(ctx)
+ if e.closer != nil {
+ err = errors.Join(err, e.closer.Close())
+ }
+ return err
+}
+
+func newAdbcFileExporter(driverName, folderPath string)
(sdktrace.SpanExporter, error) {
fullyQualifiedDriverName := strings.ToLower(driverNamespace + "." +
driverName)
writerOpts :=
[]rotatingFileWriterOption{WithLogNamePrefix(fullyQualifiedDriverName)}
if strings.TrimSpace(folderPath) != "" {
@@ -433,7 +451,12 @@ func newAdbcFileExporter(driverName, folderPath string)
(*stdouttrace.Exporter,
if err != nil {
return nil, err
}
- return stdouttrace.New(stdouttrace.WithWriter(fileWriter))
+ exporter, err := stdouttrace.New(stdouttrace.WithWriter(fileWriter))
+ if err != nil {
+ _ = fileWriter.Close()
+ return nil, err
+ }
+ return &closableSpanExporter{exporter: exporter, closer: fileWriter},
nil
}
func newTracerProvider(exporters ...sdktrace.SpanExporter)
(*sdktrace.TracerProvider, error) {
