Merge pull request ClickHouse#55485 from guoxiaolongzte/keeper_client…

…_add_get_all_children_number_command

keeper_client add get_all_children_number command

docs/en/operations/utilities/clickhouse-keeper-client.md

@@ -55,6 +55,7 @@ keeper foo bar
 -   `rmr <path>` -- Recursively deletes path. Confirmation required
 -   `flwc <command>` -- Executes four-letter-word command
 -   `help` -- Prints this message
+-   `get_all_children_number [path]` -- Get all numbers of children nodes under a specific path
 -   `get_stat [path]` -- Returns the node's stat (default `.`)
 -   `find_super_nodes <threshold> [path]` -- Finds nodes with number of children larger than some threshold for the given path (default `.`)
 -   `delete_stale_backups` -- Deletes ClickHouse nodes used for backups that are now inactive

programs/keeper-client/Commands.cpp

@@ -324,8 +324,8 @@ void FindBigFamily::execute(const ASTKeeperQuery * query, KeeperClient * client)
         queue.pop();
 
         auto children = client->zookeeper->getChildren(next_path);
-        std::transform(children.cbegin(), children.cend(), children.begin(), [&](const String & child) { return next_path / child; });
-
+        for (auto & child : children)
+            child = next_path / child;
         auto response = client->zookeeper->get(children);
 
         for (size_t i = 0; i < response.size(); ++i)
@@ -475,4 +475,45 @@ void FourLetterWordCommand::execute(const ASTKeeperQuery * query, KeeperClient *
     std::cout << client->executeFourLetterCommand(query->args[0].safeGet<String>()) << "\n";
 }
 
+bool GetAllChildrenNumberCommand::parse(IParser::Pos & pos, std::shared_ptr<ASTKeeperQuery> & node, Expected & expected) const
+{
+    String path;
+    if (!parseKeeperPath(pos, expected, path))
+        path = ".";
+
+    node->args.push_back(std::move(path));
+
+    return true;
+}
+
+void GetAllChildrenNumberCommand::execute(const ASTKeeperQuery * query, KeeperClient * client) const
+{
+    auto path = client->getAbsolutePath(query->args[0].safeGet<String>());
+
+    std::queue<fs::path> queue;
+    queue.push(path);
+    Coordination::Stat stat;
+    client->zookeeper->get(path, &stat);
+
+    int totalNumChildren = stat.numChildren;
+    while (!queue.empty())
+    {
+        auto next_path = queue.front();
+        queue.pop();
+
+        auto children = client->zookeeper->getChildren(next_path);
+        for (auto & child : children)
+            child = next_path / child;
+        auto response = client->zookeeper->get(children);
+
+        for (size_t i = 0; i < response.size(); ++i)
+        {
+            totalNumChildren += response[i].stat.numChildren;
+            queue.push(children[i]);
+        }
+    }
+
+    std::cout << totalNumChildren << "\n";
+}
+
 }

programs/keeper-client/Commands.h

@@ -238,4 +238,18 @@ class FourLetterWordCommand : public IKeeperClientCommand
     String getHelpMessage() const override { return "{} <command> -- Executes four-letter-word command"; }
 };
 
+class GetAllChildrenNumberCommand : public IKeeperClientCommand
+{
+    String getName() const override { return "get_all_children_number"; }
+
+    bool parse(IParser::Pos & pos, std::shared_ptr<ASTKeeperQuery> & node, Expected & expected) const override;
+
+    void execute(const ASTKeeperQuery * query, KeeperClient * client) const override;
+
+    String getHelpMessage() const override
+    {
+        return "{} [path] -- Get all numbers of children nodes under a specific path";
+    }
+};
+
 }

programs/keeper-client/KeeperClient.cpp

@@ -2,6 +2,7 @@
 #include "Commands.h"
 #include <Client/ReplxxLineReader.h>
 #include <Client/ClientBase.h>
+#include "Common/VersionNumber.h"
 #include <Common/Config/ConfigProcessor.h>
 #include <Common/EventNotifier.h>
 #include <Common/filesystemHelpers.h>
@@ -206,6 +207,7 @@ void KeeperClient::initialize(Poco::Util::Application & /* self */)
         std::make_shared<SyncCommand>(),
         std::make_shared<HelpCommand>(),
         std::make_shared<FourLetterWordCommand>(),
+        std::make_shared<GetAllChildrenNumberCommand>(),
     });
 
     String home_path;

tests/integration/helpers/keeper_utils.py

@@ -143,6 +143,9 @@ def find_big_family(self, path: str, n: int = 10, timeout: float = 60.0) -> str:
     def find_super_nodes(self, threshold: int, timeout: float = 60.0) -> str:
         return self.execute_query(f"find_super_nodes {threshold}", timeout)
 
+    def get_all_children_number(self, path: str, timeout: float = 60.0) -> str:
+        return self.execute_query(f"get_all_children_number {path}", timeout)
+
     def delete_stale_backups(self, timeout: float = 60.0) -> str:
         return self.execute_query("delete_stale_backups", timeout)
 

tests/integration/test_keeper_client/test.py

@@ -216,3 +216,20 @@ def test_quoted_argument_parsing(client: KeeperClient):
 
     client.execute_query(f"set '{node_path}' \"value4 with some whitespace\" 3")
     assert client.get(node_path) == "value4 with some whitespace"
+
+
+def test_get_all_children_number(client: KeeperClient):
+    client.touch("/test_get_all_children_number")
+    client.touch("/test_get_all_children_number/1")
+    client.touch("/test_get_all_children_number/1/1")
+    client.touch("/test_get_all_children_number/1/2")
+    client.touch("/test_get_all_children_number/1/3")
+    client.touch("/test_get_all_children_number/1/4")
+    client.touch("/test_get_all_children_number/1/5")
+    client.touch("/test_get_all_children_number/2")
+    client.touch("/test_get_all_children_number/2/1")
+    client.touch("/test_get_all_children_number/2/2")
+    client.touch("/test_get_all_children_number/2/3")
+    client.touch("/test_get_all_children_number/2/4")
+
+    assert client.get_all_children_number("/test_get_all_children_number") == "11"