From dc0fce9a5f8d2329861471e6508245ab9365b24e Mon Sep 17 00:00:00 2001
From: Andres Taylor <andres@planetscale.com>
Date: Fri, 8 Nov 2024 09:46:50 +0100
Subject: [PATCH] bugfix: Handle CTEs with columns named in the CTE def

Signed-off-by: Andres Taylor <andres@planetscale.com>
---
 go/test/endtoend/vtgate/vitess_tester/cte/queries.test | 9 +++++++++
 go/vt/vtgate/planbuilder/operators/SQL_builder.go      | 6 +++---
 go/vt/vtgate/planbuilder/testdata/cte_cases.json       | 4 ++--
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/go/test/endtoend/vtgate/vitess_tester/cte/queries.test b/go/test/endtoend/vtgate/vitess_tester/cte/queries.test
index 173dcaeb6ba..d4aa8b27ab5 100644
--- a/go/test/endtoend/vtgate/vitess_tester/cte/queries.test
+++ b/go/test/endtoend/vtgate/vitess_tester/cte/queries.test
@@ -30,6 +30,15 @@ WITH RECURSIVE numbers AS (SELECT 1 AS n
 SELECT *
 FROM numbers;
 
+# Simple recursive CTE using literal values, column named in the CTE def
+WITH RECURSIVE numbers(n) AS (SELECT 1
+                              UNION ALL
+                              SELECT n + 1
+                              FROM numbers
+                              WHERE n < 5)
+SELECT *
+FROM numbers;
+
 # Recursive CTE joined with a normal table
 WITH RECURSIVE emp_cte AS (SELECT id, name, manager_id
                            FROM employees
diff --git a/go/vt/vtgate/planbuilder/operators/SQL_builder.go b/go/vt/vtgate/planbuilder/operators/SQL_builder.go
index 20af2a698c3..fc91569981d 100644
--- a/go/vt/vtgate/planbuilder/operators/SQL_builder.go
+++ b/go/vt/vtgate/planbuilder/operators/SQL_builder.go
@@ -242,7 +242,7 @@ func (qb *queryBuilder) unionWith(other *queryBuilder, distinct bool) {
 	}
 }
 
-func (qb *queryBuilder) recursiveCteWith(other *queryBuilder, name, alias string, distinct bool) {
+func (qb *queryBuilder) recursiveCteWith(other *queryBuilder, name, alias string, distinct bool, columns sqlparser.Columns) {
 	cteUnion := &sqlparser.Union{
 		Left:     qb.stmt.(sqlparser.SelectStatement),
 		Right:    other.stmt.(sqlparser.SelectStatement),
@@ -254,7 +254,7 @@ func (qb *queryBuilder) recursiveCteWith(other *queryBuilder, name, alias string
 			Recursive: true,
 			CTEs: []*sqlparser.CommonTableExpr{{
 				ID:       sqlparser.NewIdentifierCS(name),
-				Columns:  nil,
+				Columns:  columns,
 				Subquery: cteUnion,
 			}},
 		},
@@ -726,7 +726,7 @@ func buildRecursiveCTE(op *RecurseCTE, qb *queryBuilder) {
 		panic(err)
 	}
 
-	qb.recursiveCteWith(qbR, op.Def.Name, infoFor.GetAliasedTableExpr().As.String(), op.Distinct)
+	qb.recursiveCteWith(qbR, op.Def.Name, infoFor.GetAliasedTableExpr().As.String(), op.Distinct, op.Def.Columns)
 }
 
 func mergeHaving(h1, h2 *sqlparser.Where) *sqlparser.Where {
diff --git a/go/vt/vtgate/planbuilder/testdata/cte_cases.json b/go/vt/vtgate/planbuilder/testdata/cte_cases.json
index 016159806d7..12f92f5a9f4 100644
--- a/go/vt/vtgate/planbuilder/testdata/cte_cases.json
+++ b/go/vt/vtgate/planbuilder/testdata/cte_cases.json
@@ -2260,8 +2260,8 @@
           "Name": "main",
           "Sharded": false
         },
-        "FieldQuery": "with recursive cte as (select 1 from dual where 1 != 1 union all select n + 1 from cte where 1 != 1) select n from cte where 1 != 1",
-        "Query": "with recursive cte as (select 1 from dual union all select n + 1 from cte where n < 5) select n from cte",
+        "FieldQuery": "with recursive cte(n) as (select 1 from dual where 1 != 1 union all select n + 1 from cte where 1 != 1) select n from cte where 1 != 1",
+        "Query": "with recursive cte(n) as (select 1 from dual union all select n + 1 from cte where n < 5) select n from cte",
         "Table": "dual"
       },
       "TablesUsed": [