From a55df2870329ce28b1eae8c7c92e6d634e8f3e31 Mon Sep 17 00:00:00 2001
From: 924060929 <924060929@qq.com>
Date: Thu, 3 Jan 2019 10:29:36 +0800
Subject: [PATCH] [parser] parser: implement Restore for Assignment (#146)

---
 parser/ast/dml.go      |  9 ++++++++-
 parser/ast/dml_test.go | 11 +++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/parser/ast/dml.go b/parser/ast/dml.go
index 70deb7886a2c1..17f415bf0ebdf 100644
--- a/parser/ast/dml.go
+++ b/parser/ast/dml.go
@@ -932,7 +932,14 @@ type Assignment struct {
 
 // Restore implements Node interface.
 func (n *Assignment) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	if err := n.Column.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore Assignment.Column")
+	}
+	ctx.WritePlain("=")
+	if err := n.Expr.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore Assignment.Expr")
+	}
+	return nil
 }
 
 // Accept implements Node Accept interface.
diff --git a/parser/ast/dml_test.go b/parser/ast/dml_test.go
index 27e85818fb774..9dfb802559aa2 100644
--- a/parser/ast/dml_test.go
+++ b/parser/ast/dml_test.go
@@ -289,6 +289,17 @@ func (tc *testDMLSuite) TestOrderByClauseRestore(c *C) {
 	RunNodeRestoreTest(c, testCases, "SELECT 1 FROM t1 UNION SELECT 2 FROM t2 %s", extractNodeFromUnionStmtFunc)
 }
 
+func (tc *testDMLSuite) TestAssignmentRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"a=1", "`a`=1"},
+		{"b=1+2", "`b`=1+2"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*UpdateStmt).List[0]
+	}
+	RunNodeRestoreTest(c, testCases, "UPDATE t1 SET %s", extractNodeFunc)
+}
+
 func (ts *testDMLSuite) TestHavingClauseRestore(c *C) {
 	testCases := []NodeRestoreTestCase{
 		{"HAVING a", "HAVING `a`"},