diff --git a/pkg/parser/ast/dml.go b/pkg/parser/ast/dml.go index c13b836b18331..fad988a29eb6a 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 76a1f7134f5be..fb5434872ee67 100644 --- a/pkg/parser/parser.go +++ b/pkg/parser/parser.go @@ -19931,6 +19931,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} @@ -19938,8 +19940,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}} @@ -19955,6 +19959,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} @@ -19962,8 +19968,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}} @@ -19979,6 +19987,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} @@ -19986,8 +19996,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}} @@ -19998,48 +20010,39 @@ yynewstate: case 1842: { 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 1843: { 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 1844: { 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 @@ -20071,7 +20074,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 ef5bebda7edf1..5a5ec90db8fc9 100644 --- a/pkg/parser/parser.y +++ b/pkg/parser/parser.y @@ -10260,6 +10260,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} @@ -10267,8 +10269,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}} @@ -10284,6 +10288,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} @@ -10291,8 +10297,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}} @@ -10308,6 +10316,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} @@ -10315,8 +10325,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}} @@ -10327,48 +10339,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}, 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 } | 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 @@ -10405,7 +10408,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 a1ef0c55127ec..dadc5e77f121b 100644 --- a/pkg/parser/parser_test.go +++ b/pkg/parser/parser_test.go @@ -5155,6 +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) }