From 1be9e6066724d3e31c5a96703cfc449438d5b8b0 Mon Sep 17 00:00:00 2001 From: Kyle Brandt Date: Wed, 11 Jun 2025 10:36:28 -0400 Subject: [PATCH] SQL Expressions: Exclude CTEs from required Tables (#106479) Fixes #105030 --------- Co-authored-by: Sam Jewell <2903904+samjewell@users.noreply.github.com> --- pkg/expr/graph.go | 8 ------ pkg/expr/sql/parser.go | 22 ++++++++++++-- pkg/expr/sql/parser_test.go | 2 +- pkg/registry/apis/query/parser.go | 10 ++----- pkg/registry/apis/query/parser_test.go | 40 ++++++++++++++++++++++++++ 5 files changed, 62 insertions(+), 20 deletions(-) diff --git a/pkg/expr/graph.go b/pkg/expr/graph.go index 4d760dff930..e19b73dd15e 100644 --- a/pkg/expr/graph.go +++ b/pkg/expr/graph.go @@ -328,14 +328,6 @@ func buildGraphEdges(dp *simple.DirectedGraph, registry map[string]Node) error { for _, neededVar := range cmdNode.Command.NeedsVars() { neededNode, ok := registry[neededVar] if !ok { - _, ok := cmdNode.Command.(*SQLCommand) - // If the SSE is a SQL expression, and the node can't be found, it might be a CTE table name - // CTEs are calculated during the evaluation of the SQL, so we won't have a node for them - // So we `continue` in order to support CTE functionality - // TODO: remove CTE table names from the list of table names during parsing of the SQL - if ok { - continue - } return fmt.Errorf("unable to find dependent node '%v'", neededVar) } diff --git a/pkg/expr/sql/parser.go b/pkg/expr/sql/parser.go index c65e4b89131..0e401c94a15 100644 --- a/pkg/expr/sql/parser.go +++ b/pkg/expr/sql/parser.go @@ -3,6 +3,7 @@ package sql import ( "fmt" "sort" + "strings" "github.com/dolthub/vitess/go/vt/sqlparser" "github.com/grafana/grafana/pkg/infra/log" @@ -10,7 +11,8 @@ import ( var logger = log.New("sql_expr") -// TablesList returns a list of tables for the sql statement +// TablesList returns a list of tables for the sql statement excluding +// CTEs and the 'dual' table. The list is sorted alphabetically. func TablesList(rawSQL string) ([]string, error) { stmt, err := sqlparser.Parse(rawSQL) if err != nil { @@ -19,10 +21,18 @@ func TablesList(rawSQL string) ([]string, error) { } tables := make(map[string]struct{}) + cteNames := make(map[string]struct{}) walkSubtree := func(node sqlparser.SQLNode) error { err = sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) { switch v := node.(type) { + case *sqlparser.CommonTableExpr: + // Track CTE name from the As field + cteName := v.As.String() + if cteName != "" { + cteNames[strings.ToLower(cteName)] = struct{}{} + } + case *sqlparser.AliasedTableExpr: if tableName, ok := v.Expr.(sqlparser.TableName); ok { tables[tableName.Name.String()] = struct{}{} @@ -49,9 +59,15 @@ func TablesList(rawSQL string) ([]string, error) { // Remove 'dual' table if it exists // This is a special table in MySQL that always returns a single row with a single column // See: https://dev.mysql.com/doc/refman/5.7/en/select.html#:~:text=You%20are%20permitted%20to%20specify%20DUAL%20as%20a%20dummy%20table%20name%20in%20situations%20where%20no%20tables%20are%20referenced - if table != "dual" { - result = append(result, table) + if table == "dual" { + continue + } + + // Skip CTEs + if _, ok := cteNames[strings.ToLower(table)]; ok { + continue } + result = append(result, table) } sort.Strings(result) diff --git a/pkg/expr/sql/parser_test.go b/pkg/expr/sql/parser_test.go index 17235684388..44ac1ccb102 100644 --- a/pkg/expr/sql/parser_test.go +++ b/pkg/expr/sql/parser_test.go @@ -105,7 +105,7 @@ func TestTablesList(t *testing.T) { ) SELECT name, price FROM top_products;`, - expected: []string{"products", "top_products"}, + expected: []string{"products"}, }, { name: "with quote", diff --git a/pkg/registry/apis/query/parser.go b/pkg/registry/apis/query/parser.go index 965a0a173c3..96089f763d6 100644 --- a/pkg/registry/apis/query/parser.go +++ b/pkg/registry/apis/query/parser.go @@ -167,14 +167,8 @@ func (p *queryParser) parseRequest(ctx context.Context, input *query.QueryDataRe q, ok := queryRefIDs[refId] if !ok { - _, isSQLCMD := target.Command.(*expr.SQLCommand) - if isSQLCMD { - continue - } else { - target, ok = expressions[refId] - if !ok { - return rsp, makeDependencyError(exp.RefID, refId) - } + if target, ok = expressions[refId]; !ok { + return rsp, makeDependencyError(exp.RefID, refId) } } diff --git a/pkg/registry/apis/query/parser_test.go b/pkg/registry/apis/query/parser_test.go index e3e189f7ab1..f43e5295631 100644 --- a/pkg/registry/apis/query/parser_test.go +++ b/pkg/registry/apis/query/parser_test.go @@ -201,6 +201,46 @@ func TestSqlInputs(t *testing.T) { require.Equal(t, parsedRequestInfo.SqlInputs["B"], struct{}{}) } +func TestSqlCTE(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": `WITH CTE AS ( + SELECT + Month + FROM A + ) + + SELECT * FROM CTE`, + }), + }, + }, + }) + require.NoError(t, err) + require.Equal(t, parsedRequestInfo.SqlInputs["B"], struct{}{}) +} + func TestGrafanaDS(t *testing.T) { ctx := context.Background() parser := newQueryParser(expr.NewExpressionQueryReader(featuremgmt.WithFeatures()),