Add get_direct_children_number command to keeper-client (#55898)

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_direct_children_number [path]` -- Get numbers of direct children nodes under a specific path
 -   `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 `.`)

programs/keeper-client/Commands.cpp

@@ -475,6 +475,27 @@ void FourLetterWordCommand::execute(const ASTKeeperQuery * query, KeeperClient *
     std::cout << client->executeFourLetterCommand(query->args[0].safeGet<String>()) << "\n";
 }
 
+bool GetDirectChildrenNumberCommand::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 GetDirectChildrenNumberCommand::execute(const ASTKeeperQuery * query, KeeperClient * client) const
+{
+    auto path = client->getAbsolutePath(query->args[0].safeGet<String>());
+
+    Coordination::Stat stat;
+    client->zookeeper->get(path, &stat);
+
+    std::cout << stat.numChildren << "\n";
+}
+
 bool GetAllChildrenNumberCommand::parse(IParser::Pos & pos, std::shared_ptr<ASTKeeperQuery> & node, Expected & expected) const
 {
     String path;

programs/keeper-client/Commands.h

@@ -238,6 +238,20 @@ class FourLetterWordCommand : public IKeeperClientCommand
     String getHelpMessage() const override { return "{} <command> -- Executes four-letter-word command"; }
 };
 
+class GetDirectChildrenNumberCommand : public IKeeperClientCommand
+{
+    String getName() const override { return "get_direct_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 numbers of direct children nodes under a specific path";
+    }
+};
+
 class GetAllChildrenNumberCommand : public IKeeperClientCommand
 {
     String getName() const override { return "get_all_children_number"; }

programs/keeper-client/KeeperClient.cpp

@@ -207,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<GetDirectChildrenNumberCommand>(),
         std::make_shared<GetAllChildrenNumberCommand>(),
     });
 

tests/integration/helpers/keeper_utils.py

@@ -143,6 +143,9 @@ class KeeperClient(object):
     def find_super_nodes(self, threshold: int, timeout: float = 60.0) -> str:
         return self.execute_query(f"find_super_nodes {threshold}", timeout)
 
+    def get_direct_children_number(self, path: str, timeout: float = 60.0) -> str:
+        return self.execute_query(f"get_direct_children_number {path}", 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)
 

tests/integration/test_keeper_client/test.py

@@ -218,6 +218,18 @@ def test_quoted_argument_parsing(client: KeeperClient):
     assert client.get(node_path) == "value4 with some whitespace"
 
 
+def get_direct_children_number(client: KeeperClient):
+    client.touch("/get_direct_children_number")
+    client.touch("/get_direct_children_number/1")
+    client.touch("/get_direct_children_number/1/1")
+    client.touch("/get_direct_children_number/1/2")
+    client.touch("/get_direct_children_number/2")
+    client.touch("/get_direct_children_number/2/1")
+    client.touch("/get_direct_children_number/2/2")
+
+    assert client.get_direct_children_number("/get_direct_children_number") == "2"
+
+
 def test_get_all_children_number(client: KeeperClient):
     client.touch("/test_get_all_children_number")
     client.touch("/test_get_all_children_number/1")