Support update pump or drainer status (#243)

ast/misc.go

@@ -61,6 +61,8 @@ const (
 	// Valid formats for explain statement.
 	ExplainFormatROW = "row"
 	ExplainFormatDOT = "dot"
+	PumpType         = "PUMP"
+	DrainerType      = "DRAINER"
 )
 
 var (
@@ -721,6 +723,41 @@ func (n *SetPwdStmt) Accept(v Visitor) (Node, bool) {
 	return v.Leave(n)
 }
 
+type ChangeStmt struct {
+	stmtNode
+
+	NodeType string
+	State    string
+	NodeID   string
+}
+
+// Restore implements Node interface.
+func (n *ChangeStmt) Restore(ctx *RestoreCtx) error {
+	ctx.WriteKeyWord("CHANGE ")
+	ctx.WriteKeyWord(n.NodeType)
+	ctx.WriteKeyWord(" TO NODE_STATE ")
+	ctx.WritePlain("=")
+	ctx.WriteString(n.State)
+	ctx.WriteKeyWord(" FOR NODE_ID ")
+	ctx.WriteString(n.NodeID)
+	return nil
+}
+
+// SecureText implements SensitiveStatement interface.
+func (n *ChangeStmt) SecureText() string {
+	return fmt.Sprintf("change %s to node_state='%s' for node_id '%s'", strings.ToLower(n.NodeType), n.State, n.NodeID)
+}
+
+// Accept implements Node Accept interface.
+func (n *ChangeStmt) Accept(v Visitor) (Node, bool) {
+	newNode, skipChildren := v.Enter(n)
+	if skipChildren {
+		return v.Leave(newNode)
+	}
+	n = newNode.(*ChangeStmt)
+	return v.Leave(n)
+}
+
 // UserSpec is used for parsing create user statement.
 type UserSpec struct {
 	User    *auth.UserIdentity

ast/misc_test.go

@@ -127,6 +127,20 @@ load data infile '/tmp/t.csv' into table t fields terminated by 'ab' enclosed by
 	}
 }
 
+// test Change Pump or drainer status sql parser
+func (ts *testMiscSuite) TestChangeStmt(c *C) {
+	sql := `change pump to node_state='paused' for node_id '127.0.0.1:8249';
+change drainer to node_state='paused' for node_id '127.0.0.1:8249';`
+
+	p := parser.New()
+	stmts, _, err := p.Parse(sql, "", "")
+	c.Assert(err, IsNil)
+	for _, stmt := range stmts {
+		stmt.Accept(visitor{})
+		stmt.Accept(visitor1{})
+	}
+}
+
 func (ts *testMiscSuite) TestSensitiveStatement(c *C) {
 	positive := []StmtNode{
 		&SetPwdStmt{},
@@ -203,3 +217,13 @@ func (ts *testMiscSuite) TestTableOptimizerHintRestore(c *C) {
 	}
 	RunNodeRestoreTest(c, testCases, "select /*+ %s */ * from t1 join t2", extractNodeFunc)
 }
+func (ts *testMiscSuite) TestChangeStmtRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"CHANGE PUMP TO NODE_STATE ='paused' FOR NODE_ID '127.0.0.1:9090'", "CHANGE PUMP TO NODE_STATE ='paused' FOR NODE_ID '127.0.0.1:9090'"},
+		{"CHANGE DRAINER TO NODE_STATE ='paused' FOR NODE_ID '127.0.0.1:9090'", "CHANGE DRAINER TO NODE_STATE ='paused' FOR NODE_ID '127.0.0.1:9090'"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*ChangeStmt)
+	}
+	RunNodeRestoreTest(c, testCases, "%s", extractNodeFunc)
+}

misc.go

@@ -364,6 +364,8 @@ var tokenMap = map[string]int{
 	"NEXT_ROW_ID":              next_row_id,
 	"NO":                       no,
 	"NO_WRITE_TO_BINLOG":       noWriteToBinLog,
+	"NODE_ID":                  node_id,
+	"NODE_STATE":               node_state,
 	"NONE":                     none,
 	"NOT":                      not,
 	"NOW":                      now,