SQL Expressions: Add sql expression specific timeout and output limit (#104834)

Adds settings for SQL expressions: sql_expression_cell_output_limit Set the maximum number of cells that can be returned from a SQL expression. Default is 100000. sql_expression_timeout The duration a SQL expression will run before being cancelled. The default is 10s.
2025-05-13 15:22:20 -04:00
parent 02d977e1af
commit 5e056c2a3f
12 changed files with 183 additions and 22 deletions
@@ -4,7 +4,9 @@ package sql

 import (
 	"context"
+	"errors"
 	"fmt"
+	"time"

 	sqle "github.com/dolthub/go-mysql-server"
 	mysql "github.com/dolthub/go-mysql-server/sql"
@@ -53,11 +55,30 @@ func isFunctionNotFoundError(err error) bool {
 	return mysql.ErrFunctionNotFound.Is(err)
 }

+type QueryOption func(*QueryOptions)
+
+type QueryOptions struct {
+	Timeout        time.Duration
+	MaxOutputCells int64
+}
+
+func WithTimeout(d time.Duration) QueryOption {
+	return func(o *QueryOptions) {
+		o.Timeout = d
+	}
+}
+
+func WithMaxOutputCells(n int64) QueryOption {
+	return func(o *QueryOptions) {
+		o.MaxOutputCells = n
+	}
+}
+
 // QueryFrames runs the sql query query against a database created from frames, and returns the frame.
 // 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, tracer tracing.Tracer, 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, opts ...QueryOption) (*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 {
@@ -66,6 +87,16 @@ func (db *DB) QueryFrames(ctx context.Context, tracer tracing.Tracer, name strin
 		return nil, err
 	}

+	QueryOptions := &QueryOptions{}
+	for _, opt := range opts {
+		opt(QueryOptions)
+	}
+
+	if QueryOptions.Timeout != 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, QueryOptions.Timeout)
+		defer cancel()
+	}
 	_, span := tracer.Start(ctx, "SSE.ExecuteGMSQuery")
 	defer span.End()

@@ -88,15 +119,35 @@ func (db *DB) QueryFrames(ctx context.Context, tracer tracing.Tracer, name strin
 		IsReadOnly: true,
 	})

+	contextErr := func(err error) error {
+		switch {
+		case errors.Is(err, context.DeadlineExceeded):
+			return fmt.Errorf("SQL expression for refId %v did not complete within the timeout of %v: %w", name, QueryOptions.Timeout, err)
+		case errors.Is(err, context.Canceled):
+			return fmt.Errorf("SQL expression for refId %v was cancelled before it completed: %w", name, err)
+		default:
+			return fmt.Errorf("SQL expression for refId %v ended unexpectedly: %w", name, err)
+		}
+	}
+
+	// Execute the query (planning + iterator construction)
 	schema, iter, _, err := engine.Query(mCtx, query)
 	if err != nil {
+		if ctx.Err() != nil {
+			return nil, contextErr(ctx.Err())
+		}
 		return nil, WrapGoMySQLServerError(err)
 	}

-	f, err := convertToDataFrame(mCtx, iter, schema)
+	// Convert the iterator into a Grafana data.Frame
+	f, err := convertToDataFrame(mCtx, iter, schema, QueryOptions.MaxOutputCells)
 	if err != nil {
+		if ctx.Err() != nil {
+			return nil, contextErr(ctx.Err())
+		}
 		return nil, err
 	}
+
 	f.Name = name
 	f.RefID = name

@@ -286,6 +286,51 @@ func TestQueryFrames_JSONFilter(t *testing.T) {
 	}
 }

