CTEs on unions parse

go/vt/sqlparser/ast.go

@@ -16,6 +16,8 @@ limitations under the License.
 
 package sqlparser
 
+//go:generate goyacc -o sql.go sql.y
+
 import (
 	"encoding/hex"
 	"encoding/json"
@@ -421,6 +423,7 @@ type SelectStatement interface {
 	SetLimit(*Limit)
 	SetLock(string)
 	SetOrderBy(OrderBy)
+	SetWith(*With)
 	SetInto(*Into) error
 	GetInto() *Into
 	WalkableSQLNode
@@ -478,6 +481,10 @@ func (node *Select) SetOrderBy(orderBy OrderBy) {
 	node.OrderBy = orderBy
 }
 
+func (node *Select) SetWith(w *With) {
+	node.With = w
+}
+
 func (node *Select) SetLock(lock string) {
 	node.Lock = lock
 }
@@ -592,6 +599,10 @@ func (node *ParenSelect) SetOrderBy(orders OrderBy) {
 	panic("unreachable")
 }
 
+func (node *ParenSelect) SetWith(w *With) {
+	panic("unreachable")
+}
+
 func (node *ParenSelect) SetLock(lock string) {
 	panic("unreachable")
 }
@@ -650,6 +661,7 @@ type Union struct {
 	Type        string
 	Left, Right SelectStatement
 	OrderBy     OrderBy
+	With        *With
 	Limit       *Limit
 	Lock        string
 	Into        *Into
@@ -671,6 +683,10 @@ func (node *Union) SetOrderBy(orderBy OrderBy) {
 	node.OrderBy = orderBy
 }
 
+func (node *Union) SetWith(w *With) {
+	node.With = w
+}
+
 // SetLimit sets the limit clause
 func (node *Union) SetLimit(limit *Limit) {
 	node.Limit = limit
@@ -701,7 +717,7 @@ func (node *Union) GetInto() *Into {
 
 // Format formats the node.
 func (node *Union) Format(buf *TrackedBuffer) {
-	buf.Myprintf("%v %s %v%v%v%s%v", node.Left, node.Type, node.Right,
+	buf.Myprintf("%v%v %s %v%v%v%s%v", node.With, node.Left, node.Type, node.Right,
 		node.OrderBy, node.Limit, node.Lock, node.Into)
 }
 

go/vt/sqlparser/parse_test.go

@@ -2707,6 +2707,18 @@ BEGIN
 END`,
 			output: "create definer = `root`@`localhost` procedure film_not_in_stock (in p_film_id INT, in p_store_id INT, out p_film_count INT) reads sql data begin\nselect inventory_id from inventory where film_id = p_film_id and store_id = p_store_id;\nselect COUNT(*) from inventory where film_id = p_film_id and store_id = p_store_id into p_film_count;\nend",
 		},
+		{
+			input:  "with a(j) as (select 1), b(i) as (select 2) (select j from a union select i from b order by j desc limit 1) union select j from a;",
+			output: "with a (j) as (select 1 from dual), b (i) as (select 2 from dual) (select j from a union select i from b order by j desc limit 1) union select j from a",
+		},
+		{
+			input:  "with a(j) as (select 1) ( with c(k) as (select 3) select k from c union select 6) union select k from c;",
+			output: "with a (j) as (select 1 from dual) (with c (k) as (select 3 from dual) select k from c union select 6 from dual) union select k from c",
+		},
+		{
+			input:  "with a(j) as (select 1) ( with c(k) as (select 3) select (select k from c union select 6 limit 1) as b) union select k from c;",
+			output: "with a (j) as (select 1 from dual) (with c (k) as (select 3 from dual) select (select k from c union select 6 from dual limit 1) as b from dual) union select k from c",
+		},
 	}
 )