From 6a6ba723a94a61a0fd5368389024058956e325f5 Mon Sep 17 00:00:00 2001 From: Kyle Brandt Date: Fri, 9 May 2025 14:48:47 -0400 Subject: [PATCH] SQL Expressions: Add internal GMS traces (#104836) get trace from inside GMS --- pkg/expr/sql/db.go | 10 +++++++-- pkg/expr/sql/db_test.go | 45 +++++++++++++++++++++++++++++++++------ pkg/expr/sql/dummy_arm.go | 3 ++- pkg/expr/sql_command.go | 2 +- 4 files changed, 49 insertions(+), 11 deletions(-) diff --git a/pkg/expr/sql/db.go b/pkg/expr/sql/db.go index a1badd44974..e92ae6fc331 100644 --- a/pkg/expr/sql/db.go +++ b/pkg/expr/sql/db.go @@ -11,6 +11,7 @@ import ( "github.com/dolthub/go-mysql-server/sql/analyzer" "github.com/grafana/grafana-plugin-sdk-go/data" + "github.com/grafana/grafana/pkg/infra/tracing" ) // DB is a database that can execute SQL queries against a set of Frames. @@ -56,7 +57,7 @@ func isFunctionNotFoundError(err error) bool { // The RefID of each frame becomes a table in the database. // It is expected that there is only one frame per RefID. // The name becomes the name and RefID of the returned frame. -func (db *DB) QueryFrames(ctx context.Context, name string, query string, frames []*data.Frame) (*data.Frame, error) { +func (db *DB) QueryFrames(ctx context.Context, tracer tracing.Tracer, name string, query string, frames []*data.Frame) (*data.Frame, error) { // We are parsing twice due to TablesList, but don't care fow now. We can save the parsed query and reuse it later if we want. if allow, err := AllowQuery(query); err != nil || !allow { if err != nil { @@ -65,9 +66,14 @@ func (db *DB) QueryFrames(ctx context.Context, name string, query string, frames return nil, err } + _, span := tracer.Start(ctx, "SSE.ExecuteGMSQuery") + defer span.End() + pro := NewFramesDBProvider(frames) session := mysql.NewBaseSession() - mCtx := mysql.NewContext(ctx, mysql.WithSession(session)) + + // Create a new context with the session and tracer + mCtx := mysql.NewContext(ctx, mysql.WithSession(session), mysql.WithTracer(tracer)) // Select the database in the context mCtx.SetCurrentDatabase(dbName) diff --git a/pkg/expr/sql/db_test.go b/pkg/expr/sql/db_test.go index bb2fa28b335..b7774cbba48 100644 --- a/pkg/expr/sql/db_test.go +++ b/pkg/expr/sql/db_test.go @@ -5,12 +5,15 @@ package sql import ( "context" "encoding/json" + "net/http" "testing" "time" "github.com/google/go-cmp/cmp" "github.com/grafana/grafana-plugin-sdk-go/data" "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" ) func TestQueryFrames(t *testing.T) { @@ -61,7 +64,7 @@ func TestQueryFrames(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - frame, err := db.QueryFrames(context.Background(), "sqlExpressionRefId", tt.query, tt.input_frames) + frame, err := db.QueryFrames(context.Background(), &testTracer{}, "sqlExpressionRefId", tt.query, tt.input_frames) require.NoError(t, err) if diff := cmp.Diff(tt.expected, frame, data.FrameTestCompareOptions()...); diff != "" { @@ -123,7 +126,7 @@ func TestQueryFramesInOut(t *testing.T) { db := DB{} qry := `SELECT * from a` - resultFrame, err := db.QueryFrames(context.Background(), "a", qry, []*data.Frame{frameA}) + resultFrame, err := db.QueryFrames(context.Background(), &testTracer{}, "a", qry, []*data.Frame{frameA}) require.NoError(t, err) if diff := cmp.Diff(frameA, resultFrame, data.FrameTestCompareOptions()...); diff != "" { @@ -163,7 +166,7 @@ func TestQueryFramesNumericSelect(t *testing.T) { 4294967295 AS 'intUnsigned', 18446744073709551615 AS 'bigUnsigned'` - resultFrame, err := db.QueryFrames(context.Background(), "a", qry, []*data.Frame{}) + resultFrame, err := db.QueryFrames(context.Background(), &testTracer{}, "a", qry, []*data.Frame{}) require.NoError(t, err) if diff := cmp.Diff(expectedFrame, resultFrame, data.FrameTestCompareOptions()...); diff != "" { @@ -186,7 +189,7 @@ func TestQueryFramesDateTimeSelect(t *testing.T) { qry := `SELECT str_to_date('2025-02-03T03:00:00','%Y-%m-%dT%H:%i:%s') as ts` - f, err := db.QueryFrames(context.Background(), "a", qry, nil) + f, err := db.QueryFrames(context.Background(), &testTracer{}, "a", qry, nil) require.NoError(t, err) if diff := cmp.Diff(expectedFrame, f, data.FrameTestCompareOptions()...); diff != "" { @@ -201,7 +204,7 @@ func TestErrorsFromGoMySQLServerAreFlagged(t *testing.T) { query := `SELECT ` + GmsNotImplemented + `(123.456, 2);` - _, err := db.QueryFrames(context.Background(), "sqlExpressionRefId", query, nil) + _, err := db.QueryFrames(context.Background(), &testTracer{}, "sqlExpressionRefId", query, nil) require.Error(t, err) require.Contains(t, err.Error(), "error in go-mysql-server") } @@ -223,7 +226,7 @@ func TestFrameToSQLAndBack_JSONRoundtrip(t *testing.T) { query := `SELECT * FROM json_test` - resultFrame, err := db.QueryFrames(context.Background(), "json_test", query, data.Frames{expectedFrame}) + resultFrame, err := db.QueryFrames(context.Background(), &testTracer{}, "json_test", query, data.Frames{expectedFrame}) require.NoError(t, err) // Use custom compare options that ignore Name and RefID @@ -267,7 +270,7 @@ func TestQueryFrames_JSONFilter(t *testing.T) { query := `SELECT title, labels FROM A WHERE json_contains(labels, '"type/bug"')` - result, err := db.QueryFrames(context.Background(), "B", query, data.Frames{input}) + result, err := db.QueryFrames(context.Background(), &testTracer{}, "B", query, data.Frames{input}) require.NoError(t, err) // Use custom compare options that ignore Name and RefID @@ -287,3 +290,31 @@ func TestQueryFrames_JSONFilter(t *testing.T) { func p[T any](v T) *T { return &v } + +type testTracer struct { + trace.Tracer +} + +func (t *testTracer) Start(ctx context.Context, name string, s ...trace.SpanStartOption) (context.Context, trace.Span) { + return ctx, &testSpan{} +} +func (t *testTracer) Inject(context.Context, http.Header, trace.Span) { + +} + +type testSpan struct { + trace.Span +} + +func (ts *testSpan) End(opt ...trace.SpanEndOption) { +} + +func (ts *testSpan) SetAttributes(attr ...attribute.KeyValue) { +} + +func (ts *testSpan) IsRecording() bool { + return true +} + +func (ts *testSpan) AddEvent(name string, options ...trace.EventOption) { +} diff --git a/pkg/expr/sql/dummy_arm.go b/pkg/expr/sql/dummy_arm.go index 95d20159e08..879ba965f7d 100644 --- a/pkg/expr/sql/dummy_arm.go +++ b/pkg/expr/sql/dummy_arm.go @@ -7,12 +7,13 @@ import ( "fmt" "github.com/grafana/grafana-plugin-sdk-go/data" + "github.com/grafana/grafana/pkg/infra/tracing" ) type DB struct{} // Stub out the QueryFrames method for ARM builds // See github.com/dolthub/go-mysql-server/issues/2837 -func (db *DB) QueryFrames(_ context.Context, _, _ string, _ []*data.Frame) (*data.Frame, error) { +func (db *DB) QueryFrames(_ context.Context, _ tracing.Tracer, _, _ string, _ []*data.Frame) (*data.Frame, error) { return nil, fmt.Errorf("sql expressions not supported in arm") } diff --git a/pkg/expr/sql_command.go b/pkg/expr/sql_command.go index 25207accf4a..2bbac8fcb04 100644 --- a/pkg/expr/sql_command.go +++ b/pkg/expr/sql_command.go @@ -143,7 +143,7 @@ func (gr *SQLCommand) Execute(ctx context.Context, now time.Time, vars mathexp.V logger.Debug("Executing query", "query", gr.query, "frames", len(allFrames)) db := sql.DB{} - frame, err := db.QueryFrames(ctx, gr.refID, gr.query, allFrames) + frame, err := db.QueryFrames(ctx, tracer, gr.refID, gr.query, allFrames) rsp := mathexp.Results{} if err != nil {