Skip to content

Commit

Permalink
This is an automated cherry-pick of pingcap#49898
Browse files Browse the repository at this point in the history
Signed-off-by: ti-chi-bot <ti-community-prow-bot@tidb.io>
  • Loading branch information
Defined2014 authored and ti-chi-bot committed Dec 29, 2023
1 parent 721bb1a commit 46ee554
Show file tree
Hide file tree
Showing 4 changed files with 24,952 additions and 22 deletions.
29 changes: 29 additions & 0 deletions parser/ast/dml.go
Original file line number Diff line number Diff line change
Expand Up @@ -1507,6 +1507,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:
Expand All @@ -1528,6 +1529,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
}

Expand All @@ -1552,6 +1567,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)
}

Expand Down
47 changes: 25 additions & 22 deletions parser/parser.y
Original file line number Diff line number Diff line change
Expand Up @@ -9341,15 +9341,19 @@ 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}
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 = $2.(*ast.SetOprType)
setOprList := append(setOprList1, nextSetOprList)
setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}}
Expand All @@ -9365,15 +9369,19 @@ 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}
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 = $2.(*ast.SetOprType)
setOprList := append(setOprList1, nextSetOprList)
setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}}
Expand All @@ -9389,15 +9397,19 @@ 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}
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 = $2.(*ast.SetOprType)
setOprList := append(setOprList1, nextSetOprList)
setOpr := &ast.SetOprStmt{SelectList: &ast.SetOprSelectList{Selects: setOprList}}
Expand All @@ -9408,48 +9420,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
Expand Down Expand Up @@ -9486,7 +9489,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
}
Expand Down
10 changes: 10 additions & 0 deletions parser/parser_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -4460,6 +4460,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)
}
Expand Down
Loading

0 comments on commit 46ee554

Please sign in to comment.