[parser] parser: implement Restore for Limit (#100)

* parser: implement Restore for Limit * fix * change test content
pingcap · Oct 9, 2021 · 20ff6cc · 20ff6cc
1 parent 8862bb5
commit 20ff6cc
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/parser/ast/dml.go b/parser/ast/dml.go
@@ -1084,7 +1084,17 @@ type Limit struct {
 
 // Restore implements Node interface.
 func (n *Limit) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	ctx.WriteKeyWord("LIMIT ")
+	if n.Offset != nil {
+		if err := n.Offset.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore Limit.Offset")
+		}
+		ctx.WritePlain(",")
+	}
+	if err := n.Count.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore Limit.Count")
+	}
+	return nil
 }
 
 // Accept implements Node Accept interface.

diff --git a/parser/ast/dml_test.go b/parser/ast/dml_test.go
@@ -126,3 +126,16 @@ func (tc *testDMLSuite) TestTableNameIndexHintsRestore(c *C) {
 	}
 	RunNodeRestoreTest(c, testCases, "SELECT * FROM %s", extractNodeFunc)
 }
+
+func (tc *testDMLSuite) TestLimitRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"limit 10", "LIMIT 10"},
+		{"limit 10,20", "LIMIT 10,20"},
+		{"limit 20 offset 10", "LIMIT 10,20"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*SelectStmt).Limit
+	}
+	RunNodeRestoreTest(c, testCases, "SELECT 1 %s", extractNodeFunc)
+}
+