[parser] parser: implement Restore for ColumnPosition (#110)

pingcap · Oct 9, 2021 · 8c816f7 · 8c816f7
1 parent ef194cc
commit 8c816f7
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/parser/ast/ddl.go b/parser/ast/ddl.go
@@ -955,7 +955,20 @@ type ColumnPosition struct {
 
 // Restore implements Node interface.
 func (n *ColumnPosition) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	switch n.Tp {
+	case ColumnPositionNone:
+		// do nothing
+	case ColumnPositionFirst:
+		ctx.WriteKeyWord("FIRST")
+	case ColumnPositionAfter:
+		ctx.WriteKeyWord("AFTER ")
+		if err := n.RelativeColumn.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore ColumnPosition.RelativeColumn")
+		}
+	default:
+		return errors.Errorf("invalid ColumnPositionType: %d", n.Tp)
+	}
+	return nil
 }
 
 // Accept implements Node Accept interface.

diff --git a/parser/ast/ddl_test.go b/parser/ast/ddl_test.go
@@ -154,3 +154,15 @@ func (ts *testDDLSuite) TestDDLTruncateTableStmtRestore(c *C) {
 	}
 	RunNodeRestoreTest(c, testCases, "%s", extractNodeFunc)
 }
+
+func (ts *testDDLSuite) TestColumnPositionRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"", ""},
+		{"first", "FIRST"},
+		{"after b", "AFTER `b`"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*AlterTableStmt).Specs[0].Position
+	}
+	RunNodeRestoreTest(c, testCases, "alter table t add column a varchar(255) %s", extractNodeFunc)
+}