parser: fix limit lost in SetOprStmt (pingcap#49898) (pingcap#49931)

close pingcap#49874
ti-chi-bot · Jan 4, 2024 · ae491e2 · ae491e2
1 parent 4d53073
commit ae491e2
Show file tree

Hide file tree

Showing 4 changed files with 106 additions and 61 deletions.
diff --git a/parser/ast/dml.go b/parser/ast/dml.go
@@ -1562,6 +1562,7 @@ func (n *SetOprSelectList) Restore(ctx *format.RestoreCtx) error {
 			return errors.Annotate(err, "An error occurred while restore SetOprSelectList.With")
 		}
 	}
+
 	for i, stmt := range n.Selects {
 		switch selectStmt := stmt.(type) {
 		case *SelectStmt:
@@ -1583,6 +1584,20 @@ func (n *SetOprSelectList) Restore(ctx *format.RestoreCtx) error {
 			ctx.WritePlain(")")
 		}
 	}
+
+	if n.OrderBy != nil {
+		ctx.WritePlain(" ")
+		if err := n.OrderBy.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore SetOprSelectList.OrderBy")
+		}
+	}
+
+	if n.Limit != nil {
+		ctx.WritePlain(" ")
+		if err := n.Limit.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore SetOprSelectList.Limit")
+		}
+	}
 	return nil
 }
 
@@ -1607,6 +1622,20 @@ func (n *SetOprSelectList) Accept(v Visitor) (Node, bool) {
 		}
 		n.Selects[i] = node
 	}
+	if n.OrderBy != nil {
+		node, ok := n.OrderBy.Accept(v)
+		if !ok {
+			return n, false
+		}
+		n.OrderBy = node.(*OrderByClause)
+	}
+	if n.Limit != nil {
+		node, ok := n.Limit.Accept(v)
+		if !ok {
+			return n, false
+		}
+		n.Limit = node.(*Limit)
+	}
 	return v.Leave(n)
 }
 

diff --git a/parser/parser.go b/parser/parser.go
@@ -19393,15 +19393,19 @@ yynewstate:
 			}
 			var setOprList2 []ast.Node
 			var with2 *ast.WithClause
+			var limit2 *ast.Limit
+			var orderBy2 *ast.OrderByClause
 			switch x := yyS[yypt-1].expr.(*ast.SubqueryExpr).Query.(type) {
 			case *ast.SelectStmt:
 				setOprList2 = []ast.Node{x}
 				with2 = x.With
 			case *ast.SetOprStmt:
 				setOprList2 = x.SelectList.Selects
 				with2 = x.With
+				limit2 = x.Limit
+				orderBy2 = x.OrderBy
 			}
-			nextSetOprList := &ast.SetOprSelectList{Selects: setOprList2, With: with2}
+			nextSetOprList := &ast.SetOprSelectList{Selects: setOprList2, With: with2, Limit: limit2, OrderBy: orderBy2}
 			nextSetOprList.AfterSetOperator = yyS[yypt-2].item.(*ast.SetOprType)
 			setOprList := append(setOprList1, nextSetOprList)
 			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}}
@@ -19417,15 +19421,19 @@ yynewstate:
 			}
 			var setOprList2 []ast.Node
 			var with2 *ast.WithClause
+			var limit2 *ast.Limit
+			var orderBy2 *ast.OrderByClause
 			switch x := yyS[yypt-1].expr.(*ast.SubqueryExpr).Query.(type) {
 			case *ast.SelectStmt:
 				setOprList2 = []ast.Node{x}
 				with2 = x.With
 			case *ast.SetOprStmt:
 				setOprList2 = x.SelectList.Selects
 				with2 = x.With
+				limit2 = x.Limit
+				orderBy2 = x.OrderBy
 			}
-			nextSetOprList := &ast.SetOprSelectList{Selects: setOprList2, With: with2}
+			nextSetOprList := &ast.SetOprSelectList{Selects: setOprList2, With: with2, Limit: limit2, OrderBy: orderBy2}
 			nextSetOprList.AfterSetOperator = yyS[yypt-2].item.(*ast.SetOprType)
 			setOprList := append(setOprList1, nextSetOprList)
 			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}}
@@ -19441,15 +19449,19 @@ yynewstate:
 			}
 			var setOprList2 []ast.Node
 			var with2 *ast.WithClause
