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

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

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" {
+		if table == "dual" {
-			result = append(result, table)
+			continue
 		}
+
+		// Skip CTEs
+		if _, ok := cteNames[strings.ToLower(table)]; ok {
+			continue
+		}
+		result = append(result, table)
 	}
 
 	sort.Strings(result)

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",

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 target, ok = expressions[refId]; !ok {
-					if isSQLCMD {
+						return rsp, makeDependencyError(exp.RefID, refId)
-						continue
-					} else {
-						target, ok = expressions[refId]
-						if !ok {
-							return rsp, makeDependencyError(exp.RefID, refId)
-						}
 					}
 				}
 

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()),