fetch/push: use dvc-data TransferResult

dvc/data_cloud.py

@@ -11,6 +11,7 @@ if TYPE_CHECKING:
     from dvc_data.hashfile.db import HashFileDB
     from dvc_data.hashfile.hash_info import HashInfo
     from dvc_data.hashfile.status import CompareStatusResult
+    from dvc_data.hashfile.transfer import TransferResult
 
 logger = logging.getLogger(__name__)
 
@@ -106,14 +107,10 @@ class DataCloud:
         dest_odb: "HashFileDB",
         objs: Iterable["HashInfo"],
         **kwargs,
-    ):
-        from dvc.exceptions import FileTransferError
-        from dvc_data.hashfile.transfer import TransferError, transfer
+    ) -> "TransferResult":
+        from dvc_data.hashfile.transfer import transfer
 
-        try:
-            return transfer(src_odb, dest_odb, objs, **kwargs)
-        except TransferError as exc:
-            raise FileTransferError(exc.fails) from exc
+        return transfer(src_odb, dest_odb, objs, **kwargs)
 
     def push(
         self,
@@ -121,7 +118,7 @@ class DataCloud:
         jobs: Optional[int] = None,
         remote: Optional[str] = None,
         odb: Optional["HashFileDB"] = None,
-    ):
+    ) -> "TransferResult":
         """Push data items in a cloud-agnostic way.
 
         Args:
@@ -148,7 +145,7 @@ class DataCloud:
         jobs: Optional[int] = None,
         remote: Optional[str] = None,
         odb: Optional["HashFileDB"] = None,
-    ):
+    ) -> "TransferResult":
         """Pull data items in a cloud-agnostic way.
 
         Args:

dvc/repo/fetch.py

@@ -3,12 +3,15 @@ from contextlib import suppress
 from typing import TYPE_CHECKING, Optional
 
 from dvc.config import NoRemoteError
-from dvc.exceptions import DownloadError, FileTransferError
+from dvc.exceptions import DownloadError
 from dvc.fs import Schemes
 
 from . import locked
 
 if TYPE_CHECKING:
+    from dvc.repo import Repo
+    from dvc.types import TargetType
+    from dvc_data.hashfile.transfer import TransferResult
     from dvc_objects.db.base import ObjectDB
 
 logger = logging.getLogger(__name__)
