From d01e2355d494019413d02bbbfe153c7e7f4b866e Mon Sep 17 00:00:00 2001
From: Peter Rowlands <peter@pmrowla.com>
Date: Thu, 3 Aug 2023 17:37:37 +0900
Subject: [PATCH] push/pull: generate explicit progress bars for dvc-data
 transfers

---
 dvc/data_cloud.py | 49 +++++++++++++++++++++++++++++------------------
 dvc/output.py     | 21 ++++++++++++--------
 2 files changed, 43 insertions(+), 27 deletions(-)

diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py
index 6af282691..e8622b34c 100644
--- a/dvc/data_cloud.py
+++ b/dvc/data_cloud.py
@@ -4,6 +4,7 @@ import logging
 from typing import TYPE_CHECKING, Iterable, Optional, Set, Tuple
 
 from dvc.config import NoRemoteError, RemoteConfigError
+from dvc.fs.callbacks import Callback
 from dvc.utils.objects import cached_property
 from dvc_data.hashfile.db import get_index
 from dvc_data.hashfile.transfer import TransferResult
@@ -209,15 +210,20 @@ class DataCloud:
             cache = self.repo.cache.legacy
         else:
             cache = self.repo.cache.local
-        return self.transfer(
-            cache,
-            odb,
-            objs,
-            jobs=jobs,
-            dest_index=get_index(odb),
-            cache_odb=cache,
-            validate_status=self._log_missing,
-        )
+        with Callback.as_tqdm_callback(
+            desc=f"Pushing to {odb.fs.unstrip_protocol(odb.path)}",
+            unit="file",
+        ) as cb:
+            return self.transfer(
+                cache,
+                odb,
+                objs,
+                jobs=jobs,
+                dest_index=get_index(odb),
+                cache_odb=cache,
+                validate_status=self._log_missing,
+                callback=cb,
+            )
 
     def pull(
         self,
@@ -263,16 +269,21 @@ class DataCloud:
             cache = self.repo.cache.legacy
         else:
             cache = self.repo.cache.local
-        return self.transfer(
-            odb,
-            cache,
-            objs,
-            jobs=jobs,
-            src_index=get_index(odb),
-            cache_odb=cache,
-            verify=odb.verify,
-            validate_status=self._log_missing,
-        )
+        with Callback.as_tqdm_callback(
+            desc=f"Fetching from {odb.fs.unstrip_protocol(odb.path)}",
+            unit="file",
+        ) as cb:
+            return self.transfer(
+                odb,
+                cache,
+                objs,
+                jobs=jobs,
+                src_index=get_index(odb),
+                cache_odb=cache,
+                verify=odb.verify,
+                validate_status=self._log_missing,
+                callback=cb,
+            )
 
     def status(
         self,
diff --git a/dvc/output.py b/dvc/output.py
index 5dfd17e00..539ff531d 100644
--- a/dvc/output.py
+++ b/dvc/output.py
@@ -1048,14 +1048,19 @@ class Output:
             upload=upload,
             no_progress_bar=no_progress_bar,
         )
-        otransfer(
-            staging,
-            odb,
-            {obj.hash_info},
-            jobs=jobs,
-            hardlink=False,
-            shallow=False,
-        )
+        with Callback.as_tqdm_callback(
+            desc=f"Transferring to {odb.fs.unstrip_protocol(odb.path)}",
+            unit="file",
+        ) as cb:
+            otransfer(
+                staging,
+                odb,
+                {obj.hash_info},
+                jobs=jobs,
+                hardlink=False,
+                shallow=False,
+                callback=cb,
+            )
 
         self.hash_info = obj.hash_info
         self.files = None