added ssh keepalive and pseudo terminal support to orchestrator

docs/Configuration.md

@@ -127,8 +127,10 @@ backup_grace_period_in_days = 10
 [ssh]
 ;username = <SSH username to use for restoring clusters>
 ;key_file = <SSH key for use for restoring clusters. Expected in PEM unencrypted format.>
-;port = <SSH port for use for restoring clusters. Default to port 22.
+;port = <SSH port for use for restoring clusters. Default to port 22.>
 ;cert_file = <Path of public key signed certificate file to use for authentication. The corresponding private key must also be provided via key_file parameter>
+;keepalive_seconds = <seconds between ssh keepalive messages to the ssh server. Default to 60 seconds. Due to a limitation in parallel-ssh, if 'cert_file' is defined, then 'keepalive_seconds' will be ignored and no keep alive messages will be sent>
+;use_pty = <Boolean: Allocates pseudo-terminal. Default to False. Useful if sudo settings require a tty>
 
 [checks]
 ;health_check = <Which ports to check when verifying a node restored properly. Options are 'cql' (default), 'thrift', 'all'.>

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'] = {

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)
 

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__':