diff --git a/pkg/parser/ast/dml.go b/pkg/parser/ast/dml.go index 8306d09688c28..a1da415611e39 100644 --- a/pkg/parser/ast/dml.go +++ b/pkg/parser/ast/dml.go @@ -1569,6 +1569,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: @@ -1590,6 +1591,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 } @@ -1614,6 +1629,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/pkg/parser/parser.go b/pkg/parser/parser.go index ee851f07448ed..d61073ea8a173 100644 --- a/pkg/parser/parser.go +++ b/pkg/parser/parser.go @@ -19908,6 +19908,8 @@ 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} @@ -19915,8 +19917,10 @@ yynewstate: 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}} @@ -19932,6 +19936,8 @@ 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} @@ -19939,8 +19945,10 @@ yynewstate: 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}} @@ -19956,6 +19964,8 @@ 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} @@ -19963,8 +19973,10 @@ yynewstate: 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}} @@ -19975,48 +19987,39 @@ yynewstate: case 1839: { 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}}} 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 1840: { 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 1841: { 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 @@ -20048,7 +20051,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 index e58ce2eb13e08..bb3b3f04282fa 100644 --- a/pkg/parser/parser.y +++ b/pkg/parser/parser.y @@ -10274,6 +10274,8 @@ SetOprStmtWithLimitOrderBy: } var setOprList2 []ast.Node var with2 *ast.WithClause + var limit2 *ast.Limit + var orderBy2 *ast.OrderByClause switch x := $3.(*ast.SubqueryExpr).Query.(type) { case *ast.SelectStmt: setOprList2 = []ast.Node{x} @@ -10281,8 +10283,10 @@ SetOprStmtWithLimitOrderBy: 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 = $2.(*ast.SetOprType) setOprList := append(setOprList1, nextSetOprList) setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}} @@ -10298,6 +10302,8 @@ SetOprStmtWithLimitOrderBy: } var setOprList2 []ast.Node var with2 *ast.WithClause + var limit2 *ast.Limit + var orderBy2 *ast.OrderByClause switch x := $3.(*ast.SubqueryExpr).Query.(type) { case *ast.SelectStmt: setOprList2 = []ast.Node{x} @@ -10305,8 +10311,10 @@ SetOprStmtWithLimitOrderBy: 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 = $2.(*ast.SetOprType) setOprList := append(setOprList1, nextSetOprList) setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}} @@ -10322,6 +10330,8 @@ SetOprStmtWithLimitOrderBy: } var setOprList2 []ast.Node var with2 *ast.WithClause + var limit2 *ast.Limit + var orderBy2 *ast.OrderByClause switch x := $3.(*ast.SubqueryExpr).Query.(type) { case *ast.SelectStmt: setOprList2 = []ast.Node{x} @@ -10329,8 +10339,10 @@ SetOprStmtWithLimitOrderBy: 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 = $2.(*ast.SetOprType) setOprList := append(setOprList1, nextSetOprList) setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}} @@ -10341,48 +10353,39 @@ SetOprStmtWithLimitOrderBy: | SubSelect OrderBy { var setOprList []ast.Node - var with *ast.WithClause switch x := $1.(*ast.SubqueryExpr).Query.(type) { case *ast.SelectStmt: - setOprList = []ast.Node{x} - with = x.With + setOprList = []ast.Node{&ast.SetOprSelectList{Selects: []ast.Node{x}}} 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 = $2.(*ast.OrderByClause) $$ = setOpr } | SubSelect SelectStmtLimit { var setOprList []ast.Node - var with *ast.WithClause switch x := $1.(*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 = $2.(*ast.Limit) $$ = setOpr } | SubSelect OrderBy SelectStmtLimit { var setOprList []ast.Node - var with *ast.WithClause switch x := $1.(*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 = $2.(*ast.OrderByClause) setOpr.Limit = $3.(*ast.Limit) $$ = setOpr @@ -10419,7 +10422,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 index e698b42f55fd6..5315282408493 100644 --- a/pkg/parser/parser_test.go +++ b/pkg/parser/parser_test.go @@ -5147,6 +5147,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) }