add more test cases

pingcap · Dec 29, 2023 · 502e421 · 502e421
1 parent c3b4ccd
commit 502e421
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 8 deletions.
diff --git a/pkg/parser/ast/dml.go b/pkg/parser/ast/dml.go
@@ -1554,6 +1554,7 @@ func (n *SelectStmt) Accept(v Visitor) (Node, bool) {
 type SetOprSelectList struct {
 	node
 
+	IsInBraces       bool
 	With             *WithClause
 	AfterSetOperator *SetOprType
 	Selects          []Node
@@ -1570,6 +1571,13 @@ func (n *SetOprSelectList) Restore(ctx *format.RestoreCtx) error {
 		}
 	}
 
+	if n.IsInBraces {
+		ctx.WritePlain("(")
+		defer func() {
+			ctx.WritePlain(")")
+		}()
+	}
+
 	for i, stmt := range n.Selects {
 		switch selectStmt := stmt.(type) {
 		case *SelectStmt:

diff --git a/pkg/parser/parser.go b/pkg/parser/parser.go
@@ -20023,7 +20023,7 @@ yynewstate:
 				limit = x.Limit
 				orderBy = x.OrderBy
 			}
-			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}, With: with, Limit: limit, OrderBy: orderBy}
+			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList, With: with, Limit: limit, OrderBy: orderBy, IsInBraces: true}}
 			setOpr.OrderBy = yyS[yypt-0].item.(*ast.OrderByClause)
 			parser.yyVAL.statement = setOpr
 		}
@@ -20043,7 +20043,7 @@ yynewstate:
 				limit = x.Limit
 				orderBy = x.OrderBy
 			}
-			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}, With: with, Limit: limit, OrderBy: orderBy}
+			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList, With: with, Limit: limit, OrderBy: orderBy, IsInBraces: true}}
 			setOpr.Limit = yyS[yypt-0].item.(*ast.Limit)
 			parser.yyVAL.statement = setOpr
 		}
@@ -20063,7 +20063,7 @@ yynewstate:
 				limit = x.Limit
 				orderBy = x.OrderBy
 			}
-			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}, With: with, Limit: limit, OrderBy: orderBy}
+			setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList, With: with, Limit: limit, OrderBy: orderBy, IsInBraces: true}}
 			setOpr.OrderBy = yyS[yypt-1].item.(*ast.OrderByClause)
 			setOpr.Limit = yyS[yypt-0].item.(*ast.Limit)
 			parser.yyVAL.statement = setOpr
@@ -20095,7 +20095,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
 		}

diff --git a/pkg/parser/parser.y b/pkg/parser/parser.y
@@ -10352,7 +10352,7 @@ SetOprStmtWithLimitOrderBy:
 			limit = x.Limit
 			orderBy = x.OrderBy
 		}
-		setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}, With: with, Limit: limit, OrderBy: orderBy}
+		setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList, With: with, Limit: limit, OrderBy: orderBy, IsInBraces: true}}
 		setOpr.OrderBy = $2.(*ast.OrderByClause)
 		$$ = setOpr
 	}
@@ -10372,7 +10372,7 @@ SetOprStmtWithLimitOrderBy:
 			limit = x.Limit
 			orderBy = x.OrderBy
 		}
-		setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}, With: with, Limit: limit, OrderBy: orderBy}
+		setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList, With: with, Limit: limit, OrderBy: orderBy, IsInBraces: true}}
 		setOpr.Limit = $2.(*ast.Limit)
 		$$ = setOpr
 	}
@@ -10392,7 +10392,7 @@ SetOprStmtWithLimitOrderBy:
 			limit = x.Limit
 			orderBy = x.OrderBy
 		}
-		setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}, With: with, Limit: limit, OrderBy: orderBy}
+		setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList, With: with, Limit: limit, OrderBy: orderBy, IsInBraces: true}}
 		setOpr.OrderBy = $2.(*ast.OrderByClause)
 		setOpr.Limit = $3.(*ast.Limit)
 		$$ = setOpr
@@ -10429,7 +10429,7 @@ SetOprClause:
 		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}}
 		}
 		$$ = setOprList
 	}

diff --git a/pkg/parser/parser_test.go b/pkg/parser/parser_test.go
@@ -5155,10 +5155,16 @@ func TestSetOperator(t *testing.T) {
 		{"(select c1 from t1) union all (select c2 from t2 except select c3 from t3) order by c1 limit 1", true, "(SELECT `c1` FROM `t1`) UNION ALL (SELECT `c2` FROM `t2` EXCEPT SELECT `c3` FROM `t3`) ORDER BY `c1` LIMIT 1"},
 		{"((select c1 from t1) except select c2 from t2) intersect all (select c3 from t3) order by c1 limit 1", true, "((SELECT `c1` FROM `t1`) EXCEPT SELECT `c2` FROM `t2`) INTERSECT ALL (SELECT `c3` FROM `t3`) ORDER BY `c1` LIMIT 1"},
 		{"select 1 union distinct (select 1 except all select 1 intersect select 1)", true, "SELECT 1 UNION (SELECT 1 EXCEPT ALL SELECT 1 INTERSECT SELECT 1)"},
+
+		// https://github.com/pingcap/tidb/issues/49874
 		{"select * from a where PK = 0 union all (select * from b where PK = 0 union all (select * from b where PK != 0) order by pk limit 1)", true,
 			"SELECT * FROM `a` WHERE `PK`=0 UNION ALL (SELECT * FROM `b` WHERE `PK`=0 UNION ALL (SELECT * FROM `b` WHERE `PK`!=0) ORDER BY `pk` LIMIT 1)"},
 		{"select * from a where PK = 0 union all (select * from b where PK = 0 union all (select * from b where PK != 0) order by pk limit 1) order by pk limit 2", true,
 			"SELECT * FROM `a` WHERE `PK`=0 UNION ALL (SELECT * FROM `b` WHERE `PK`=0 UNION ALL (SELECT * FROM `b` WHERE `PK`!=0) ORDER BY `pk` LIMIT 1) ORDER BY `pk` LIMIT 2"},
+		{"(select * from b where pk= 0 union all (select * from b where pk !=0) order by pk limit 1) order by pk limit 2", true,
+			"(SELECT * FROM `b` WHERE `pk`=0 UNION ALL (SELECT * FROM `b` WHERE `pk`!=0) ORDER BY `pk` LIMIT 1) ORDER BY `pk` LIMIT 2"},
+		{"(select * from b where pk= 0 union all (select * from b where pk !=0) order by pk limit 1) order by pk", true,
+			"(SELECT * FROM `b` WHERE `pk`=0 UNION ALL (SELECT * FROM `b` WHERE `pk`!=0) ORDER BY `pk` LIMIT 1) ORDER BY `pk`"},
 	}
 	RunTest(t, table, false)
 }