Balance stop (vesoft-inc#1238)

* balance stop * address dangleptr's comments * update log, fix ut
yixinglu · Nov 12, 2019 · e3a9d90 · e3a9d90
1 parent d539a9f
commit e3a9d90
Show file tree

Hide file tree

Showing 6 changed files with 33 additions and 4 deletions.
diff --git a/src/executor/BalanceExecutor.cpp b/src/executor/BalanceExecutor.cpp
@@ -27,6 +27,9 @@ void BalanceExecutor::execute() {
         case BalanceSentence::SubType::kData:
             balanceData();
             break;
+        case BalanceSentence::SubType::kDataStop:
+            balanceData(true);
+            break;
         case BalanceSentence::SubType::kShowBalancePlan:
             showBalancePlan();
             break;
@@ -66,8 +69,8 @@ void BalanceExecutor::balanceLeader() {
     std::move(future).via(runner).thenValue(cb).thenError(error);
 }
 
-void BalanceExecutor::balanceData() {
-    auto future = ectx()->getMetaClient()->balance();
+void BalanceExecutor::balanceData(bool isStop) {
+    auto future = ectx()->getMetaClient()->balance(isStop);
     auto *runner = ectx()->rctx()->runner();
 
     auto cb = [this] (auto &&resp) {

diff --git a/src/executor/BalanceExecutor.h b/src/executor/BalanceExecutor.h
@@ -27,7 +27,9 @@ class BalanceExecutor final : public Executor {
 
     void balanceLeader();
 
-    void balanceData();
+    void balanceData(bool isStop = false);
+
+    void stopBalanceData();
 
     void showBalancePlan();
 

diff --git a/src/parser/AdminSentences.h b/src/parser/AdminSentences.h
@@ -309,6 +309,7 @@ class BalanceSentence final : public Sentence {
         kUnknown,
         kLeader,
         kData,
+        kDataStop,
         kShowBalancePlan,
     };
 

diff --git a/src/parser/parser.yy b/src/parser/parser.yy
@@ -109,7 +109,7 @@ class GraphScanner;
 %token KW_COUNT KW_COUNT_DISTINCT KW_SUM KW_AVG KW_MAX KW_MIN KW_STD KW_BIT_AND KW_BIT_OR KW_BIT_XOR
 %token KW_FETCH KW_PROP KW_UPDATE KW_UPSERT KW_WHEN
 %token KW_DISTINCT KW_ALL KW_OF
-%token KW_BALANCE KW_LEADER KW_DATA
+%token KW_BALANCE KW_LEADER KW_DATA KW_STOP
 %token KW_SHORTEST KW_PATH
 
 /* symbols */
@@ -1680,6 +1680,9 @@ balance_sentence
     | KW_BALANCE KW_DATA INTEGER {
         $$ = new BalanceSentence($3);
     }
+    | KW_BALANCE KW_DATA KW_STOP {
+        $$ = new BalanceSentence(BalanceSentence::SubType::kDataStop);
+    }
     ;
 
 mutate_sentence

diff --git a/src/parser/scanner.lex b/src/parser/scanner.lex
@@ -123,6 +123,7 @@ LEADER                      ([Ll][Ee][Aa][Dd][Ee][Rr])
 UUID                        ([Uu][Uu][Ii][Dd])
 OF                          ([Oo][Ff])
 DATA                        ([Dd][Aa][Tt][Aa])
+STOP                        ([Ss][Tt][Oo][Pp])
 SHORTEST                    ([Ss][Hh][Oo][Rr][Tt][Ee][Ss][Tt])
 PATH                        ([Pp][Aa][Tt][Hh])
 LIMIT                       ([Ll][Ii][Mm][Ii][Tt])
@@ -250,6 +251,7 @@ IP_OCTET                    ([0-9]|[1-9][0-9]|1[0-9][0-9]|2[0-4][0-9]|25[0-5])
 {LEADER}                    { return TokenType::KW_LEADER; }
 {UUID}                      { return TokenType::KW_UUID; }
 {DATA}                      { return TokenType::KW_DATA; }
+{STOP}                      { return TokenType::KW_STOP; }
 {SHORTEST}                  { return TokenType::KW_SHORTEST; }
 {PATH}                      { return TokenType::KW_PATH; }
 {LIMIT}                     { return TokenType::KW_LIMIT; }

diff --git a/src/parser/test/ParserTest.cpp b/src/parser/test/ParserTest.cpp
@@ -1370,6 +1370,24 @@ TEST(Parser, BalanceOperation) {
         auto result = parser.parse(query);
         ASSERT_TRUE(result.ok()) << result.status();
     }
+    {
+        GQLParser parser;
+        std::string query = "BALANCE DATA";
+        auto result = parser.parse(query);
+        ASSERT_TRUE(result.ok()) << result.status();
+    }
+    {
+        GQLParser parser;
+        std::string query = "BALANCE DATA 1234567890";
+        auto result = parser.parse(query);
+        ASSERT_TRUE(result.ok()) << result.status();
+    }
+    {
+        GQLParser parser;
+        std::string query = "BALANCE DATA STOP";
+        auto result = parser.parse(query);
+        ASSERT_TRUE(result.ok()) << result.status();
+    }
 }
 
 TEST(Parser, CrashByFuzzer) {