Add per-host copy_args to parallel copy functions (#111)

Add per-host copy_args to parallel copy functions

Author: Shados

pssh/base_pssh.py

@@ -201,7 +201,7 @@ def get_exit_code(self, host_output):
             return
         return self._get_exit_code(host_output.channel)
 
-    def copy_file(self, local_file, remote_file, recurse=False):
+    def copy_file(self, local_file, remote_file, recurse=False, copy_args=None):
         """Copy local file to remote file in parallel
 
         This function returns a list of greenlets which can be
@@ -224,11 +224,19 @@ def copy_file(self, local_file, remote_file, recurse=False):
         :type remote_file: str
         :param recurse: Whether or not to descend into directories recursively.
         :type recurse: bool
+        :param copy_args: (Optional) format local_file and remote_file strings
+          with per-host arguments in ``copy_args``. ``copy_args`` length must
+          equal length of host list -
+          :py:class:`pssh.exceptions.HostArgumentException` is raised otherwise
+        :type copy_args: tuple or list
+
         :rtype: List(:py:class:`gevent.Greenlet`) of greenlets for remote copy
           commands
 
         :raises: :py:class:`ValueError` when a directory is supplied to
           local_file and recurse is not set
+        :raises: :py:class:`pssh.exceptions.HostArgumentException` on number of
+          per-host copy arguments not equal to number of hosts
         :raises: :py:class:`IOError` on I/O errors writing files
         :raises: :py:class:`OSError` on OS errors like permission denied
 
@@ -238,9 +246,21 @@ def copy_file(self, local_file, remote_file, recurse=False):
           created as long as permissions allow.
 
         """
-        return [self.pool.spawn(self._copy_file, host, local_file, remote_file,
-                                {'recurse': recurse})
-                for host in self.hosts]
+        if copy_args:
+            try:
+                return [self.pool.spawn(self._copy_file, host,
+                                        local_file % copy_args[host_i],
+                                        remote_file % copy_args[host_i],
+                                        {'recurse': recurse})
+                        for host_i, host in enumerate(self.hosts)]
+            except IndexError:
+                raise HostArgumentException(
+                    "Number of per-host copy arguments provided does not match "
+                    "number of hosts")
+        else:
+            return [self.pool.spawn(self._copy_file, host, local_file,
+                                    remote_file, {'recurse': recurse})
+                    for host in self.hosts]
 
     def _copy_file(self, host, local_file, remote_file, recurse=False):
         """Make sftp client, copy file"""
@@ -249,7 +269,7 @@ def _copy_file(self, host, local_file, remote_file, recurse=False):
                                                  recurse=recurse)
 
     def copy_remote_file(self, remote_file, local_file, recurse=False,
-                         suffix_separator='_', **kwargs):
+                         suffix_separator='_', copy_args=None, **kwargs):
         """Copy remote file(s) in parallel as
         <local_file><suffix_separator><host>
 
@@ -281,13 +301,22 @@ def copy_remote_file(self, remote_file, local_file, recurse=False,
           filename and host, defaults to ``_``. For example, for a
           ``local_file`` value of ``myfile`` and default separator the
           resulting filename will be ``myfile_myhost`` for the file from
-          host ``myhost``
+          host ``myhost``. ``suffix_separator`` has no meaning if
+          ``copy_args`` is provided
         :type suffix_separator: str
+        :param copy_args: (Optional) format remote_file and local_file strings
+          with per-host arguments in ``copy_args``.   ``copy_args`` length must
+          equal length of host list -
+          :py:class:`pssh.exceptions.HostArgumentException` is raised otherwise
+        :type copy_args: tuple or list
+
         :rtype: list(:py:class:`gevent.Greenlet`) of greenlets for remote copy
           commands
 
         :raises: :py:class:`ValueError` when a directory is supplied to
           local_file and recurse is not set
+        :raises: :py:class:`pssh.exceptions.HostArgumentException` on number of
+          per-host copy arguments not equal to number of hosts
         :raises: :py:class:`IOError` on I/O errors writing files
         :raises: :py:class:`OSError` on OS errors like permission denied
 
@@ -300,17 +329,29 @@ def copy_remote_file(self, remote_file, local_file, recurse=False,
           filepath separated by ``suffix_separator``.
 
         """
-        return [self.pool.spawn(
-            self._copy_remote_file, host, remote_file,
-            local_file, recurse, suffix_separator=suffix_separator,
-            **kwargs)
-            for host in self.hosts]
+        if copy_args:
+            try:
+                return [self.pool.spawn(
+                    self._copy_remote_file, host,
+                    remote_file % copy_args[host_i],
+                    local_file % copy_args[host_i], {'recurse': recurse},
+                    **kwargs)
+                    for host_i, host in enumerate(self.hosts)]
+            except IndexError:
+                raise HostArgumentException(
+                    "Number of per-host copy arguments provided does not match "
+                    "number of hosts")
+        else:
+            return [self.pool.spawn(
+                self._copy_remote_file, host, remote_file,
+                suffix_separator.join([local_file, host]), recurse,
+                **kwargs)
+                for host in self.hosts]
 
     def _copy_remote_file(self, host, remote_file, local_file, recurse,
-                          suffix_separator='_', **kwargs):
+                          **kwargs):
         """Make sftp client, copy file to local"""
