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`"},