@@ -63,7 +66,17 @@ def fetch(
         if _remote.worktree:
             return _fetch_worktree(self, _remote)
 
-    used = self.used_objs(
+    downloaded_count = 0
+    failed_count = 0
+
+    try:
+        if run_cache:
+            self.stage_cache.pull(remote)
+    except DownloadError as exc:
+        failed_count += exc.amount
+
+    result = _fetch(
+        self,
         targets,
         all_branches=all_branches,
         all_tags=all_tags,
@@ -74,45 +87,10 @@ def fetch(
         jobs=jobs,
         recursive=recursive,
         revs=revs,
+        odb=odb,
     )
-
-    downloaded = 0
-    failed = 0
-
-    try:
-        if run_cache:
-            self.stage_cache.pull(remote)
-    except DownloadError as exc:
-        failed += exc.amount
-
-    if odb:
-        all_ids = set()
-        for _odb, obj_ids in used.items():
-            all_ids.update(obj_ids)
-        d, f = _fetch(
-            self,
-            all_ids,
-            jobs=jobs,
-            remote=remote,
-            odb=odb,
-        )
-        downloaded += d
-        failed += f
-    else:
-        for src_odb, obj_ids in sorted(
-            used.items(),
-            key=lambda item: item[0] is not None
-            and item[0].fs.protocol == Schemes.MEMORY,
-        ):
-            d, f = _fetch(
-                self,
-                obj_ids,
-                jobs=jobs,
-                remote=remote,
-                odb=src_odb,
-            )
-            downloaded += d
-            failed += f
+    downloaded_count = len(result.transferred)
+    failed_count += len(result.failed)
 
     d, f = _fetch_partial_imports(
         self,
@@ -123,23 +101,59 @@ def fetch(
         recursive=recursive,
         revs=revs,
     )
-    downloaded += d
-    failed += f
+    downloaded_count += d
+    failed_count += f
 
-    if failed:
-        raise DownloadError(failed)
+    if failed_count:
+        raise DownloadError(failed_count)
 
-    return downloaded
+    return downloaded_count
 
 
-def _fetch(repo, obj_ids, **kwargs):
-    downloaded = 0
-    failed = 0
-    try:
-        downloaded += repo.cloud.pull(obj_ids, **kwargs)
-    except FileTransferError as exc:
-        failed += exc.amount
-    return downloaded, failed
+def _fetch(
+    repo: "Repo",
+    targets: "TargetType",
+    remote: Optional[str] = None,
+    jobs: Optional[int] = None,
+    odb: Optional["ObjectDB"] = None,
+    **kwargs,
+) -> "TransferResult":
+    from dvc_data.hashfile.transfer import TransferResult
+
+    result = TransferResult(set(), set())
+    used = repo.used_objs(
+        targets,
+        remote=remote,
+        jobs=jobs,
+        **kwargs,
+    )
+    if odb:
+        all_ids = set()
+        for _odb, obj_ids in used.items():
+            all_ids.update(obj_ids)
+        d, f = repo.cloud.pull(
+            all_ids,
+            jobs=jobs,
+            remote=remote,
+            odb=odb,
+        )
+        result.transferred.update(d)
+        result.failed.update(f)
+    else:
+        for src_odb, obj_ids in sorted(
+            used.items(),
+            key=lambda item: item[0] is not None
+            and item[0].fs.protocol == Schemes.MEMORY,
+        ):
+            d, f = repo.cloud.pull(
+                obj_ids,
+                jobs=jobs,
+                remote=remote,
+                odb=src_odb,
+            )
+            result.transferred.update(d)
+            result.failed.update(f)
+    return result
 
 
 def _fetch_partial_imports(repo, targets, **kwargs):

dvc/repo/push.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, Optional
 
-from dvc.exceptions import FileTransferError, UploadError
+from dvc.exceptions import UploadError
 
 from ..utils import glob_targets
 from . import locked
@@ -103,18 +103,18 @@ def push(
             if not include_imports and dest_odb and dest_odb.read_only:
                 continue
             all_ids.update(obj_ids)
-        try:
-            pushed += self.cloud.push(all_ids, jobs, remote=remote, odb=odb)
-        except FileTransferError as exc:
-            raise UploadError(exc.amount)
+        result = self.cloud.push(all_ids, jobs, remote=remote, odb=odb)
+        if result.failed:
+            raise UploadError(len(result.failed))
+        pushed += len(result.transferred)
     else:
         for dest_odb, obj_ids in used.items():
             if dest_odb and dest_odb.read_only:
                 continue
-            try:
-                pushed += self.cloud.push(
-                    obj_ids, jobs, remote=remote, odb=odb or dest_odb
-                )
-            except FileTransferError as exc:
-                raise UploadError(exc.amount)
+            result = self.cloud.push(
+                obj_ids, jobs, remote=remote, odb=odb or dest_odb
+            )
+            if result.failed:
+                raise UploadError(len(result.failed))
+            pushed += len(result.transferred)
     return pushed

dvc/stage/cache.py

@@ -9,7 +9,6 @@ from funcy import cached_property, first
 from dvc import fs
 from dvc.exceptions import DvcException
 from dvc.utils import dict_sha256, relpath
-from dvc_data.hashfile.transfer import _log_exceptions
 
 if TYPE_CHECKING:
     from dvc_objects.db import ObjectDB
@@ -222,7 +221,7 @@ class StageCache:
 
         from_fs = from_odb.fs
         to_fs = to_odb.fs
-        func = _log_exceptions(fs.generic.copy)
+        func = fs.generic.log_exceptions(fs.generic.copy)
         runs = from_fs.path.join(from_odb.path, "runs")
 
         http_odb = next(