+func TestQueryFrames_Limits(t *testing.T) {
+	tests := []struct {
+		name        string
+		query       string
+		opts        []QueryOption
+		expectRows  int
+		expectError string
+	}{
+		{
+			name:       "respects max output cells",
+			query:      `SELECT 1 as x UNION ALL SELECT 2 UNION ALL SELECT 3`,
+			opts:       []QueryOption{WithMaxOutputCells(2)},
+			expectRows: 2,
+		},
+		{
+			name: "timeout with large cross join",
+			query: `
+				SELECT a.val + b.val AS sum
+				FROM (SELECT 1 AS val UNION ALL SELECT 2 UNION ALL SELECT 3 UNION ALL SELECT 4 UNION ALL SELECT 5) a
+				CROSS JOIN (SELECT 1 AS val UNION ALL SELECT 2 UNION ALL SELECT 3 UNION ALL SELECT 4 UNION ALL SELECT 5) b
+			`,
+			opts:        []QueryOption{WithTimeout(5 * time.Microsecond)},
+			expectError: "did not complete within the timeout",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			db := DB{}
+			ctx := context.Background()
+			frame, err := db.QueryFrames(ctx, &testTracer{}, "test", tt.query, nil, tt.opts...)
+
+			if tt.expectError != "" {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tt.expectError)
+				return
+			}
+
+			require.NoError(t, err)
+			require.NotNil(t, frame)
+			require.Equal(t, tt.expectRows, frame.Rows())
+		})
+	}
+}
+
 // p is a utility for pointers from constants
 func p[T any](v T) *T {
 	return &v
@@ -5,6 +5,7 @@ package sql
 import (
 	"context"
 	"fmt"
+	"time"

 	"github.com/grafana/grafana-plugin-sdk-go/data"
 	"github.com/grafana/grafana/pkg/infra/tracing"
@@ -14,6 +15,22 @@ 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, _ tracing.Tracer, _, _ string, _ []*data.Frame) (*data.Frame, error) {
+func (db *DB) QueryFrames(_ context.Context, _ tracing.Tracer, _, _ string, _ []*data.Frame, _...QueryOption) (*data.Frame, error) {
 	return nil, fmt.Errorf("sql expressions not supported in arm")
 }
+
+func WithTimeout(_ time.Duration) QueryOption {
+	return func(_ *QueryOptions) {
+		// no-op
+	}
+}
+
+func WithMaxOutputCells(_ int64) QueryOption {
+	return func(_ *QueryOptions) {
+		// no-op
+	}
+}
+
+type QueryOptions struct{}
+
+type QueryOption func(*QueryOptions)
@@ -16,8 +16,9 @@ import (
 )

 // TODO: Should this accept a row limit and converters, like sqlutil.FrameFromRows?
-func convertToDataFrame(ctx *mysql.Context, iter mysql.RowIter, schema mysql.Schema) (*data.Frame, error) {
+func convertToDataFrame(ctx *mysql.Context, iter mysql.RowIter, schema mysql.Schema, maxOutputCells int64) (*data.Frame, error) {
 	f := &data.Frame{}
+
 	// Create fields based on the schema
 	for _, col := range schema {
 		fT, err := MySQLColToFieldType(col)
@@ -29,8 +30,17 @@ func convertToDataFrame(ctx *mysql.Context, iter mysql.RowIter, schema mysql.Sch
 		f.Fields = append(f.Fields, field)
 	}

+	cellCount := int64(0)
+
 	// Iterate through the rows and append data to fields
 	for {
+		// Check for context cancellation or timeout
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+		}
+
 		row, err := iter.Next(ctx)
 		if errors.Is(err, io.EOF) {
 			break
@@ -39,6 +49,20 @@ func convertToDataFrame(ctx *mysql.Context, iter mysql.RowIter, schema mysql.Sch
 			return nil, fmt.Errorf("error reading row: %v", err)
 		}

+		// We check the cell count here to avoid appending an incomplete row, so the
+		// the number returned may be less than the maxOutputCells.
+		// If the maxOutputCells is 0, we don't check the cell count.
+		if maxOutputCells > 0 {
+			cellCount += int64(len(row))
+			if cellCount > maxOutputCells {
+				f.AppendNotices(data.Notice{
+					Severity: data.NoticeSeverityWarning,
+					Text:     fmt.Sprintf("Query exceeded max output cells (%d). Only %d cells returned.", maxOutputCells, cellCount-int64(len(row))),
+				})
+				return f, nil
+			}
+		}
+
 		for i, val := range row {
 			// Run val through mysql.Type.Convert to normalize underlying value
 			// of the interface