Skip to content

Commit

Permalink
Add ability to specify user in ydbd_slice (ydb-platform#3819)
Browse files Browse the repository at this point in the history
  • Loading branch information
@maximyurchuk
maximyurchuk authored Apr 17, 2024
1 parent 62e55d6 commit 3b65df0
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 21 deletions.
35 changes: 24 additions & 11 deletions ydb/tools/ydbd_slice/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -299,7 +299,7 @@ def deduce_components_from_args(args, cluster_details):
return result


def deduce_nodes_from_args(args, walle_provider):
def deduce_nodes_from_args(args, walle_provider, ssh_user):
cluster_hosts = safe_load_cluster_details(args.cluster, walle_provider).hosts_names
result = cluster_hosts

Expand All @@ -315,7 +315,7 @@ def deduce_nodes_from_args(args, walle_provider):
sys.exit("unable to deduce hosts")

logger.info("use nodes '%s'", result)
return nodes.Nodes(result, args.dry_run)
return nodes.Nodes(result, args.dry_run, ssh_user=ssh_user)


def ya_build(arcadia_root, artifact, opts, dry_run):
Expand Down Expand Up @@ -491,6 +491,19 @@ def component_args():
return args


def ssh_args():
current_user = os.environ["USER"]
args = argparse.ArgumentParser(add_help=False)
args.add_argument(
"--ssh-user",
metavar="SSH_USER",
default=[current_user],
help="user for ssh interaction with slice. Default value is $USER "
"(which equals {user} now)".format(user=current_user),
)
return args


def add_explain_mode(modes, walle_provider):
def _run(args):
logger.debug("run func explain with cmd args is '%s'", args)
Expand Down Expand Up @@ -539,7 +552,7 @@ def dispatch_run(func, args, walle_provider):
cluster_details = safe_load_cluster_details(args.cluster, walle_provider)
components = deduce_components_from_args(args, cluster_details)

nodes = deduce_nodes_from_args(args, walle_provider)
nodes = deduce_nodes_from_args(args, walle_provider, args.ssh_user)

temp_dir = deduce_temp_dir_from_args(args)
clear_tmp = not args.dry_run and args.temp_dir is None
Expand All @@ -564,7 +577,7 @@ def dispatch_run(func, args, walle_provider):
configurator,
clear_logs,
yav_version,
walle_provider
walle_provider,
)
func(slice)

Expand All @@ -580,7 +593,7 @@ def _run(args):
mode = modes.add_parser(
"install",
conflict_handler='resolve',
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args(), log_args()],
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args(), log_args(), ssh_args()],
description="Full installation of the cluster from scratch. "
"You can use --hosts to specify particular hosts. But it is tricky."
)
Expand All @@ -594,7 +607,7 @@ def _run(args):
mode = modes.add_parser(
"update",
conflict_handler='resolve',
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args(), log_args()],
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args(), log_args(), ssh_args()],
description="Minor cluster update, just binary and cfg. No additional configuration is performed."
"Stop all kikimr instances at the nodes, sync binary and cfg, start the instances. "
"Use --hosts to specify particular hosts."
Expand All @@ -610,7 +623,7 @@ def _run(args):
mode = modes.add_parser(
"update-raw-cfg",
conflict_handler='resolve',
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args()],
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args(), ssh_args()],
description=""
)
mode.add_argument(
Expand All @@ -628,7 +641,7 @@ def _run(args):

mode = modes.add_parser(
"stop",
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args()],
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args(), ssh_args()],
description="Stop kikimr static instaneces at the nodes. "
"If option components specified, try to stop particular component. "
"Use --hosts to specify particular hosts."
Expand All @@ -642,7 +655,7 @@ def _run(args):

mode = modes.add_parser(
"start",
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args()],
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args(), ssh_args()],
description="Start all kikimr instances at the nodes. "
"If option components specified, try to start particular component. "
"Otherwise only kikimr-multi-all will be started. "
Expand All @@ -657,7 +670,7 @@ def _run(args):

