From 0e06020118f4d99ca9c28ef3b246529111416524 Mon Sep 17 00:00:00 2001
From: Dmitrii Stepanov <d.stepanov@yadro.com>
Date: Wed, 13 Dec 2023 17:52:18 +0300
Subject: [PATCH] [#107] preset_grpc: Allow to create local containers

Signed-off-by: Dmitrii Stepanov <d.stepanov@yadro.com>
---
 scenarios/preset/helpers/frostfs_cli.py | 70 ++++++++++++++++++++++---
 scenarios/preset/preset_grpc.py         |  3 +-
 2 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/scenarios/preset/helpers/frostfs_cli.py b/scenarios/preset/helpers/frostfs_cli.py
index 5a7632d..0f3dcaa 100644
--- a/scenarios/preset/helpers/frostfs_cli.py
+++ b/scenarios/preset/helpers/frostfs_cli.py
@@ -1,11 +1,14 @@
 import re
 from helpers.cmd import execute_cmd, log
 
-def create_container(endpoint, policy, wallet_file, wallet_config):
-    if wallet_file:
-        wallet_file = "--wallet " + wallet_file
-    if wallet_config:
-        wallet_config = "--config " + wallet_config
+def create_container(endpoint, policy, wallet_path, config, local=False, depth=0):
+    if depth > 20:
+        raise ValueError(f"unable to create container: too many unsuccessful attempts")
+
+    if wallet_path:
+        wallet_file = "--wallet " + wallet_path
+    if config:
+        wallet_config = "--config " + config
     cmd_line = f"frostfs-cli --rpc-endpoint {endpoint} container create {wallet_file} {wallet_config} " \
                f" --policy '{policy}' --basic-acl public-read-write --await"
 
@@ -27,10 +30,63 @@ def create_container(endpoint, policy, wallet_file, wallet_config):
     splitted = fst_str.split(": ")
     if len(splitted) != 2:
         raise ValueError(f"no CID was parsed from command output:\t{fst_str}")
+    cid = splitted[1]
 
-    log(f"Created container {splitted[1]}", endpoint)
+    log(f"Created container {cid}", endpoint)
 
-    return splitted[1]
+    if not local:
+        return cid
+    
+    cmd_line = f"frostfs-cli netmap nodeinfo --rpc-endpoint {endpoint} {wallet_file} {wallet_config}"
+    output, success = execute_cmd(cmd_line)
+
+    if not success:
+        log(f"{cmd_line}\n"
+            f"Failed to get nodeinfo\n"
+            f"{output}", endpoint)
+        return False
+    
+    try:
+        fst_str = output.split('\n')[0]
+    except Exception:
+        log(f"{cmd_line}\n"
+            f"Incorrect output\n"
+            f"Output: {output or '<empty>'}", endpoint)
+        return False
+    splitted = fst_str.split(": ")
+    if len(splitted) != 2 or len(splitted[1]) == 0:
+        raise ValueError(f"no node key was parsed from command output:\t{fst_str}")
+    
+    node_key = splitted[1]
+
+    cmd_line = f"frostfs-cli container nodes --rpc-endpoint {endpoint} {wallet_file} {wallet_config} --cid {cid}"
+    output, success = execute_cmd(cmd_line)
+
+    if not success:
+        log(f"{cmd_line}\n"
+            f"Failed to get container nodes\n"
+            f"{output}", endpoint)
+        return False
+
+    for output_str in output.split('\n'):
+        output_str = output_str.lstrip().rstrip()
+        if not output_str.startswith("Node "):
+            continue
+        splitted = output_str.split(": ")
+        if len(splitted) != 2 or len(splitted[1]) == 0:
+            continue
+        try:
+            k = splitted[1].split(" ")[0]
+        except Exception:
+            log(f"{cmd_line}\n"
+                f"Incorrect output\n"
+                f"Output: {output or '<empty>'}", endpoint)
+            continue
+        if k == node_key:
+            return cid
+    
+    log(f"Created container {cid} is not stored on {endpoint}, creating another one...", endpoint)
+    return create_container(endpoint, policy, wallet_path, config, local, depth + 1)
 
 
 def upload_object(container, payload_filepath, endpoint, wallet_file, wallet_config):
diff --git a/scenarios/preset/preset_grpc.py b/scenarios/preset/preset_grpc.py
index 89e74f9..b636104 100755
--- a/scenarios/preset/preset_grpc.py
+++ b/scenarios/preset/preset_grpc.py
@@ -34,6 +34,7 @@ parser.add_argument('--ignore-errors', help='Ignore preset errors', action='stor
 parser.add_argument('--workers', help='Count of workers in preset. Max = 50, Default = 50', default=50)
 parser.add_argument('--sleep', help='Time to sleep between containers creation and objects upload (in seconds), '
                                     'Default = 8', default=8)
+parser.add_argument('--local', help='Create containers that store data on provided endpoints. Warning: additional empty containers may be created.', default=False)
 
 args: Namespace = parser.parse_args()
 print(args)
@@ -61,7 +62,7 @@ def main():
         containers_count = int(args.containers)
         print(f"Create containers: {containers_count}")
         with ProcessPoolExecutor(max_workers=min(MAX_WORKERS, workers)) as executor:
-            containers_runs = [executor.submit(create_container, endpoint, args.policy, wallet, wallet_config)
+            containers_runs = [executor.submit(create_container, endpoint, args.policy, wallet, wallet_config, args.local)
                                for _, endpoint in 
                                zip(range(containers_count), cycle(endpoints))]