SQL Expressions: Add internal GMS traces (#104836)

get trace from inside GMS

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)

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) {
+}

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")
 }

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 {