mode = modes.add_parser(
"clear",
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args()],
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args(), ssh_args()],
description="Stop all kikimr instances at the nodes, format all kikimr drivers, shutdown dynamic slots. "
"And don't start nodes afrer it. "
"Use --hosts to specify particular hosts."
Expand All @@ -671,7 +684,7 @@ def _run(args):

mode = modes.add_parser(
"format",
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args()],
parents=[direct_nodes_args(), cluster_description_args(), binaries_args(), component_args(), ssh_args()],
description="Stop all kikimr instances at the nodes, format all kikimr drivers at the nodes, start the instances. "
"If you call format for all cluster, you will spoil it. "
"Additional dynamic configuration will required after it. "
Expand Down
32 changes: 22 additions & 10 deletions ydb/tools/ydbd_slice/nodes.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,21 +8,27 @@


class Nodes(object):
def __init__(self, nodes, dry_run=False):
def __init__(self, nodes, dry_run=False, ssh_user=None):
assert isinstance(nodes, list)
assert len(nodes) > 0
assert isinstance(nodes[0], str)
self._nodes = nodes
self._dry_run = bool(dry_run)
self._ssh_user = ssh_user
self._logger = logger.getChild(self.__class__.__name__)

@property
def nodes_list(self):
return self._nodes

@staticmethod
def _wrap_ssh_cmd(cmd, host):
return ['ssh', '-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null', '-A', host, cmd]
def _get_ssh_command_prefix(self, host):
command = []
command.extend(['ssh', '-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null', '-A'])
if (self._ssh_user):
command.extend(['-l', self._ssh_user])

command.extend([host])
return command

def _check_async_execution(self, running_jobs, check_retcode=True, results=None):
if self._dry_run:
Expand Down Expand Up @@ -76,7 +82,7 @@ def execute_async_ret(self, cmd, check_retcode=True, nodes=None, results=None):
if self._dry_run:
continue

actual_cmd = self._wrap_ssh_cmd(cmd, host)
actual_cmd = self._get_ssh_command_prefix(host) + [cmd]
process = subprocess.Popen(actual_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
running_jobs.append((actual_cmd, process, host))
return running_jobs
Expand All @@ -96,6 +102,8 @@ def _copy_on_node(self, local_path, host, remote_path):
if self._dry_run:
return
destination = "{host}:{path}".format(host=host, path=remote_path)
if self._ssh_user:
destination = self._ssh_user + "@" + destination
subprocess.check_call(["rsync", "-avqLW", "--del", "--no-o", "--no-g", "--rsync-path=sudo rsync", "--progress",
local_path, destination])

Expand All @@ -105,6 +113,8 @@ def _copy_between_nodes(self, hub, hub_path, hosts, remote_path):
assert isinstance(hosts, list)

src = "{hub}:{hub_path}".format(hub=hub, hub_path=hub_path)
if self._ssh_user:
src = self._ssh_user + "@" + src
running_jobs = []
for dst in hosts:
self._logger.info(
Expand All @@ -117,11 +127,13 @@ def _copy_between_nodes(self, hub, hub_path, hosts, remote_path):
)
if self._dry_run:
continue
cmd = [
"ssh", dst, '-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null', "-A", "sudo", "rsync", "-avqW", "--del",
"--no-o", "--no-g", "--rsh='ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -l %s'" % os.getenv("USER"),
src, remote_path,
]
user = self._ssh_user or os.getenv("USER")
cmd = self._get_ssh_command_prefix
cmd.extend([
"sudo", "rsync", "-avqW", "--del", "--no-o", "--no-g",
"--rsh='ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -l %s'" % user,
src, remote_path
])
process = subprocess.Popen(cmd)
running_jobs.append((cmd, process, dst))

Expand Down

0 comments on commit 3b65df0

Please sign in to comment.