_put_file: upload file atomically (#129)

Author: skshetry

dvc_ssh/__init__.py

@@ -105,22 +105,13 @@ def _prepare_credentials(self, **config):
     @wrap_prop(threading.Lock())
     @cached_property
     def fs(self):
-        from sshfs import SSHFileSystem as _SSHFileSystem
+        from . import spec
 
-        return _SSHFileSystem(**self.fs_args)
+        return spec.SSHFileSystem(**self.fs_args)
 
     # Ensure that if an interrupt happens during the transfer, we don't
     # pollute the cache.
 
     def upload_fobj(self, fobj, to_info, **kwargs):
         with as_atomic(self, to_info) as tmp_file:
             super().upload_fobj(fobj, tmp_file, **kwargs)
-
-    def put_file(
-        self,
-        from_file,
-        to_info,
-        **kwargs,
-    ):
-        with as_atomic(self, to_info) as tmp_file:
-            super().put_file(from_file, tmp_file, **kwargs)

dvc_ssh/client.py

@@ -75,7 +75,7 @@ async def public_key_auth_requested(
                 pubkey: Optional[SSHKey] = read_public_key(pubkey_to_load)
             except (OSError, KeyImportError):
                 pubkey = None
-            return SSHLocalKeyPair(key, pubkey)
+            return SSHLocalKeyPair(key, pubkey, cert=None, enc_key=None)
         return None
 
     async def _read_private_key_interactive(self, path: "FilePath") -> "SSHKey":

dvc_ssh/spec.py

@@ -0,0 +1,28 @@
+import posixpath
+from secrets import token_urlsafe
+
+from sshfs import SSHFileSystem as _SSHFileSystem
+
+
+def tmp_fname(prefix: str = "") -> str:
+    """Temporary name for a partial download"""
+    return f"{prefix}.{token_urlsafe(16)}.tmp"
+
+
+class SSHFileSystem(_SSHFileSystem):
+    async def _put_file(self, lpath, rpath, *args, **kwargs):
+        parent = posixpath.dirname(rpath)
+        tmp_info = posixpath.join(parent, tmp_fname())
+        try:
+            await super()._put_file(lpath, tmp_info, *args, **kwargs)
+        except BaseException:
+            # Handle stuff like KeyboardInterrupt
+            # as well as other errors that might
+            # arise during file transfer.
+            try:
+                await self._rm_file(tmp_info)
+            except FileNotFoundError:
+                pass
+            raise
+        else:
+            await self._mv(tmp_info, rpath)