use custom scripts for download binaries in ydbd_slice (#16566)

2 files changed, 17 insertions, 1 deletions

diff --git a/ydb/tools/ydbd_slice/__init__.py b/ydb/tools/ydbd_slice/__init__.py
index bb701808f77..56e294d7ed1 100644
--- a/ydb/tools/ydbd_slice/__init__.py
+++ b/ydb/tools/ydbd_slice/__init__.py
@@ -471,7 +471,7 @@ def binaries_args():
         "--kikimr",
         metavar="BIN",
         default=None,
-        help="explicit path to ydbd. Can be url: 'rbtorrent:<torrent>' for rbtorrent, 'sbr:<id>' for sandbox resource or 'http(s)://<url>' for other."
+        help="explicit path to ydbd. Can be url: 'rbtorrent:<torrent>' for rbtorrent, 'sbr:<id>' for sandbox resource or 'http(s)://<url>' for http or 'script:' for custom script."
     )
     args.add_argument(
         "--binary-lz4",
diff --git a/ydb/tools/ydbd_slice/nodes.py b/ydb/tools/ydbd_slice/nodes.py
index 20b3e09f6f6..21f65637832 100644
--- a/ydb/tools/ydbd_slice/nodes.py
+++ b/ydb/tools/ydbd_slice/nodes.py
@@ -3,6 +3,7 @@ import sys
 import logging
 import subprocess
 import queue
+import random
 
 
 logger = logging.getLogger(__name__)
@@ -178,6 +179,19 @@ class Nodes(object):
         running_jobs = self.execute_async_ret(script)
         self._check_async_execution(running_jobs, retry_attemps=2)
 
+    def _download_script(self, script, remote_path):
+        user_script = script[len('script:'):]
+        self._logger.info(f"download by script '{user_script}' to '{remote_path}'")
+        tmp_path = f'tmp_{random.randint(0, 100500)}'
+        full_script = (
+            f'mkdir -p {tmp_path} && cd {tmp_path} && '
+            f'( {user_script} ) && cd - && '
+            f'for FILE in `find {tmp_path} -name *.tgz -or -name *.tar`; do tar -C {tmp_path} -xf $FILE && rm $FILE; done && '
+            f'sudo mv {tmp_path}/* {remote_path} && rm -rf {tmp_path}'
+        )
+        running_jobs = self.execute_async_ret(full_script)
+        self._check_async_execution(running_jobs, retry_attemps=2)
+
     def _download_http(self, url, remote_path):
         self._logger.info(f"download from '{url}' to '{remote_path}'")
         running_jobs = self.execute_async_ret(f'sudo curl --output {remote_path} {url}')
@@ -215,6 +229,8 @@ class Nodes(object):
         self.execute_async("sudo mkdir -p {}".format(os.path.dirname(remote_path)))
         if local_path.startswith('rbtorrent:') or local_path.startswith('sbr:'):
             self._download_sky(local_path, remote_path)
+        elif local_path.startswith('script:'):
+            self._download_script(local_path, remote_path)
         elif local_path.startswith('http:') or local_path.startswith('https:'):
             self._download_http(local_path, remote_path)
         else: