Merge pull request #1646 from mroutis/fix-1570

ssh: upload to a tmp file to prevent data corruption

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"])
 

dvc/remote/base.py

@@ -108,13 +108,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
 

dvc/remote/gs.py

@@ -8,6 +8,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
@@ -182,7 +183,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"])
 

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"],

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"])
 

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

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"