parser: implement Restore for SelectStmt (pingcap#153)

ast/dml.go

@@ -756,7 +756,110 @@ type SelectStmt struct {
 
 // Restore implements Node interface.
 func (n *SelectStmt) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	ctx.WriteKeyWord("SELECT ")
+
+	if n.SelectStmtOpts.Priority > 0 {
+		ctx.WriteKeyWord(mysql.Priority2Str[n.SelectStmtOpts.Priority])
+		ctx.WritePlain(" ")
+	}
+
+	if !n.SelectStmtOpts.SQLCache {
+		ctx.WriteKeyWord("SQL_NO_CACHE ")
+	}
+
+	if n.TableHints != nil && len(n.TableHints) != 0 {
+		ctx.WritePlain("/*+ ")
+		for i, tableHint := range n.TableHints {
+			if err := tableHint.Restore(ctx); err != nil {
+				errors.Annotatef(err, "An error occurred while restore SelectStmt.TableHints[%d]", i)
+			}
+		}
+		ctx.WritePlain("*/ ")
+	}
+
+	if n.Distinct {
+		ctx.WriteKeyWord("DISTINCT ")
+	}
+	if n.SelectStmtOpts.StraightJoin {
+		ctx.WriteKeyWord("STRAIGHT_JOIN ")
+	}
+	if n.Fields != nil {
+		for i, field := range n.Fields.Fields {
+			if i != 0 {
+				ctx.WritePlain(",")
+			}
+			if err := field.Restore(ctx); err != nil {
+				errors.Annotatef(err, "An error occurred while restore SelectStmt.Fields[%d]", i)
+			}
+		}
+	}
+
+	if n.From != nil {
+		ctx.WriteKeyWord(" FROM ")
+		if err := n.From.Restore(ctx); err != nil {
+			errors.Annotate(err, "An error occurred while restore SelectStmt.From")
+		}
+	}
+
+	if n.From == nil && n.Where != nil {
+		ctx.WriteKeyWord(" FROM DUAL")
+	}
+	if n.Where != nil {
+		ctx.WriteKeyWord(" WHERE ")
+		if err := n.Where.Restore(ctx); err != nil {
+			errors.Annotate(err, "An error occurred while restore SelectStmt.Where")
+		}
+	}
+
+	if n.GroupBy != nil {
+		ctx.WritePlain(" ")
+		if err := n.GroupBy.Restore(ctx); err != nil {
+			errors.Annotate(err, "An error occurred while restore SelectStmt.GroupBy")
+		}
+	}
+
+	if n.Having != nil {
+		ctx.WritePlain(" ")
+		if err := n.Having.Restore(ctx); err != nil {
+			errors.Annotate(err, "An error occurred while restore SelectStmt.Having")
+		}
+	}
+
+	if n.WindowSpecs != nil {
+		ctx.WriteKeyWord(" WINDOW ")
+		for i, windowsSpec := range n.WindowSpecs {
+			if i != 0 {
+				ctx.WritePlain(",")
+			}
+			if err := windowsSpec.Restore(ctx); err != nil {
+				errors.Annotatef(err, "An error occurred while restore SelectStmt.WindowSpec[%d]", i)
+			}
+		}
+	}
+
+	if n.OrderBy != nil {
+		ctx.WritePlain(" ")
+		if err := n.OrderBy.Restore(ctx); err != nil {
+			errors.Annotate(err, "An error occurred while restore SelectStmt.OrderBy")
+		}
+	}
+
+	if n.Limit != nil {
+		ctx.WritePlain(" ")
+		if err := n.Limit.Restore(ctx); err != nil {
+			errors.Annotate(err, "An error occurred while restore SelectStmt.Limit")
+		}
+	}
+
+	switch n.LockTp {
+	case SelectLockInShareMode:
+		ctx.WriteKeyWord(" LOCK ")
+		ctx.WriteKeyWord(n.LockTp.String())
+	case SelectLockForUpdate:
+		ctx.WritePlain(" ")
+		ctx.WriteKeyWord(n.LockTp.String())
+	}
+	return nil
 }
 
 // Accept implements Node Accept interface.
@@ -855,7 +958,24 @@ type UnionSelectList struct {
 
 // Restore implements Node interface.
 func (n *UnionSelectList) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	for i, selectStmt := range n.Selects {
+		if i != 0 {
+			ctx.WriteKeyWord(" UNION ")
+			if !selectStmt.IsAfterUnionDistinct {
+				ctx.WriteKeyWord("ALL ")
+			}
+		}
+		if selectStmt.IsInBraces {
+			ctx.WritePlain("(")
+		}
+		if err := selectStmt.Restore(ctx); err != nil {
+			errors.Annotate(err, "An error occurred while restore UnionSelectList.SelectStmt")
+		}
+		if selectStmt.IsInBraces {
+			ctx.WritePlain(")")
+		}
+	}
+	return nil
 }
 
 // Accept implements Node Accept interface.
@@ -888,7 +1008,24 @@ type UnionStmt struct {
 
 // Restore implements Node interface.
 func (n *UnionStmt) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	if err := n.SelectList.Restore(ctx); err != nil {
+		errors.Annotate(err, "An error occurred while restore UnionStmt.SelectList")
+	}
+
+	if n.OrderBy != nil {
+		ctx.WritePlain(" ")
+		if err := n.OrderBy.Restore(ctx); err != nil {
+			errors.Annotate(err, "An error occurred while restore UnionStmt.OrderBy")
+		}
+	}
+
+	if n.Limit != nil {
+		ctx.WritePlain(" ")
+		if err := n.Limit.Restore(ctx); err != nil {
+			errors.Annotate(err, "An error occurred while restore UnionStmt.Limit")
+		}
+	}
+	return nil
 }
 
 // Accept implements Node Accept interface.

ast/dml_test.go

@@ -183,9 +183,8 @@ func (tc *testDMLSuite) TestTableSourceRestore(c *C) {
 	testCases := []NodeRestoreTestCase{
 		{"tbl", "`tbl`"},
 		{"tbl as t", "`tbl` AS `t`"},
-		// TODO: Once `Restore` of SelectStmt or UnionStmt is implemented, add the following test cases
-		// {"(select * from tbl) as t", "(SELECT * FROM `tbl`) AS `t`"},
-		// {"(select * from a union select * from b) as t", "(SELECT * FROM `a` UNION SELECT * FROM `b`) AS `t`"},
+		{"(select * from tbl) as t", "(SELECT * FROM `tbl`) AS `t`"},
+		{"(select * from a union select * from b) as t", "(SELECT * FROM `a` UNION SELECT * FROM `b`) AS `t`"},
 	}
 	extractNodeFunc := func(node Node) Node {
 		return node.(*SelectStmt).From.TableRefs.Left
@@ -196,7 +195,7 @@ func (tc *testDMLSuite) TestTableSourceRestore(c *C) {
 func (tc *testDMLSuite) TestOnConditionRestore(c *C) {
 	testCases := []NodeRestoreTestCase{
 		{"on t1.a=t2.a", "ON `t1`.`a`=`t2`.`a`"},
-		{"on t1.a=t2.a and t1.b=t2.b", "ON `t1`.`a`=`t2`.`a`&&`t1`.`b`=`t2`.`b`"},
+		{"on t1.a=t2.a and t1.b=t2.b", "ON `t1`.`a`=`t2`.`a` AND `t1`.`b`=`t2`.`b`"},
 	}
 	extractNodeFunc := func(node Node) Node {
 		return node.(*SelectStmt).From.TableRefs.On

ast/expressions.go

@@ -347,7 +347,12 @@ type SubqueryExpr struct {
 
 // Restore implements Node interface.
 func (n *SubqueryExpr) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	ctx.WritePlain("(")
+	if err := n.Query.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore SubqueryExpr.Query")
+	}
+	ctx.WritePlain(")")
+	return nil
 }
 
 // Format the ExprNode into a Writer.
@@ -388,7 +393,21 @@ type CompareSubqueryExpr struct {
 
 // Restore implements Node interface.
 func (n *CompareSubqueryExpr) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	if err := n.L.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore CompareSubqueryExpr.L")
+	}
+	if err := n.Op.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore CompareSubqueryExpr.Op")
+	}
+	if n.All {
+		ctx.WriteKeyWord("ALL ")
+	} else {
+		ctx.WriteKeyWord("ANY ")
+	}
+	if err := n.R.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore CompareSubqueryExpr.R")
+	}
+	return nil
 }
 
 // Format the ExprNode into a Writer.
@@ -817,6 +836,12 @@ func (n *PatternLikeExpr) Restore(ctx *RestoreCtx) error {
 		return errors.Annotate(err, "An error occurred while restore PatternLikeExpr.Pattern")
 	}
 
+	escape := string(n.Escape)
+	if escape != "\\" {
+		ctx.WriteKeyWord(" ESCAPE ")
+		ctx.WriteString(escape)
+
+	}
 	return nil
 }
 

ast/functions.go

@@ -16,6 +16,7 @@ package ast
 import (
 	"fmt"
 	"io"
+	"strings"
 
 	"github.com/pingcap/errors"
 	. "github.com/pingcap/parser/format"
@@ -340,15 +341,7 @@ func (n *FuncCallExpr) Restore(ctx *RestoreCtx) error {
 		}
 		ctx.WriteKeyWord(" USING ")
 		ctx.WriteKeyWord(n.Args[1].GetType().Charset)
-	case "adddate":
-		if err := n.Args[0].Restore(ctx); err != nil {
-			return errors.Annotatef(err, "An error occurred while restore FuncCallExpr")
-		}
-		ctx.WritePlain(", ")
-		if err := n.Args[1].Restore(ctx); err != nil {
-			return errors.Annotatef(err, "An error occurred while restore FuncCallExpr")
-		}
-	case "date_add":
+	case "adddate", "subdate", "date_add", "date_sub":
 		if err := n.Args[0].Restore(ctx); err != nil {
 			return errors.Annotatef(err, "An error occurred while restore FuncCallExpr")
 		}
@@ -413,6 +406,15 @@ func (n *FuncCallExpr) Restore(ctx *RestoreCtx) error {
 				return errors.Annotatef(err, "An error occurred while restore FuncCallExpr")
 			}
 		}
+	case "timestampdiff", "timestampadd":
+		ctx.WriteKeyWord(n.Args[0].(ValueExpr).GetString())
+		for i := 1; i < len(n.Args); {
+			ctx.WritePlain(", ")
+			if err := n.Args[i].Restore(ctx); err != nil {
+				return errors.Annotatef(err, "An error occurred while restore FuncCallExpr")
+			}
+			i++
+		}
 	default:
 		for i, argv := range n.Args {
 			if i != 0 {
@@ -644,12 +646,28 @@ func (n *AggregateFuncExpr) Restore(ctx *RestoreCtx) error {
 	if n.Distinct {
 		ctx.WriteKeyWord("DISTINCT ")
 	}
-	for i, argv := range n.Args {
-		if i != 0 {
-			ctx.WritePlain(", ")
+	switch strings.ToLower(n.F) {
+	case "group_concat":
+		for i := 0; i < len(n.Args)-1; i++ {
+			if i != 0 {
+				ctx.WritePlain(", ")
+			}
+			if err := n.Args[i].Restore(ctx); err != nil {
+				return errors.Annotatef(err, "An error occurred while restore AggregateFuncExpr.Args[%d]", i)
+			}
+		}
+		ctx.WriteKeyWord(" SEPARATOR ")
+		if err := n.Args[len(n.Args)-1].Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore AggregateFuncExpr.Args SEPARATOR")
 		}
-		if err := argv.Restore(ctx); err != nil {
-			return errors.Annotatef(err, "An error occurred while restore AggregateFuncExpr.Args %d", i)
+	default:
+		for i, argv := range n.Args {
+			if i != 0 {
+				ctx.WritePlain(", ")
+			}
+			if err := argv.Restore(ctx); err != nil {
+				return errors.Annotatef(err, "An error occurred while restore AggregateFuncExpr.Args[%d]", i)
+			}
 		}
 	}
 	ctx.WritePlain(")")

ast/functions_test.go

@@ -47,7 +47,7 @@ func (ts *testFunctionsSuite) TestFuncCallExprRestore(c *C) {
 		{"CONV('a',16,2)", "CONV('a', 16, 2)"},
 		{"COS(PI())", "COS(PI())"},
 		{"RAND()", "RAND()"},
-		{"ADDDATE('2000-01-01', 1)", "ADDDATE('2000-01-01', 1)"},
+		{"ADDDATE('2000-01-01', 1)", "ADDDATE('2000-01-01', INTERVAL 1 DAY)"},
 		{"DATE_ADD('2000-01-01', INTERVAL 1 DAY)", "DATE_ADD('2000-01-01', INTERVAL 1 DAY)"},
 		{"DATE_ADD('2000-01-01', INTERVAL '1 1:12:23.100000' DAY_MICROSECOND)", "DATE_ADD('2000-01-01', INTERVAL '1 1:12:23.100000' DAY_MICROSECOND)"},
 		{"EXTRACT(DAY FROM '2000-01-01')", "EXTRACT(DAY FROM '2000-01-01')"},

ast/misc.go

@@ -1045,7 +1045,20 @@ type TableOptimizerHint struct {
 
 // Restore implements Node interface.
 func (n *TableOptimizerHint) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	ctx.WriteKeyWord(n.HintName.String())
+	ctx.WritePlain("(")
+	if n.HintName.L == "max_execution_time" {
+		ctx.WritePlainf("%d", n.MaxExecutionTime)
+	} else {
+		for i, table := range n.Tables {
+			if i != 0 {
+				ctx.WritePlain(", ")
+			}
+			ctx.WriteName(table.String())
+		}
+	}
+	ctx.WritePlain(")")
+	return nil
 }
 
 // Accept implements Node Accept interface.

ast/misc_test.go

@@ -188,3 +188,18 @@ func (ts *testMiscSuite) TestUserSpec(c *C) {
 	c.Assert(ok, IsTrue)
 	c.Assert(pwd, Equals, "")
 }
+
+func (ts *testMiscSuite) TestTableOptimizerHintRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"TIDB_SMJ(`t1`)", "TIDB_SMJ(`t1`)"},
+		{"TIDB_SMJ(t1)", "TIDB_SMJ(`t1`)"},
+		{"TIDB_SMJ(t1,t2)", "TIDB_SMJ(`t1`, `t2`)"},
+		{"TIDB_INLJ(t1,t2)", "TIDB_INLJ(`t1`, `t2`)"},
+		{"TIDB_HJ(t1,t2)", "TIDB_HJ(`t1`, `t2`)"},
+		{"MAX_EXECUTION_TIME(3000)", "MAX_EXECUTION_TIME(3000)"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*SelectStmt).TableHints[0]
+	}
+	RunNodeRestoreTest(c, testCases, "select /*+ %s */ * from t1 join t2", extractNodeFunc)
+}

ast/util_test.go

@@ -74,6 +74,10 @@ func (checker *nodeTextCleaner) Enter(in Node) (out Node, skipChildren bool) {
 		}
 	case *AggregateFuncExpr:
 		node.F = strings.ToLower(node.F)
+	case *FieldList:
+		for _, f := range node.Fields {
+			f.Offset = 0
+		}
 	case *AlterTableSpec:
 		for _, opt := range node.Options {
 			opt.StrValue = strings.ToLower(opt.StrValue)

opcode/opcode.go

@@ -103,9 +103,9 @@ func (o Op) String() string {
 }
 
 var opsLiteral = map[Op]string{
-	LogicAnd:   "&&",
-	LogicOr:    "||",
-	LogicXor:   "^",
+	LogicAnd:   " AND ",
+	LogicOr:    " OR ",
+	LogicXor:   " XOR ",
 	LeftShift:  "<<",
 	RightShift: ">>",
 	GE:         ">=",