-        file_w_suffix = suffix_separator.join([local_file, host])
         self._make_ssh_client(host)
         return self.host_clients[host].copy_remote_file(
-            remote_file, file_w_suffix, recurse=recurse,
+            remote_file, local_file, recurse=recurse,
             **kwargs)

pssh/exceptions.py

@@ -40,7 +40,7 @@ class SSHException(Exception):
 
 
 class HostArgumentException(Exception):
-    """Raised on errors with per-host command arguments"""
+    """Raised on errors with per-host arguments to parallel functions"""
     pass
 
 

pssh/pssh2_client.py

@@ -312,7 +312,7 @@ def _make_ssh_client(self, host):
                 num_retries=self.num_retries, timeout=self.timeout,
                 allow_agent=self.allow_agent, retry_delay=self.retry_delay)
 
-    def copy_file(self, local_file, remote_file, recurse=False):
+    def copy_file(self, local_file, remote_file, recurse=False, copy_args=None):
         """Copy local file to remote file in parallel
 
         This function returns a list of greenlets which can be
@@ -335,12 +335,19 @@ def copy_file(self, local_file, remote_file, recurse=False):
         :type remote_file: str
         :param recurse: Whether or not to descend into directories recursively.
         :type recurse: bool
+        :param copy_args: (Optional) format local_file and remote_file strings
+          with per-host arguments in ``copy_args``. ``copy_args`` length must
+          equal length of host list -
+          :py:class:`pssh.exceptions.HostArgumentException` is raised otherwise
+        :type copy_args: tuple or list
 
         :rtype: list(:py:class:`gevent.Greenlet`) of greenlets for remote copy
           commands
 
         :raises: :py:class:`ValueError` when a directory is supplied to
           local_file and recurse is not set
+        :raises: :py:class:`pssh.exceptions.HostArgumentException` on number of
+          per-host copy arguments not equal to number of hosts
         :raises: :py:class:`pss.exceptions.SFTPError` on SFTP initialisation
           errors
         :raises: :py:class:`pssh.exceptions.SFTPIOError` on I/O errors writing
@@ -354,10 +361,11 @@ def copy_file(self, local_file, remote_file, recurse=False):
 
         """
         return BaseParallelSSHClient.copy_file(
-            self, local_file, remote_file, recurse=recurse)
+            self, local_file, remote_file, recurse=recurse, copy_args=copy_args)
 
     def copy_remote_file(self, remote_file, local_file, recurse=False,
-                         suffix_separator='_', encoding='utf-8'):
+                         suffix_separator='_', copy_args=None,
+                         encoding='utf-8'):
         """Copy remote file(s) in parallel as
         <local_file><suffix_separator><host>
 
@@ -389,8 +397,14 @@ def copy_remote_file(self, remote_file, local_file, recurse=False,
           filename and host, defaults to ``_``. For example, for a
           ``local_file`` value of ``myfile`` and default separator the
           resulting filename will be ``myfile_myhost`` for the file from
-          host ``myhost``
+          host ``myhost``. ``suffix_separator`` has no meaning if
+          ``copy_args`` is provided
         :type suffix_separator: str
+        :param copy_args: (Optional) format remote_file and local_file strings
+          with per-host arguments in ``copy_args``.   ``copy_args`` length must
+          equal length of host list -
+          :py:class:`pssh.exceptions.HostArgumentException` is raised otherwise
+        :type copy_args: tuple or list
         :param encoding: Encoding to use for file paths.
         :type encoding: str
 
@@ -399,6 +413,8 @@ def copy_remote_file(self, remote_file, local_file, recurse=False,
 
         :raises: :py:class:`ValueError` when a directory is supplied to
           local_file and recurse is not set
+        :raises: :py:class:`pssh.exceptions.HostArgumentException` on number of
+          per-host copy arguments not equal to number of hosts
         :raises: :py:class:`pss.exceptions.SFTPError` on SFTP initialisation
           errors
         :raises: :py:class:`pssh.exceptions.SFTPIOError` on I/O errors reading
@@ -416,4 +432,5 @@ def copy_remote_file(self, remote_file, local_file, recurse=False,
         """
         return BaseParallelSSHClient.copy_remote_file(
             self, remote_file, local_file, recurse=recurse,
-            suffix_separator=suffix_separator, encoding=encoding)
+            suffix_separator=suffix_separator, copy_args=copy_args,
+            encoding=encoding)

tests/test_pssh_ssh2_client.py

@@ -370,6 +370,60 @@ def test_pssh_copy_file(self):
             except OSError:
                 pass
 
