From 4a0ec27e5d02d6b169eca1de3e18ff4ddac8e999 Mon Sep 17 00:00:00 2001 From: Kyle Brandt Date: Thu, 3 Apr 2025 09:36:02 -0400 Subject: [PATCH] SQL Expressions: Query Service Support (#101955) --------- Co-authored-by: Adam Simpson Co-authored-by: Sarah Zinger --- pkg/expr/converter.go | 43 ++++++++++++++++++++++++++ pkg/expr/converter_test.go | 6 ++-- pkg/expr/ml.go | 2 +- pkg/expr/nodes.go | 39 ++--------------------- pkg/expr/reader.go | 13 ++------ pkg/registry/apis/query/parser.go | 31 +++++++++++++++++-- pkg/registry/apis/query/parser_test.go | 34 ++++++++++++++++++++ pkg/registry/apis/query/query.go | 7 +++-- 8 files changed, 118 insertions(+), 57 deletions(-) diff --git a/pkg/expr/converter.go b/pkg/expr/converter.go index 0e8cc461d47..c804662304d 100644 --- a/pkg/expr/converter.go +++ b/pkg/expr/converter.go @@ -23,11 +23,17 @@ type ResultConverter struct { func (c *ResultConverter) Convert(ctx context.Context, datasourceType string, frames data.Frames, + forSqlInput bool, ) (string, mathexp.Results, error) { if len(frames) == 0 { return "no-data", mathexp.Results{Values: mathexp.Values{mathexp.NewNoData()}}, nil } + if forSqlInput { + results, err := handleSqlInput(frames) + return "sql input", results, err + } + var dt data.FrameType dt, useDataplane, _ := shouldUseDataplane(frames, logger, c.Features.IsEnabled(ctx, featuremgmt.FlagDisableSSEDataplane)) if useDataplane { @@ -120,6 +126,43 @@ func (c *ResultConverter) Convert(ctx context.Context, }, nil } +// copied from pkg/expr/nodes.go from within the Execute method +func handleSqlInput(dataFrames data.Frames) (mathexp.Results, error) { + var result mathexp.Results + var needsConversion bool + // Convert it if Multi: + if len(dataFrames) > 1 { + needsConversion = true + } + + // Convert it if Wide (has labels): + if len(dataFrames) == 1 { + for _, field := range dataFrames[0].Fields { + if len(field.Labels) > 0 { + needsConversion = true + break + } + } + } + + if needsConversion { + convertedFrames, err := ConvertToFullLong(dataFrames) + if err != nil { + return result, fmt.Errorf("failed to convert data frames to long format for sql: %w", err) + } + result.Values = mathexp.Values{ + mathexp.TableData{Frame: convertedFrames[0]}, + } + return result, nil + } + + // Otherwise it is already Long format; return as is + result.Values = mathexp.Values{ + mathexp.TableData{Frame: dataFrames[0]}, + } + return result, nil +} + func getResponseFrame(logger *log.ConcreteLogger, resp *backend.QueryDataResponse, refID string) (data.Frames, error) { response, ok := resp.Responses[refID] if !ok { diff --git a/pkg/expr/converter_test.go b/pkg/expr/converter_test.go index 1509294e3de..d5bb8e64f74 100644 --- a/pkg/expr/converter_test.go +++ b/pkg/expr/converter_test.go @@ -40,7 +40,7 @@ func TestConvertDataFramesToResults(t *testing.T) { for _, dtype := range supported { t.Run(dtype, func(t *testing.T) { - resultType, res, err := converter.Convert(context.Background(), dtype, frames) + resultType, res, err := converter.Convert(context.Background(), dtype, frames, false) require.NoError(t, err) assert.Equal(t, "single frame series", resultType) require.Len(t, res.Values, 2) @@ -68,7 +68,7 @@ func TestConvertDataFramesToResults(t *testing.T) { for _, dtype := range supported { t.Run(dtype, func(t *testing.T) { - resultType, res, err := converter.Convert(context.Background(), dtype, frames) + resultType, res, err := converter.Convert(context.Background(), dtype, frames, false) require.NoError(t, err) assert.Equal(t, "multi frame series", resultType) require.Len(t, res.Values, 2) @@ -101,7 +101,7 @@ func TestConvertDataFramesToResults(t *testing.T) { for _, dtype := range supported { t.Run(dtype, func(t *testing.T) { - resultType, res, err := converter.Convert(context.Background(), dtype, frames) + resultType, res, err := converter.Convert(context.Background(), dtype, frames, false) require.NoError(t, err) assert.Equal(t, "multi frame series", resultType) require.Len(t, res.Values, 2) diff --git a/pkg/expr/ml.go b/pkg/expr/ml.go index 4c0dba76b89..9efa143e805 100644 --- a/pkg/expr/ml.go +++ b/pkg/expr/ml.go @@ -130,7 +130,7 @@ func (m *MLNode) Execute(ctx context.Context, now time.Time, _ mathexp.Vars, s * } // process the response the same way DSNode does. Use plugin ID as data source type. Semantically, they are the same. - responseType, result, err = s.converter.Convert(ctx, mlPluginID, dataFrames) + responseType, result, err = s.converter.Convert(ctx, mlPluginID, dataFrames, false) return result, err } diff --git a/pkg/expr/nodes.go b/pkg/expr/nodes.go index 141ce36e3af..e18bad3a7cb 100644 --- a/pkg/expr/nodes.go +++ b/pkg/expr/nodes.go @@ -341,7 +341,7 @@ func executeDSNodesGrouped(ctx context.Context, now time.Time, vars mathexp.Vars } var result mathexp.Results - responseType, result, err := s.converter.Convert(ctx, dn.datasource.Type, dataFrames) + responseType, result, err := s.converter.Convert(ctx, dn.datasource.Type, dataFrames, dn.isInputToSQLExpr) if err != nil { result.Error = makeConversionError(dn.RefID(), err) } @@ -409,44 +409,9 @@ func (dn *DSNode) Execute(ctx context.Context, now time.Time, _ mathexp.Vars, s } var result mathexp.Results - // If the datasource node is an input to a SQL expression, - // the data must be in the Long format - if dn.isInputToSQLExpr { - var needsConversion bool - // Convert it if Multi: - if len(dataFrames) > 1 { - needsConversion = true - } - // Convert it if Wide (has labels): - if len(dataFrames) == 1 { - for _, field := range dataFrames[0].Fields { - if len(field.Labels) > 0 { - needsConversion = true - break - } - } - } + responseType, result, err = s.converter.Convert(ctx, dn.datasource.Type, dataFrames, dn.isInputToSQLExpr) - if needsConversion { - convertedFrames, err := ConvertToFullLong(dataFrames) - if err != nil { - return result, fmt.Errorf("failed to convert data frames to long format for sql: %w", err) - } - result.Values = mathexp.Values{ - mathexp.TableData{Frame: convertedFrames[0]}, - } - return result, nil - } - - // Otherwise it is already Long format; return as is - result.Values = mathexp.Values{ - mathexp.TableData{Frame: dataFrames[0]}, - } - return result, nil - } - - responseType, result, err = s.converter.Convert(ctx, dn.datasource.Type, dataFrames) if err != nil { err = makeConversionError(dn.refID, err) } diff --git a/pkg/expr/reader.go b/pkg/expr/reader.go index ed692358d19..6d1349ae84c 100644 --- a/pkg/expr/reader.go +++ b/pkg/expr/reader.go @@ -126,9 +126,8 @@ func (h *ExpressionQueryReader) ReadQuery( } case QueryTypeSQL: - enabled := enableSqlExpressions(h) - if !enabled { - return eq, fmt.Errorf("sqlExpressions is not implemented") + if !h.features.IsEnabledGlobally(featuremgmt.FlagSqlExpressions) { + return eq, fmt.Errorf("sql expressions are disabled") } q := &SQLExpression{} err = iter.ReadVal(q) @@ -192,11 +191,3 @@ func getReferenceVar(exp string, refId string) (string, error) { } return exp, nil } - -func enableSqlExpressions(h *ExpressionQueryReader) bool { - enabled := !h.features.IsEnabledGlobally(featuremgmt.FlagSqlExpressions) - if enabled { - return false - } - return false -} diff --git a/pkg/registry/apis/query/parser.go b/pkg/registry/apis/query/parser.go index f36642c33d8..1333a348846 100644 --- a/pkg/registry/apis/query/parser.go +++ b/pkg/registry/apis/query/parser.go @@ -43,6 +43,9 @@ type parsedRequestInfo struct { // Hidden queries used as dependencies HideBeforeReturn []string `json:"hide,omitempty"` + + // SQL Inputs + SqlInputs map[string]struct{} `json:"sqlInputs,omitempty"` } type queryParser struct { @@ -71,6 +74,7 @@ func (p *queryParser) parseRequest(ctx context.Context, input *query.QueryDataRe index := make(map[string]int) // index lookup rsp := parsedRequestInfo{ RefIDTypes: make(map[string]string, len(input.Queries)), + SqlInputs: make(map[string]struct{}), } for _, q := range input.Queries { @@ -90,6 +94,7 @@ func (p *queryParser) parseRequest(ctx context.Context, input *query.QueryDataRe } // Process each query + // check if ds is expression if expr.IsDataSource(ds.UID) { // In order to process the query as a typed expression query, we // are writing it back to JSON and parsing again. Alternatively we @@ -149,20 +154,40 @@ func (p *queryParser) parseRequest(ctx context.Context, input *query.QueryDataRe // Build the graph for a request dg := simple.NewDirectedGraph() dg.AddNode(queryNode) + for _, exp := range expressions { dg.AddNode(exp) } + for _, exp := range expressions { vars := exp.Command.NeedsVars() + for _, refId := range vars { target := queryNode q, ok := queryRefIDs[refId] + if !ok { - target, ok = expressions[refId] - if !ok { - return rsp, makeDependencyError(exp.RefID, refId) + _, isSQLCMD := target.Command.(*expr.SQLCommand) + if isSQLCMD { + continue + } else { + target, ok = expressions[refId] + if !ok { + return rsp, makeDependencyError(exp.RefID, refId) + } } } + + // If the input is SQL, conversion is handled differently + if _, isSqlExp := exp.Command.(*expr.SQLCommand); isSqlExp { + if _, ifDepIsAlsoExpression := expressions[refId]; ifDepIsAlsoExpression { + // Only allow data source nodes as SQL expression inputs for now + return rsp, fmt.Errorf("only data source queries may be inputs to a sql expression, %v is the input for %v", refId, exp.RefID) + } else { + rsp.SqlInputs[refId] = struct{}{} + } + } + // Do not hide queries used in variables if q != nil && q.Hide { q.Hide = false diff --git a/pkg/registry/apis/query/parser_test.go b/pkg/registry/apis/query/parser_test.go index d775cbf55fe..f2a2d7d518d 100644 --- a/pkg/registry/apis/query/parser_test.go +++ b/pkg/registry/apis/query/parser_test.go @@ -166,6 +166,40 @@ func TestQuerySplitting(t *testing.T) { }) } +func TestSqlInputs(t *testing.T) { + parser := newQueryParser( + expr.NewExpressionQueryReader(featuremgmt.WithFeatures(featuremgmt.FlagSqlExpressions)), + nil, + tracing.InitializeTracerForTest(), + log.NewNopLogger(), + ) + + parsedRequestInfo, err := parser.parseRequest(context.Background(), &query.QueryDataRequest{ + QueryDataRequest: data.QueryDataRequest{ + Queries: []data.DataQuery{ + data.NewDataQuery(map[string]any{ + "refId": "A", + "datasource": &data.DataSourceRef{ + Type: "prometheus", + UID: "local-prom", + }, + }), + data.NewDataQuery(map[string]any{ + "refId": "B", + "datasource": &data.DataSourceRef{ + Type: "__expr__", + UID: "__expr__", + }, + "type": "sql", + "expression": "Select time, value + 10 from A", + }), + }, + }, + }) + require.NoError(t, err) + require.Equal(t, parsedRequestInfo.SqlInputs["B"], struct{}{}) +} + type legacyDataSourceRetriever struct{} func (s *legacyDataSourceRetriever) GetDataSourceFromDeprecatedFields(ctx context.Context, name string, id int64) (*data.DataSourceRef, error) { diff --git a/pkg/registry/apis/query/query.go b/pkg/registry/apis/query/query.go index a18d488f7f9..ddcf88a3d7e 100644 --- a/pkg/registry/apis/query/query.go +++ b/pkg/registry/apis/query/query.go @@ -381,11 +381,14 @@ func (b *QueryAPIBuilder) handleExpressions(ctx context.Context, req parsedReque if !ok { dr, ok := qdr.Responses[refId] if ok { - _, res, err := b.converter.Convert(ctx, req.RefIDTypes[refId], dr.Frames) + _, isSqlInput := req.SqlInputs[refId] + + _, res, err := b.converter.Convert(ctx, req.RefIDTypes[refId], dr.Frames, isSqlInput) if err != nil { expressionsLogger.Error("error converting frames for expressions", "error", err) res.Error = err } + vars[refId] = res } else { expressionsLogger.Error("missing variable in handle expressions", "refId", refId, "expressionRefId", expression.RefID) @@ -427,7 +430,7 @@ func (b *QueryAPIBuilder) convertQueryWithoutExpression(ctx context.Context, req return nil, fmt.Errorf("refID '%s' does not exist", refID) } frames := qdr.Responses[refID].Frames - _, results, err := b.converter.Convert(ctx, req.PluginId, frames) + _, results, err := b.converter.Convert(ctx, req.PluginId, frames, false) if err != nil { results.Error = err }