diff --git a/dvc/fs/base.py b/dvc/fs/base.py index cfb5306a8..f53e6bec5 100644 --- a/dvc/fs/base.py +++ b/dvc/fs/base.py @@ -217,44 +217,50 @@ class BaseFileSystem: return False return hash_.endswith(cls.CHECKSUM_DIR_SUFFIX) - def upload(self, from_info, to_info, name=None, no_progress_bar=False): - if not hasattr(self, "_upload"): - raise RemoteActionNotImplemented("upload", self.scheme) + def upload( + self, + from_info, + to_info, + total=None, + desc=None, + no_progress_bar=False, + **pbar_args, + ): + is_file_obj = hasattr(from_info, "read") + method = "upload_fobj" if is_file_obj else "_upload" + if not hasattr(self, method): + raise RemoteActionNotImplemented(method, self.scheme) if to_info.scheme != self.scheme: raise NotImplementedError + logger.debug("Uploading '%s' to '%s'", from_info, to_info) + if is_file_obj: + with Tqdm.wrapattr( + from_info, + "read", + disable=no_progress_bar, + bytes=True, + total=total, + desc=desc, + **pbar_args, + ) as wrapped: + # `size` is used to provide hints to the WebdavFileSystem + # for legacy servers. + # pylint: disable=no-member + return self.upload_fobj(wrapped, to_info, size=total) + if from_info.scheme != "local": raise NotImplementedError - logger.debug("Uploading '%s' to '%s'", from_info, to_info) - - name = name or from_info.name - - self._upload( # noqa, pylint: disable=no-member + name = desc or from_info.name + return self._upload( # noqa, pylint: disable=no-member from_info.fspath, to_info, name=name, no_progress_bar=no_progress_bar, - ) - - def upload_fobj( - self, fobj, to_info, no_progress_bar=False, size=None, **pbar_args - ): - if not hasattr(self, "_upload_fobj"): - raise RemoteActionNotImplemented("upload_fobj", self.scheme) - - with Tqdm.wrapattr( - fobj, - "read", - disable=no_progress_bar, - bytes=True, - total=size, **pbar_args, - ) as wrapped: - self._upload_fobj( # pylint: disable=no-member - wrapped, to_info, size=size - ) + ) def download( self, diff --git a/dvc/fs/fsspec_wrapper.py b/dvc/fs/fsspec_wrapper.py index e784b799b..446085f4f 100644 --- a/dvc/fs/fsspec_wrapper.py +++ b/dvc/fs/fsspec_wrapper.py @@ -130,7 +130,7 @@ class FSSpecWrapper(BaseFileSystem): self._with_bucket(path_info), exist_ok=kwargs.pop("exist_ok", True) ) - def _upload_fobj(self, fobj, to_info, size=None): + def upload_fobj(self, fobj, to_info, **kwargs): self.makedirs(to_info.parent) with self.open(to_info, "wb") as fdest: shutil.copyfileobj(fobj, fdest, length=fdest.blocksize) @@ -141,13 +141,15 @@ class FSSpecWrapper(BaseFileSystem): self.makedirs(to_info.parent) size = os.path.getsize(from_file) with open(from_file, "rb") as fobj: - self.upload_fobj( + with Tqdm.wrapattr( fobj, - to_info, - size=size, + "read", + disable=no_progress_bar, + bytes=True, + total=size, desc=name, - no_progress_bar=no_progress_bar, - ) + ) as wrapped: + self.upload_fobj(wrapped, to_info, size=size) self.fs.invalidate_cache(self._with_bucket(to_info.parent)) def _download( diff --git a/dvc/fs/gdrive.py b/dvc/fs/gdrive.py index e622e1619..db504aa6e 100644 --- a/dvc/fs/gdrive.py +++ b/dvc/fs/gdrive.py @@ -298,7 +298,7 @@ class GDriveFileSystem( return super()._with_bucket(path) - def _upload_fobj(self, fobj, to_info, size: int = None): + def upload_fobj(self, fobj, to_info, **kwargs): rpath = self._with_bucket(to_info) self.makedirs(os.path.dirname(rpath)) - return self.fs.upload_fobj(fobj, rpath, size=size) + return self.fs.upload_fobj(fobj, rpath, **kwargs) diff --git a/dvc/fs/hdfs.py b/dvc/fs/hdfs.py index 57dcf6f1a..3812155db 100644 --- a/dvc/fs/hdfs.py +++ b/dvc/fs/hdfs.py @@ -244,7 +244,7 @@ class HDFSFileSystem(BaseFileSystem): size=self.getsize(path_info), ) - def _upload_fobj(self, fobj, to_info, **kwargs): + def upload_fobj(self, fobj, to_info, **kwargs): with self.hdfs(to_info) as hdfs: with hdfs.open_output_stream(to_info.path) as fdest: shutil.copyfileobj(fobj, fdest, self.BLOCK_SIZE) diff --git a/dvc/fs/local.py b/dvc/fs/local.py index 4253282d9..5e1cbef3c 100644 --- a/dvc/fs/local.py +++ b/dvc/fs/local.py @@ -108,7 +108,7 @@ class LocalFileSystem(BaseFileSystem): self.remove(tmp_info) raise - def _upload_fobj(self, fobj, to_info, **kwargs): + def upload_fobj(self, fobj, to_info, **kwargs): self.makedirs(to_info.parent) tmp_info = to_info.parent / tmp_fname("") try: diff --git a/dvc/fs/ssh.py b/dvc/fs/ssh.py index 1947b3d6f..b3b0fbad8 100644 --- a/dvc/fs/ssh.py +++ b/dvc/fs/ssh.py @@ -120,9 +120,9 @@ class SSHFileSystem(CallbackMixin, FSSpecWrapper): # Ensure that if an interrupt happens during the transfer, we don't # pollute the cache. - def _upload_fobj(self, fobj, to_info, *args, **kwargs): + def upload_fobj(self, fobj, to_info, **kwargs): with as_atomic(self, to_info) as tmp_file: - super()._upload_fobj(fobj, tmp_file, *args, **kwargs) + super().upload_fobj(fobj, tmp_file, **kwargs) def _upload(self, from_file, to_info, *args, **kwargs): with as_atomic(self, to_info) as tmp_file: diff --git a/dvc/fs/utils.py b/dvc/fs/utils.py index c585e6fa4..e683bd1c2 100644 --- a/dvc/fs/utils.py +++ b/dvc/fs/utils.py @@ -33,7 +33,8 @@ def transfer( return from_fs.download_file(from_info, to_info) with from_fs.open(from_info, mode="rb") as fobj: - return to_fs.upload_fobj(fobj, to_info) + size = from_fs.getsize(from_info) + return to_fs.upload(fobj, to_info, total=size) except OSError as exc: # If the target file already exists, we are going to simply # ignore the exception (#4992). diff --git a/dvc/fs/webdav.py b/dvc/fs/webdav.py index 7f9d6788a..da4d451ae 100644 --- a/dvc/fs/webdav.py +++ b/dvc/fs/webdav.py @@ -1,5 +1,4 @@ import logging -import os import threading from functools import lru_cache @@ -71,15 +70,13 @@ class WebDAVFileSystem(FSSpecWrapper): # pylint:disable=abstract-method return WebdavFileSystem(**self.fs_args) - def _upload_fobj(self, fobj, to_info, size: int = None): + def upload_fobj(self, fobj, to_info, **kwargs): rpath = self.translate_path_info(to_info) - self.makedirs(os.path.dirname(rpath)) + size = kwargs.get("size") # using upload_fileobj to directly upload fileobj rather than buffering # and using overwrite=True to avoid check for an extra exists call, # as caller should ensure that the file does not exist beforehand. - return self.fs.client.upload_fileobj( - fobj, rpath, overwrite=True, size=size - ) + return self.fs.upload_fileobj(fobj, rpath, overwrite=True, size=size) @lru_cache(512) def translate_path_info(self, path): diff --git a/dvc/fs/webhdfs.py b/dvc/fs/webhdfs.py index f2cda6e8f..fdf9de1b5 100644 --- a/dvc/fs/webhdfs.py +++ b/dvc/fs/webhdfs.py @@ -141,7 +141,7 @@ class WebHDFSFileSystem(BaseFileSystem): # pylint:disable=abstract-method self.hdfs_client.makedirs(to_info.parent.path) self.hdfs_client.rename(from_info.path, to_info.path) - def _upload_fobj(self, fobj, to_info, **kwargs): + def upload_fobj(self, fobj, to_info, **kwargs): with self.hdfs_client.write(to_info.path) as fdest: shutil.copyfileobj(fobj, fdest) diff --git a/dvc/objects/db/reference.py b/dvc/objects/db/reference.py index 0c1654ac0..0332a924c 100644 --- a/dvc/objects/db/reference.py +++ b/dvc/objects/db/reference.py @@ -70,7 +70,7 @@ class ReferenceObjectDB(ObjectDB): ref_fobj = io.BytesIO(ref_file.to_bytes()) ref_fobj.seek(0) try: - self.fs.upload_fobj(ref_fobj, to_info) + self.fs.upload(ref_fobj, to_info) except OSError as exc: if isinstance(exc, FileExistsError) or ( os.name == "nt" diff --git a/dvc/objects/stage.py b/dvc/objects/stage.py index c86fc1b48..4550782d4 100644 --- a/dvc/objects/stage.py +++ b/dvc/objects/stage.py @@ -34,9 +34,8 @@ def _upload_file(path_info, fs, odb, upload_odb): tmp_info = upload_odb.path_info / tmp_fname() with fs.open(path_info, mode="rb", chunk_size=fs.CHUNK_SIZE) as stream: stream = HashedStreamReader(stream) - upload_odb.fs.upload_fobj( - stream, tmp_info, desc=path_info.name, size=fs.getsize(path_info) - ) + size = fs.getsize(path_info) + upload_odb.fs.upload(stream, tmp_info, desc=path_info.name, total=size) odb.add(tmp_info, upload_odb.fs, stream.hash_info) return path_info, odb.get(stream.hash_info) diff --git a/tests/func/test_fs.py b/tests/func/test_fs.py index c8034f253..04e016fb8 100644 --- a/tests/func/test_fs.py +++ b/tests/func/test_fs.py @@ -405,7 +405,7 @@ def test_fs_makedirs_on_upload_and_copy(dvc, cloud): fs = cls(**config) with io.BytesIO(b"foo") as stream: - fs.upload_fobj(stream, cloud / "dir" / "foo") + fs.upload(stream, cloud / "dir" / "foo") assert fs.isdir(cloud / "dir") assert fs.exists(cloud / "dir" / "foo")