Repository: iterative/dvc
Issue: remote: ssh: use temporary file when uploading

User with a nick `rsadek` on discord, reported that while he was pushing data to ssh remote, his internet connection got interrupted and so the file uploaded got corrupted as well, which resulted in `dvc pull` pulling corrupted files, detecting the corruption, removing those files and refusing to checkout a dataset directory. Need to use a temporary file when uploading to ssh remote, so we could ensure that we've uploaded it successfully and only then rename it to a proper name. This is the same as we do in `dowload()` methods.

Changed files:
dvc/remote/local.py, dvc/utils/__init__.py, dvc/remote/gs.py, dvc/remote/s3.py, dvc/remote/azure.py, dvc/remote/base.py, dvc/remote/ssh/connection.py

diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py
index 43a978d02..3853a9d67 100644
--- a/dvc/remote/azure.py
+++ b/dvc/remote/azure.py
@@ -10,6 +10,7 @@ except ImportError:
     BlockBlobService = None
 
 import dvc.logger as logger
+from dvc.utils import tmp_fname
 from dvc.utils.compat import urlparse, makedirs
 from dvc.progress import progress
 from dvc.config import Config
@@ -170,7 +171,7 @@ class RemoteAzure(RemoteBase):
                 )
             )
 
-            tmp_file = self.tmp_file(to_info["path"])
+            tmp_file = tmp_fname(to_info["path"])
             if not name:
                 name = os.path.basename(to_info["path"])
 
diff --git a/dvc/remote/base.py b/dvc/remote/base.py
index 9952ce824..982e26c79 100644
--- a/dvc/remote/base.py
+++ b/dvc/remote/base.py
@@ -112,13 +112,6 @@ class RemoteBase(object):
             return None
         return m.group(name)
 
-    @staticmethod
-    def tmp_file(fname):
-        """ Temporary name for a partial download """
-        import uuid
-
-        return fname + "." + str(uuid.uuid4())
-
     def save_info(self, path_info):
         raise NotImplementedError
 
diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py
index aa7d8ffd6..79737ec9e 100644
--- a/dvc/remote/gs.py
+++ b/dvc/remote/gs.py
@@ -9,6 +9,7 @@ except ImportError:
     storage = None
 
 import dvc.logger as logger
+from dvc.utils import tmp_fname
 from dvc.utils.compat import urlparse, makedirs
 from dvc.remote.base import RemoteBase
 from dvc.config import Config
@@ -183,7 +184,7 @@ class RemoteGS(RemoteBase):
             )
             logger.debug(msg)
 
-            tmp_file = self.tmp_file(to_info["path"])
+            tmp_file = tmp_fname(to_info["path"])
             if not name:
                 name = os.path.basename(to_info["path"])
 
diff --git a/dvc/remote/local.py b/dvc/remote/local.py
index bc698041b..ae29921ab 100644
--- a/dvc/remote/local.py
+++ b/dvc/remote/local.py
@@ -14,7 +14,7 @@ from operator import itemgetter
 import dvc.logger as logger
 from dvc.system import System
 from dvc.remote.base import RemoteBase, STATUS_MAP, STATUS_NEW, STATUS_DELETED
-from dvc.utils import remove, move, copyfile, dict_md5, to_chunks
+from dvc.utils import remove, move, copyfile, dict_md5, to_chunks, tmp_fname
 from dvc.utils import LARGE_DIR_SIZE
 from dvc.config import Config
 from dvc.exceptions import DvcException
@@ -522,7 +522,7 @@ class RemoteLOCAL(RemoteBase):
                 name = os.path.basename(to_info["path"])
 
             makedirs(os.path.dirname(to_info["path"]), exist_ok=True)
-            tmp_file = self.tmp_file(to_info["path"])
+            tmp_file = tmp_fname(to_info["path"])
             try:
                 copyfile(
                     from_info["path"],
diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py
index e38369d52..a9216ec62 100644
--- a/dvc/remote/s3.py
+++ b/dvc/remote/s3.py
@@ -9,6 +9,7 @@ except ImportError:
     boto3 = None
 
 import dvc.logger as logger
+from dvc.utils import tmp_fname
 from dvc.utils.compat import urlparse, makedirs
 from dvc.progress import progress
 from dvc.config import Config
@@ -222,7 +223,7 @@ class RemoteS3(RemoteBase):
             )
             logger.debug(msg)
 
-            tmp_file = self.tmp_file(to_info["path"])
+            tmp_file = tmp_fname(to_info["path"])
             if not name:
                 name = os.path.basename(to_info["path"])
 
diff --git a/dvc/remote/ssh/connection.py b/dvc/remote/ssh/connection.py
index 10f714f8c..cc28692cd 100644
--- a/dvc/remote/ssh/connection.py
+++ b/dvc/remote/ssh/connection.py
@@ -7,6 +7,7 @@ except ImportError:
     paramiko = None
 
 import dvc.logger as logger
+from dvc.utils import tmp_fname
 from dvc.utils.compat import makedirs
 from dvc.progress import progress
 from dvc.exceptions import DvcException
@@ -145,12 +146,10 @@ class SSHConnection:
         self._sftp.remove(path)
 
     def download(self, src, dest, no_progress_bar=False, progress_title=None):
-        from uuid import uuid4
-
         self._sftp_connect()
 
         makedirs(posixpath.dirname(dest), exist_ok=True)
-        tmp_file = dest + "." + str(uuid4())
+        tmp_file = tmp_fname(dest)
 
         if no_progress_bar:
             self._sftp.get(src, tmp_file)
@@ -167,16 +166,19 @@ class SSHConnection:
         self._sftp_connect()
 
         self.makedirs(posixpath.dirname(dest))
+        tmp_file = tmp_fname(dest)
 
         if no_progress_bar:
-            self._sftp.put(src, dest)
+            self._sftp.put(src, tmp_file)
         else:
             if not progress_title:
                 progress_title = os.path.basename(dest)
 
-            self._sftp.put(src, dest, callback=create_cb(progress_title))
+            self._sftp.put(src, tmp_file, callback=create_cb(progress_title))
             progress.finish_target(progress_title)
 
+        self._sftp.rename(tmp_file, dest)
+
     def execute(self, cmd):
         stdin, stdout, stderr = self._ssh.exec_command(cmd)
         channel = stdout.channel
diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py
index bb1fa0321..093e291db 100644
--- a/dvc/utils/__init__.py
+++ b/dvc/utils/__init__.py
@@ -212,3 +212,10 @@ def convert_to_unicode(data):
         return type(data)(map(convert_to_unicode, data))
     else:
         return data
+
+
+def tmp_fname(fname):
+    """ Temporary name for a partial download """
+    from uuid import uuid4
+
+    return fname + "." + str(uuid4()) + ".tmp"