pass size to upload_fobj (#6187)

This allows us to select between chunked and non-chunked uploading
in Webdav, and hopefully in the http itself.

dvc/fs/base.py

@@ -235,14 +235,23 @@ class BaseFileSystem:
             no_progress_bar=no_progress_bar,
         )
 
-    def upload_fobj(self, fobj, to_info, no_progress_bar=False, **pbar_args):
+    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, **pbar_args
+            fobj,
+            "read",
+            disable=no_progress_bar,
+            bytes=True,
+            total=size,
+            **pbar_args,
         ) as wrapped:
-            self._upload_fobj(wrapped, to_info)  # pylint: disable=no-member
+            self._upload_fobj(  # pylint: disable=no-member
+                wrapped, to_info, size=size
+            )
 
     def download(
         self,

dvc/fs/fsspec_wrapper.py

@@ -118,20 +118,20 @@ class FSSpecWrapper(BaseFileSystem):
         info["name"] = self._strip_bucket(info["name"])
         return info
 
-    def _upload_fobj(self, fobj, to_info):
+    def _upload_fobj(self, fobj, to_info, size=None):
         with self.open(to_info, "wb") as fdest:
             shutil.copyfileobj(fobj, fdest, length=fdest.blocksize)
 
     def _upload(
         self, from_file, to_info, name=None, no_progress_bar=False, **kwargs
     ):
-        total = os.path.getsize(from_file)
+        size = os.path.getsize(from_file)
         with open(from_file, "rb") as fobj:
             self.upload_fobj(
                 fobj,
                 self._with_bucket(to_info),
+                size=size,
                 desc=name,
-                total=total,
                 no_progress_bar=no_progress_bar,
             )
         self.fs.invalidate_cache(self._with_bucket(to_info.parent))

dvc/fs/gdrive.py

@@ -618,22 +618,27 @@ class GDriveFileSystem(BaseFileSystem):  # pylint:disable=abstract-method
         gdrive_file.FetchMetadata(fields="fileSize")
         return {"size": gdrive_file.get("fileSize")}
 
-    def _upload_fobj(self, fobj, to_info):
+    def _upload_fobj(self, fobj, to_info, **kwargs):
         dirname = to_info.parent
         assert dirname
         parent_id = self._get_item_id(dirname, create=True)
         self._gdrive_upload_fobj(fobj, parent_id, to_info.name)
 
     def _upload(
-        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
+        self,
+        from_file,
+        to_info,
+        name=None,
+        no_progress_bar=False,
+        **_kwargs,
     ):
         with open(from_file, "rb") as fobj:
             self.upload_fobj(
                 fobj,
                 to_info,
+                size=os.path.getsize(from_file),
                 no_progress_bar=no_progress_bar,
                 desc=name or to_info.name,
-                total=os.path.getsize(from_file),
             )
 
     def _download(self, from_info, to_file, name=None, no_progress_bar=False):

dvc/fs/hdfs.py

@@ -230,7 +230,7 @@ class HDFSFileSystem(BaseFileSystem):
             size=self.getsize(path_info),
         )
 
-    def _upload_fobj(self, fobj, to_info):
+    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)

dvc/fs/http.py

@@ -147,7 +147,7 @@ class HTTPFileSystem(BaseFileSystem):  # pylint:disable=abstract-method
         size = self._content_length(resp)
         return {"etag": etag, "size": size}
 
-    def _upload_fobj(self, fobj, to_info):
+    def _upload_fobj(self, fobj, to_info, **kwargs):
         def chunks(fobj):
             while True:
                 chunk = fobj.read(self.CHUNK_SIZE)
@@ -184,9 +184,9 @@ class HTTPFileSystem(BaseFileSystem):  # pylint:disable=abstract-method
             self.upload_fobj(
                 fobj,
                 to_info,
+                size=None if no_progress_bar else os.path.getsize(from_file),
                 no_progress_bar=no_progress_bar,
                 desc=name or to_info.url,
-                total=None if no_progress_bar else os.path.getsize(from_file),
             )
 
     @staticmethod

dvc/fs/local.py

@@ -103,7 +103,7 @@ class LocalFileSystem(BaseFileSystem):
             self.remove(tmp_info)
             raise
 
-    def _upload_fobj(self, fobj, to_info):
+    def _upload_fobj(self, fobj, to_info, **kwargs):
         self.makedirs(to_info.parent)
         tmp_info = to_info.parent / tmp_fname("")
         try:

dvc/fs/oss.py

@@ -104,7 +104,7 @@ class OSSFileSystem(BaseFileSystem):  # pylint:disable=abstract-method
         logger.debug(f"Removing oss://{path_info}")
         self._get_bucket(path_info.bucket).delete_object(path_info.path)
 
-    def _upload_fobj(self, fobj, to_info):
+    def _upload_fobj(self, fobj, to_info, **kwargs):
         self._get_bucket(to_info.bucket).put_object(to_info.path, fobj)
 
     def _upload(

dvc/fs/ssh/__init__.py

@@ -237,7 +237,7 @@ class SSHFileSystem(BaseFileSystem):  # pylint:disable=abstract-method
         with self.ssh(path_info) as ssh:
             return ssh.info(path_info.path)
 
-    def _upload_fobj(self, fobj, to_info):
+    def _upload_fobj(self, fobj, to_info, **kwargs):
         self.makedirs(to_info.parent)
         with self.open(to_info, mode="wb") as fdest:
             shutil.copyfileobj(fobj, fdest)

dvc/fs/webdav.py

@@ -71,13 +71,15 @@ class WebDAVFileSystem(FSSpecWrapper):  # pylint:disable=abstract-method
 
         return WebdavFileSystem(**self.fs_args)
 
-    def _upload_fobj(self, fobj, to_info):
+    def _upload_fobj(self, fobj, to_info, size: int = None):
         rpath = self.translate_path_info(to_info)
         self.makedirs(os.path.dirname(rpath))
         # 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)
+        return self.fs.client.upload_fileobj(
+            fobj, rpath, overwrite=True, size=size
+        )
 
     def makedirs(self, path_info):
         path = self.translate_path_info(path_info)

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):
+    def _upload_fobj(self, fobj, to_info, **kwargs):
         with self.hdfs_client.write(to_info.path) as fdest:
             shutil.copyfileobj(fobj, fdest)
 

dvc/objects/stage.py

@@ -19,7 +19,7 @@ def _upload_file(path_info, fs, odb):
     with fs.open(path_info, mode="rb", chunk_size=fs.CHUNK_SIZE) as stream:
         stream = HashedStreamReader(stream)
         odb.fs.upload_fobj(
-            stream, tmp_info, desc=path_info.name, total=fs.getsize(path_info)
+            stream, tmp_info, desc=path_info.name, size=fs.getsize(path_info)
         )
 
     obj = HashFile(tmp_info, odb.fs, stream.hash_info)

setup.py

@@ -106,7 +106,7 @@ ssh = ["paramiko[invoke]>=2.7.0"]
 # Remove the env marker if/when pyarrow is available for Python3.9
 hdfs = ["pyarrow>=2.0.0"]
 webhdfs = ["hdfs==2.5.8"]
-webdav = ["webdav4>=0.7.0"]
+webdav = ["webdav4>=0.8.1"]
 # gssapi should not be included in all_remotes, because it doesn't have wheels
 # for linux and mac, so it will fail to compile if user doesn't have all the
 # requirements, including kerberos itself. Once all the wheels are available,