+    def test_pssh_copy_file_per_host_args(self):
+        """Test parallel copy file with per-host arguments"""
+        host2, host3 = '127.0.0.6', '127.0.0.7'
+        server2 = OpenSSHServer(host2)
+        server3 = OpenSSHServer(host3)
+        servers = [server2, server3]
+        for server in servers:
+            server.start_server()
+        time.sleep(1)
+        hosts = [self.host, host2, host3]
+
+        local_file_prefix = 'test_file_'
+        remote_file_prefix = 'test_remote_'
+
+        copy_args = [dict(zip(('local_file', 'remote_file',),
+                              (local_file_prefix + str(i + 1),
+                               remote_file_prefix + str(i + 1),)
+                              ))
+                     for i, _ in enumerate(hosts)]
+
+        test_file_data = 'test'
+        for i, _ in enumerate(hosts):
+            test_file = open(local_file_prefix + str(i + 1), 'w')
+            test_file.writelines([test_file_data + os.linesep])
+            test_file.close()
+
+        client = ParallelSSHClient(hosts, port=self.port, pkey=self.user_key,
+                                   num_retries=2)
+        greenlets = client.copy_file('%(local_file)s', '%(remote_file)s',
+                                     copy_args=copy_args)
+        gevent.joinall(greenlets)
+
+        self.assertRaises(HostArgumentException, client.copy_file,
+                          '%(local_file)s', '%(remote_file)s',
+                          copy_args=[copy_args[0]])
+        try:
+            for i, _ in enumerate(hosts):
+                remote_file_abspath = os.path.expanduser(
+                    '~/' + remote_file_prefix + str(i + 1))
+                self.assertTrue(os.path.isfile(remote_file_abspath))
+                remote_file_data = open(
+                    remote_file_abspath, 'r').readlines()
+                self.assertEqual(
+                    remote_file_data[0].strip(), test_file_data)
+        except Exception:
+            raise
+        finally:
+            for i, _ in enumerate(hosts):
+                remote_file_abspath = os.path.expanduser(
+                        '~/' + remote_file_prefix + str(i + 1))
+                local_file_path = local_file_prefix + str(i + 1)
+                os.unlink(remote_file_abspath)
+                os.unlink(local_file_path)
+
     def test_pssh_client_directory_relative_path(self):
         """Tests copying multiple directories with SSH client. Copy all the files from
         local directory to server, then make sure they are all present."""
@@ -581,6 +635,59 @@ def test_pssh_copy_remote_file(self):
             shutil.rmtree(new_local_copied_dir)
             shutil.rmtree(remote_test_path_abs)
 
+    def test_pssh_copy_remote_file_per_host_args(self):
+        """Test parallel remote copy file with per-host arguments"""
+        host2, host3 = '127.0.0.10', '127.0.0.11'
+        server2 = OpenSSHServer(host2)
+        server3 = OpenSSHServer(host3)
+        servers = [server2, server3]
+        for server in servers:
+            server.start_server()
+        time.sleep(1)
+        hosts = [self.host, host2, host3]
+
+        remote_file_prefix = 'test_file_'
+        local_file_prefix = 'test_local_'
+
+        copy_args = [dict(zip(('remote_file', 'local_file',),
+                              (remote_file_prefix + str(i + 1),
+                               local_file_prefix + str(i + 1),)
+                              ))
+                     for i, _ in enumerate(hosts)]
+
+        test_file_data = 'test'
+        for i, _ in enumerate(hosts):
+            remote_file_abspath = os.path.expanduser(
+                '~/' + remote_file_prefix + str(i + 1))
+            test_file = open(remote_file_abspath, 'w')
+            test_file.writelines([test_file_data + os.linesep])
+            test_file.close()
+
+        client = ParallelSSHClient(hosts, port=self.port, pkey=self.user_key,
+                                   num_retries=2)
+        greenlets = client.copy_remote_file('%(remote_file)s', '%(local_file)s',
+                                            copy_args=copy_args)
+        gevent.joinall(greenlets)
+
+        self.assertRaises(HostArgumentException, client.copy_remote_file,
+                          '%(remote_file)s', '%(local_file)s',
+                          copy_args=[copy_args[0]])
+        try:
+            for i, _ in enumerate(hosts):
+                local_file_path = local_file_prefix + str(i + 1)
+                self.assertTrue(os.path.isfile(local_file_path))
+                local_file_data = open(local_file_path, 'r').readlines()
+                self.assertEqual(local_file_data[0].strip(), test_file_data)
+        except Exception:
+            raise
+        finally:
+            for i, _ in enumerate(hosts):
+                remote_file_abspath = os.path.expanduser(
+                        '~/' + remote_file_prefix + str(i + 1))
+                local_file_path = local_file_prefix + str(i + 1)
+                os.unlink(remote_file_abspath)
+                os.unlink(local_file_path)
+
     def test_pssh_pool_size(self):
         """Test setting pool size to non default values"""
         hosts = ['host-%01d' % d for d in range(5)]