+			var limit2 *ast.Limit
+			var orderBy2 *ast.OrderByClause
 			switch x := yyS[yypt-2].expr.(*ast.SubqueryExpr).Query.(type) {
 			case *ast.SelectStmt:
 				setOprList2 = []ast.Node{x}
 				with2 = x.With
 			case *ast.SetOprStmt:
 				setOprList2 = x.SelectList.Selects
 				with2 = x.With
+				limit2 = x.Limit
+				orderBy2 = x.OrderBy
 			}
-			nextSetOprList := &ast.SetOprSelectList{Selects: setOprList2, With: with2}
+			nextSetOprList := &ast.SetOprSelectList{Selects: setOprList2, With: with2, Limit: limit2, OrderBy: orderBy2}
 			nextSetOprList.AfterSetOperator = yyS[yypt-3].item.(*ast.SetOprType)
 			setOprList := append(setOprList1, nextSetOprList)
 			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}}
@@ -19460,48 +19472,39 @@ yynewstate:
 	case 1793:
 		{
 			var setOprList []ast.Node
-			var with *ast.WithClause
 			switch x := yyS[yypt-1].expr.(*ast.SubqueryExpr).Query.(type) {
 			case *ast.SelectStmt:
-				setOprList = []ast.Node{x}
-				with = x.With
+				setOprList = []ast.Node{&ast.SetOprSelectList{Selects: []ast.Node{x}, With: x.With}}
 			case *ast.SetOprStmt:
-				setOprList = x.SelectList.Selects
-				with = x.With
+				setOprList = []ast.Node{&ast.SetOprSelectList{Selects: x.SelectList.Selects, With: x.With, Limit: x.Limit, OrderBy: x.OrderBy}}
 			}
-			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}, With: with}
+			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}}
 			setOpr.OrderBy = yyS[yypt-0].item.(*ast.OrderByClause)
 			parser.yyVAL.statement = setOpr
 		}
 	case 1794:
 		{
 			var setOprList []ast.Node
-			var with *ast.WithClause
 			switch x := yyS[yypt-1].expr.(*ast.SubqueryExpr).Query.(type) {
 			case *ast.SelectStmt:
-				setOprList = []ast.Node{x}
-				with = x.With
+				setOprList = []ast.Node{&ast.SetOprSelectList{Selects: []ast.Node{x}, With: x.With}}
 			case *ast.SetOprStmt:
-				setOprList = x.SelectList.Selects
-				with = x.With
+				setOprList = []ast.Node{&ast.SetOprSelectList{Selects: x.SelectList.Selects, With: x.With, Limit: x.Limit, OrderBy: x.OrderBy}}
 			}
-			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}, With: with}
+			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}}
 			setOpr.Limit = yyS[yypt-0].item.(*ast.Limit)
 			parser.yyVAL.statement = setOpr
 		}
 	case 1795:
 		{
 			var setOprList []ast.Node
-			var with *ast.WithClause
 			switch x := yyS[yypt-2].expr.(*ast.SubqueryExpr).Query.(type) {
 			case *ast.SelectStmt:
-				setOprList = []ast.Node{x}
-				with = x.With
+				setOprList = []ast.Node{&ast.SetOprSelectList{Selects: []ast.Node{x}, With: x.With}}
 			case *ast.SetOprStmt:
-				setOprList = x.SelectList.Selects
-				with = x.With
+				setOprList = []ast.Node{&ast.SetOprSelectList{Selects: x.SelectList.Selects, With: x.With, Limit: x.Limit, OrderBy: x.OrderBy}}
 			}
-			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}, With: with}
+			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}}
 			setOpr.OrderBy = yyS[yypt-1].item.(*ast.OrderByClause)
 			setOpr.Limit = yyS[yypt-0].item.(*ast.Limit)
 			parser.yyVAL.statement = setOpr
@@ -19533,7 +19536,7 @@ yynewstate:
 			case *ast.SelectStmt:
 				setOprList = []ast.Node{&ast.SetOprSelectList{Selects: []ast.Node{x}}}
 			case *ast.SetOprStmt:
-				setOprList = []ast.Node{&ast.SetOprSelectList{Selects: x.SelectList.Selects, With: x.With}}
+				setOprList = []ast.Node{&ast.SetOprSelectList{Selects: x.SelectList.Selects, With: x.With, Limit: x.Limit, OrderBy: x.OrderBy}}
 			}
 			parser.yyVAL.item = setOprList
 		}