From 047e0606e41a2ae78029c87db0aecb71a0e577ad Mon Sep 17 00:00:00 2001 From: alberto-bortolan Date: Mon, 16 Sep 2024 16:53:04 +0100 Subject: [PATCH] added ssh keepalive and pseudo terminal support to orchestrator --- docs/Configuration.md | 4 +++- medusa/config.py | 6 +++-- medusa/orchestration.py | 48 +++++++++++++++++++++++++------------ tests/orchestration_test.py | 10 ++++---- 4 files changed, 46 insertions(+), 22 deletions(-) diff --git a/docs/Configuration.md b/docs/Configuration.md index 85b8f7d39..c80643f73 100644 --- a/docs/Configuration.md +++ b/docs/Configuration.md @@ -127,8 +127,10 @@ backup_grace_period_in_days = 10 [ssh] ;username = ;key_file = -;port = ;cert_file = +;keepalive_seconds = +;use_pty = [checks] ;health_check = diff --git a/medusa/config.py b/medusa/config.py index 90ab35338..7aae58073 100644 --- a/medusa/config.py +++ b/medusa/config.py @@ -45,7 +45,7 @@ SSHConfig = collections.namedtuple( 'SSHConfig', - ['username', 'key_file', 'port', 'cert_file'] + ['username', 'key_file', 'port', 'cert_file', 'use_pty', 'keepalive_seconds'] ) ChecksConfig = collections.namedtuple( @@ -146,7 +146,9 @@ def _build_default_config(): 'username': os.environ.get('USER') or '', 'key_file': '', 'port': '22', - 'cert_file': '' + 'cert_file': '', + 'use_pty': 'False', + 'keepalive_seconds': '60' } config['checks'] = { diff --git a/medusa/orchestration.py b/medusa/orchestration.py index 86a1b2dc4..37a72f065 100644 --- a/medusa/orchestration.py +++ b/medusa/orchestration.py @@ -15,7 +15,8 @@ import logging -from pssh.clients.ssh import ParallelSSHClient +from pssh.clients.native.parallel import ParallelSSHClient as PsshNativeClient +from pssh.clients.ssh.parallel import ParallelSSHClient as PsshSSHClient import medusa.utils from medusa.storage import divide_chunks @@ -39,33 +40,50 @@ def pssh_run(self, hosts, command, hosts_variables=None, ssh_client=None): Runs a command on hosts list using pssh under the hood Return: True (success) or False (error) """ + username = self.config.ssh.username if self.config.ssh.username != '' else None + port = int(self.config.ssh.port) + pkey = self.config.ssh.key_file if self.config.ssh.key_file != '' else None + cert_file = self.config.ssh.cert_file if self.config.ssh.cert_file != '' else None + keepalive_seconds = int(self.config.ssh.keepalive_seconds) + use_pty = medusa.utils.evaluate_boolean(self.config.ssh.use_pty) + if ssh_client is None: - ssh_client = ParallelSSHClient + if cert_file is None: + ssh_client = PsshNativeClient + else: + ssh_client = PsshSSHClient pssh_run_success = False success = [] error = [] i = 1 - username = self.config.ssh.username if self.config.ssh.username != '' else None - port = int(self.config.ssh.port) - pkey = self.config.ssh.key_file if self.config.ssh.key_file != '' else None - cert_file = self.config.ssh.cert_file if self.config.ssh.cert_file != '' else None - logging.info('Executing "{command}" on following nodes {hosts} with a parallelism/pool size of {pool_size}' .format(command=command, hosts=hosts, pool_size=self.pool_size)) for parallel_hosts in divide_chunks(hosts, self.pool_size): - client = ssh_client(parallel_hosts, - forward_ssh_agent=True, - pool_size=len(parallel_hosts), - user=username, - port=port, - pkey=pkey, - cert_file=cert_file) + if cert_file is None: + client = ssh_client(parallel_hosts, + forward_ssh_agent=True, + pool_size=len(parallel_hosts), + user=username, + port=port, + pkey=pkey, + keepalive_seconds=keepalive_seconds) + else: + logging.debug('The ssh parameter "cert_file" is defined. Due to limitations in parallel-ssh ' + '"keep_alive" will be ignored and no ServerAlive messages will be generated') + client = ssh_client(parallel_hosts, + forward_ssh_agent=True, + pool_size=len(parallel_hosts), + user=username, + port=port, + pkey=pkey, + cert_file=cert_file) + logging.debug('Batch #{i}: Running "{command}" on nodes {hosts} parallelism of {pool_size}' .format(i=i, command=command, hosts=parallel_hosts, pool_size=len(parallel_hosts))) - output = client.run_command(command, host_args=hosts_variables, + output = client.run_command(command, host_args=hosts_variables, use_pty=use_pty, sudo=medusa.utils.evaluate_boolean(self.config.cassandra.use_sudo)) client.join(output) diff --git a/tests/orchestration_test.py b/tests/orchestration_test.py index 3c3d1a6f9..913c87711 100644 --- a/tests/orchestration_test.py +++ b/tests/orchestration_test.py @@ -67,7 +67,9 @@ def _build_config_parser(): 'username': '', 'key_file': '', 'port': '22', - 'cert_file': '' + 'cert_file': '', + 'keepalive_seconds': '60', + 'use_pty': 'False' } return config @@ -91,7 +93,7 @@ def test_pssh_with_sudo(self): self.mock_pssh.run_command.return_value = output assert self.orchestration.pssh_run(list(self.hosts.keys()), 'fake command', ssh_client=self.fake_ssh_client_factory) - self.mock_pssh.run_command.assert_called_with('fake command', host_args=None, sudo=True) + self.mock_pssh.run_command.assert_called_with('fake command', host_args=None, use_pty=False, sudo=True) def test_pssh_without_sudo(self): """Ensure that Parallel SSH honors configuration when we don't want to use sudo in commands""" @@ -105,7 +107,7 @@ def test_pssh_without_sudo(self): assert orchestration_no_sudo.pssh_run(list(self.hosts.keys()), 'fake command', ssh_client=self.fake_ssh_client_factory) - self.mock_pssh.run_command.assert_called_with('fake command', host_args=None, sudo=False) + self.mock_pssh.run_command.assert_called_with('fake command', host_args=None, use_pty=False, sudo=False) def test_pssh_run_failure(self): """Ensure that Parallel SSH detects a failed command on a host""" @@ -118,7 +120,7 @@ def test_pssh_run_failure(self): self.mock_pssh.run_command.return_value = output assert not self.orchestration.pssh_run(list(self.hosts.keys()), 'fake command', ssh_client=self.fake_ssh_client_factory) - self.mock_pssh.run_command.assert_called_with('fake command', host_args=None, sudo=True) + self.mock_pssh.run_command.assert_called_with('fake command', host_args=None, use_pty=False, sudo=True) if __name__ == '__main__':