Pass through remote temp filename with file up/download commands

This means the temporary filename is generated at operation generation
stage rather than execution, meaning any in-deploy config overrides are
correctly used.

Author: Fizzadar

pyinfra/api/command.py

@@ -127,10 +127,11 @@ def execute(self, state, host, executor_kwargs):
 
 
 class FileUploadCommand(PyinfraCommand):
-    def __init__(self, src, dest, **kwargs):
+    def __init__(self, src, dest, remote_temp_filename=None, **kwargs):
         super(FileUploadCommand, self).__init__(**kwargs)
         self.src = src
         self.dest = dest
+        self.remote_temp_filename = remote_temp_filename
 
     def __repr__(self):
         return 'FileUploadCommand({0}, {1})'.format(self.src, self.dest)
@@ -140,17 +141,19 @@ def execute(self, state, host, executor_kwargs):
 
         return host.put_file(
             self.src, self.dest,
+            remote_temp_filename=self.remote_temp_filename,
             print_output=state.print_output,
             print_input=state.print_input,
             **executor_kwargs
         )
 
 
 class FileDownloadCommand(PyinfraCommand):
-    def __init__(self, src, dest, **kwargs):
+    def __init__(self, src, dest, remote_temp_filename=None, **kwargs):
         super(FileDownloadCommand, self).__init__(**kwargs)
         self.src = src
         self.dest = dest
+        self.remote_temp_filename = remote_temp_filename
 
     def __repr__(self):
         return 'FileDownloadCommand({0}, {1})'.format(self.src, self.dest)
