From a4bdcab31117b4b4cbfbdb792307e3332534950d Mon Sep 17 00:00:00 2001
From: exialin <exialin37@gmail.com>
Date: Fri, 28 Dec 2018 10:29:38 +0800
Subject: [PATCH] parser: implement Restore for OnCondition and TableRefsClause
 (#132)

---
 ast/dml.go      | 15 +++++++++++----
 ast/dml_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/ast/dml.go b/ast/dml.go
index eb452a994..0a2e7fde7 100644
--- a/ast/dml.go
+++ b/ast/dml.go
@@ -118,8 +118,8 @@ func (n *Join) Restore(ctx *RestoreCtx) error {
 	ctx.JoinLevel--
 
 	if n.On != nil {
-		ctx.WriteKeyWord(" ON ")
-		if err := n.On.Expr.Restore(ctx); err != nil {
+		ctx.WritePlain(" ")
+		if err := n.On.Restore(ctx); err != nil {
 			return errors.Annotate(err, "An error occurred while restore Join.On")
 		}
 	}
@@ -325,7 +325,11 @@ type OnCondition struct {
 
 // Restore implements Node interface.
 func (n *OnCondition) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	ctx.WriteKeyWord("ON ")
+	if err := n.Expr.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore OnCondition.Expr")
+	}
+	return nil
 }
 
 // Accept implements Node Accept interface.
@@ -551,7 +555,10 @@ type TableRefsClause struct {
 
 // Restore implements Node interface.
 func (n *TableRefsClause) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	if err := n.TableRefs.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore TableRefsClause.TableRefs")
+	}
+	return nil
 }
 
 // Accept implements Node Accept interface.
diff --git a/ast/dml_test.go b/ast/dml_test.go
index 3ffec4715..b0a04b669 100644
--- a/ast/dml_test.go
+++ b/ast/dml_test.go
@@ -193,6 +193,17 @@ func (tc *testDMLSuite) TestTableSourceRestore(c *C) {
 	RunNodeRestoreTest(c, testCases, "select * from %s", extractNodeFunc)
 }
 
+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`"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*SelectStmt).From.TableRefs.On
+	}
+	RunNodeRestoreTest(c, testCases, "select * from t1 join t2 %s", extractNodeFunc)
+}
+
 func (tc *testDMLSuite) TestJoinRestore(c *C) {
 	testCases := []NodeRestoreTestCase{
 		{"t1 natural join t2", "`t1` NATURAL JOIN `t2`"},
@@ -215,6 +226,19 @@ func (tc *testDMLSuite) TestJoinRestore(c *C) {
 	}
 	RunNodeRestoreTest(c, testCases, "select * from %s", extractNodeFunc)
 }
+
+func (ts *testDMLSuite) TestTableRefsClauseRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"t", "`t`"},
+		{"t1 join t2", "`t1` JOIN `t2`"},
+		{"t1, t2", "(`t1`) JOIN `t2`"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*SelectStmt).From
+	}
+	RunNodeRestoreTest(c, testCases, "select * from %s", extractNodeFunc)
+}
+
 func (tc *testDMLSuite) TestDeleteTableListRestore(c *C) {
 	testCases := []NodeRestoreTestCase{
 		{"t1,t2", "`t1`,`t2`"},
@@ -225,6 +249,7 @@ func (tc *testDMLSuite) TestDeleteTableListRestore(c *C) {
 	RunNodeRestoreTest(c, testCases, "DELETE %s FROM t1, t2;", extractNodeFunc)
 	RunNodeRestoreTest(c, testCases, "DELETE FROM %s USING t1, t2;", extractNodeFunc)
 }
+
 func (tc *testExpressionsSuite) TestByItemRestore(c *C) {
 	testCases := []NodeRestoreTestCase{
 		{"a", "`a`"},