@@ -160,6 +163,7 @@ def execute(self, state, host, executor_kwargs):
 
         return host.get_file(
             self.src, self.dest,
+            remote_temp_filename=self.remote_temp_filename,
             print_output=state.print_output,
             print_input=state.print_input,
             **executor_kwargs

pyinfra/api/connectors/chroot.py

@@ -93,6 +93,7 @@ def put_file(
     host,
     filename_or_io,
     remote_filename,
+    remote_temp_filename=None,  # ignored
     print_output=False,
     print_input=False,
     **kwargs  # ignored (sudo/etc)
@@ -146,6 +147,7 @@ def get_file(
     host,
     remote_filename,
     filename_or_io,
+    remote_temp_filename=None,  # ignored
     print_output=False,
     print_input=False,
     **kwargs  # ignored (sudo/etc)

pyinfra/api/connectors/docker.py

@@ -121,6 +121,7 @@ def run_shell_command(
 
 def put_file(
     state, host, filename_or_io, remote_filename,
+    remote_temp_filename=None,  # ignored
     print_output=False, print_input=False,
     **kwargs  # ignored (sudo/etc)
 ):
@@ -171,6 +172,7 @@ def put_file(
 
 def get_file(
     state, host, remote_filename, filename_or_io,
+    remote_temp_filename=None,  # ignored
     print_output=False, print_input=False,
     **kwargs  # ignored (sudo/etc)
 ):

pyinfra/api/connectors/dockerssh.py

@@ -138,6 +138,7 @@ def run_shell_command(
 
 def put_file(
     state, host, filename_or_io, remote_filename,
+    remote_temp_filename=None,
     print_output=False, print_input=False,
     **kwargs  # ignored (sudo/etc)
 ):
@@ -146,12 +147,12 @@ def put_file(
     temporary location and then uploading it into the container using ``docker cp``.
     '''
 
-    fd, temp_filename = mkstemp()
-    remote_temp_filename = state.get_temp_filename(temp_filename)
+    fd, local_temp_filename = mkstemp()
+    remote_temp_filename = remote_temp_filename or state.get_temp_filename(local_temp_filename)
 
     # Load our file or IO object and write it to the temporary file
     with get_file_io(filename_or_io) as file_io:
-        with open(temp_filename, 'wb') as temp_f:
+        with open(local_temp_filename, 'wb') as temp_f:
             data = file_io.read()
 
             if isinstance(data, six.text_type):
@@ -160,7 +161,7 @@ def put_file(
             temp_f.write(data)
 
     # upload file to remote server
-    ssh_status = ssh.put_file(state, host, temp_filename, remote_temp_filename)
+    ssh_status = ssh.put_file(state, host, local_temp_filename, remote_temp_filename)
     if not ssh_status:
         raise IOError('Failed to copy file over ssh')
 
@@ -180,9 +181,9 @@ def put_file(
         )
     finally:
         os.close(fd)
-        os.remove(temp_filename)
+        os.remove(local_temp_filename)
         remote_remove(
-            state, host, temp_filename,
+            state, host, local_temp_filename,
             print_output=print_output,
             print_input=print_input,
         )
@@ -200,6 +201,7 @@ def put_file(
 
 def get_file(
     state, host, remote_filename, filename_or_io,
+    remote_temp_filename=None,
     print_output=False, print_input=False,
     **kwargs  # ignored (sudo/etc)
 ):
@@ -208,14 +210,14 @@ def get_file(
     location and then reading that into our final file/IO object.
     '''
 
-    temp_filename = state.get_temp_filename(remote_filename)
+    remote_temp_filename = remote_temp_filename or state.get_temp_filename(remote_filename)
 
     try:
         docker_id = host.host_data['docker_container_id']
         docker_command = 'docker cp {0}:{1} {2}'.format(
             docker_id,
             remote_filename,
-            temp_filename,
+            remote_temp_filename,
         )
 
         status, _, stderr = ssh.run_shell_command(
@@ -225,9 +227,9 @@ def get_file(
             print_input=print_input,
         )
 
-        ssh_status = ssh.get_file(state, host, temp_filename, filename_or_io)
+        ssh_status = ssh.get_file(state, host, remote_temp_filename, filename_or_io)
     finally:
-        remote_remove(state, host, temp_filename, print_output=print_output,
+        remote_remove(state, host, remote_temp_filename, print_output=print_output,
                       print_input=print_input)
 
     if not ssh_status:

pyinfra/api/connectors/local.py

@@ -96,6 +96,7 @@ def execute_command():
 
 def put_file(
     state, host, filename_or_io, remote_filename,
+    remote_temp_filename=None,  # ignored
     print_output=False, print_input=False,
     **command_kwargs
 ):
@@ -141,6 +142,7 @@ def put_file(
 
 def get_file(
     state, host, remote_filename, filename_or_io,
+    remote_temp_filename=None,  # ignored
     print_output=False, print_input=False,
     **command_kwargs
 ):

pyinfra/api/connectors/ssh.py

@@ -339,6 +339,7 @@ def _get_file(host, remote_filename, filename_or_io):
 
 def get_file(
     state, host, remote_filename, filename_or_io,
+    remote_temp_filename=None,
     sudo=False, sudo_user=None, su_user=None,
     print_output=False, print_input=False,
     **command_kwargs
@@ -351,7 +352,7 @@ def get_file(
 
     if sudo or su_user:
         # Get temp file location
-        temp_file = state.get_temp_filename(remote_filename)
+        temp_file = remote_temp_filename or state.get_temp_filename(remote_filename)
 
         # Copy the file to the tempfile location and add read permissions
         command = 'cp {0} {1} && chmod +r {0}'.format(remote_filename, temp_file)
@@ -406,6 +407,7 @@ def _put_file(host, filename_or_io, remote_location):
 
 def put_file(
     state, host, filename_or_io, remote_filename,
+    remote_temp_filename=None,
     sudo=False, sudo_user=None, su_user=None,
     print_output=False, print_input=False,
     **command_kwargs
@@ -419,7 +421,7 @@ def put_file(
     # user connected, so upload to tmp and copy/chown w/sudo and/or su_user
     if sudo or su_user:
         # Get temp file location
-        temp_file = state.get_temp_filename(remote_filename)
+        temp_file = remote_temp_filename or state.get_temp_filename(remote_filename)
         _put_file(host, filename_or_io, temp_file)
 
         # Make sure our sudo/su user can access the file

pyinfra/api/connectors/winrm.py

@@ -191,6 +191,7 @@ def run_shell_command(
 
 def get_file(
     state, host, remote_filename, filename_or_io,
+    remote_temp_filename=None,
     **command_kwargs
 ):
     raise PyinfraError('Not implemented')
@@ -222,6 +223,7 @@ def _put_file(state, host, filename_or_io, remote_location, chunk_size=2048):
 def put_file(
     state, host, filename_or_io, remote_filename,
     print_output=False, print_input=False,
+    remote_temp_filename=None,  # ignored
     **command_kwargs
 ):
     '''

pyinfra/operations/files.py

@@ -706,11 +706,11 @@ def get(
 
     # No remote file, so assume exists and download it "blind"
     if not remote_file or force:
-        yield FileDownloadCommand(src, dest)
+        yield FileDownloadCommand(src, dest, remote_temp_filename=state.get_temp_filename(dest))
 
     # No local file, so always download
     elif not os_path.exists(dest):
-        yield FileDownloadCommand(src, dest)
+        yield FileDownloadCommand(src, dest, remote_temp_filename=state.get_temp_filename(dest))
 
     # Remote file exists - check if it matches our local
     else:
@@ -719,7 +719,7 @@ def get(
 
         # Check sha1sum, upload if needed
         if local_sum != remote_sum:
-            yield FileDownloadCommand(src, dest)
+            yield FileDownloadCommand(src, dest, remote_temp_filename=state.get_temp_filename(dest))
 
 
 @operation(pipeline_facts={
@@ -799,7 +799,11 @@ def put(
 
     # No remote file, always upload and user/group/mode if supplied
     if not remote_file or force:
-        yield FileUploadCommand(local_file, dest)
+        yield FileUploadCommand(
+            local_file,
+            dest,
+            remote_temp_filename=state.get_temp_filename(dest),
+        )
 
         if user or group:
             yield chown(dest, user, group)
@@ -813,7 +817,11 @@ def put(
 
         # Check sha1sum, upload if needed
         if local_sum != remote_sum:
-            yield FileUploadCommand(local_file, dest)
+            yield FileUploadCommand(
+                local_file,
+                dest,
+                remote_temp_filename=state.get_temp_filename(dest),
+            )
 
             if user or group:
                 yield chown(dest, user, group)

pyinfra/operations/windows_files.py

@@ -209,7 +209,11 @@ def put(
 
     # No remote file, always upload and user/group/mode if supplied
     if not remote_file or force:
-        yield FileUploadCommand(local_file, dest)
+        yield FileUploadCommand(
+            local_file,
+            dest,
+            remote_temp_filename=state.get_temp_filename(dest),
+        )
 
         # if user or group:
         #    yield chown(dest, user, group)
@@ -224,7 +228,11 @@ def put(
 
         # Check sha1sum, upload if needed
         if local_sum != remote_sum:
-            yield FileUploadCommand(local_file, dest)
+            yield FileUploadCommand(
+                local_file,
+                dest,
+                remote_temp_filename=state.get_temp_filename(dest),
+            )
 
             # if user or group:
             #    yield chown(dest, user, group)

tests/test_connectors/test_ssh.py

@@ -568,9 +568,6 @@ def test_run_shell_command_sudo_password_automatic_prompt(
             "sudo -H -A -k sh -c 'echo Šablony'"
         ), get_pty=False)
 
-    # SSH file put/get tests
-    #
-
     @patch('pyinfra.api.connectors.ssh.SSHClient')
     @patch('pyinfra.api.connectors.util._get_sudo_password')
     def test_run_shell_command_retry_for_sudo_password(
@@ -606,6 +603,9 @@ def test_run_shell_command_retry_for_sudo_password(
             get_pty=False,
         )
 
+    # SSH file put/get tests
+    #
+
     @patch('pyinfra.api.connectors.ssh.SSHClient')
     @patch('pyinfra.api.connectors.ssh.SFTPClient')
     def test_put_file(self, fake_sftp_client, fake_ssh_client):
@@ -727,6 +727,39 @@ def test_put_file_su_user_fail_copy(self, fake_sftp_client, fake_ssh_client):
             fake_open(), '/tmp/pyinfra-43db9984686317089fefcf2e38de527e4cb44487',
         )
 
+    @patch('pyinfra.api.connectors.ssh.SSHClient')
+    @patch('pyinfra.api.connectors.ssh.SFTPClient')
+    def test_put_file_sudo_custom_temp_file(self, fake_sftp_client, fake_ssh_client):
+        inventory = make_inventory(hosts=('anotherhost',))
+        State(inventory, Config())
+        host = inventory.get_host('anotherhost')
+        host.connect()
+
+        stdout_mock = MagicMock()
+        stdout_mock.channel.recv_exit_status.return_value = 0
+        fake_ssh_client().exec_command.return_value = MagicMock(), stdout_mock, MagicMock()
+
+        fake_open = mock_open(read_data='test!')
+        with patch('pyinfra.api.util.open', fake_open, create=True):
+            status = host.put_file(
+                'not-a-file', 'not another file',
+                print_output=True,
+                sudo=True,
+                sudo_user='ubuntu',
+                remote_temp_filename='/a-different-tempfile',
+            )
+
+        assert status is True
+
+        fake_ssh_client().exec_command.assert_called_with((
+            "sh -c 'rm -f "
+            "/a-different-tempfile'"
+        ), get_pty=False)
+
+        fake_sftp_client.from_transport().putfo.assert_called_with(
+            fake_open(), '/a-different-tempfile',
+        )
+
     @patch('pyinfra.api.connectors.ssh.SSHClient')
     @patch('pyinfra.api.connectors.ssh.SFTPClient')
     def test_get_file(self, fake_sftp_client, fake_ssh_client):