Sync lakectl download recursive summary with local sync

treeverse · Feb 11, 2025 · 6fe84bc · 6fe84bc
1 parent a1ee81d
commit 6fe84bc
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 12 deletions.
diff --git a/cmd/lakectl/cmd/fs_download.go b/cmd/lakectl/cmd/fs_download.go
@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/go-openapi/swag"
@@ -41,12 +42,12 @@ var fsDownloadCmd = &cobra.Command{
 		downloader := helpers.NewDownloader(client, syncFlags.Presign)
 		downloader.PartSize = downloadPartSize
 
+		remotePath := remote.GetPath()
 		if !recursive {
 			// if dest is a directory, add the file name
 			if s, _ := os.Stat(dest); s != nil && s.IsDir() {
 				dest += uri.PathSeparator
 			}
-			remotePath := remote.GetPath()
 			if remotePath != "" && strings.HasSuffix(dest, uri.PathSeparator) {
 				dest += filepath.Base(remotePath)
 			}
@@ -63,19 +64,8 @@ var fsDownloadCmd = &cobra.Command{
 		pw := newDownloadProgressWriter(syncFlags.NoProgress)
 		// ProgressRender start render progress and return callback waiting for the progress to finish.
 		go pw.Render()
-		defer func() {
-			for pw.IsRenderInProgress() {
-				// for manual-stop mode, stop when there are no more active trackers
-				if pw.LengthActive() == 0 {
-					pw.Stop()
-				}
-				const waitForRender = 100 * time.Millisecond
-				time.Sleep(waitForRender)
-			}
-		}()
 
 		ch := make(chan string, filesChanSize)
-		remotePath := remote.GetPath()
 		if remotePath != "" && !strings.HasSuffix(remotePath, uri.PathSeparator) {
 			*remote.Path += uri.PathSeparator
 		}
@@ -116,6 +106,7 @@ var fsDownloadCmd = &cobra.Command{
 		// download files in parallel
 		var wg sync.WaitGroup
 		wg.Add(syncFlags.Parallelism)
+		var downloaded int64
 		for i := 0; i < syncFlags.Parallelism; i++ {
 			go func() {
 				defer wg.Done()
@@ -139,11 +130,32 @@ var fsDownloadCmd = &cobra.Command{
 						DieErr(err)
 					}
 					tracker.MarkAsDone()
+					atomic.AddInt64(&downloaded, 1)
 				}
 			}()
 		}
 		// wait for all downloads to finish
 		wg.Wait()
+
+		// wait for progress to finish render
+		for pw.IsRenderInProgress() {
+			// for manual-stop mode, stop when there are no more active trackers
+			if pw.LengthActive() == 0 {
+				pw.Stop()
+			}
+			const waitForRender = 100 * time.Millisecond
+			time.Sleep(waitForRender)
+		}
+
+		Write(localSummaryTemplate, struct {
+			Operation  string
+			Downloaded int64
+			Removed    int
+			Uploaded   int
+		}{
+			Operation:  "Download",
+			Downloaded: downloaded,
+		})
 	},
 }
 

diff --git a/esti/lakectl_test.go b/esti/lakectl_test.go
@@ -642,6 +642,10 @@ func TestLakectlFsDownload(t *testing.T) {
 		require.Contains(t, sanitizedResult, "download ro/ro_1k.2")
 		require.Contains(t, sanitizedResult, "download ro/ro_1k.3")
 		require.Contains(t, sanitizedResult, "download ro/ro_1k.4")
+		require.Contains(t, sanitizedResult, "Download Summary:")
+		require.Contains(t, sanitizedResult, "Downloaded: 5")
+		require.Contains(t, sanitizedResult, "Uploaded: 0")
+		require.Contains(t, sanitizedResult, "Removed: 0")
 	})
 
 	t.Run("directory_with_dest", func(t *testing.T) {
@@ -652,6 +656,10 @@ func TestLakectlFsDownload(t *testing.T) {
 		require.Contains(t, sanitizedResult, "download ro/ro_1k.2")
 		require.Contains(t, sanitizedResult, "download ro/ro_1k.3")
 		require.Contains(t, sanitizedResult, "download ro/ro_1k.4")
+		require.Contains(t, sanitizedResult, "Download Summary:")
+		require.Contains(t, sanitizedResult, "Downloaded: 5")
+		require.Contains(t, sanitizedResult, "Uploaded: 0")
+		require.Contains(t, sanitizedResult, "Removed: 0")
 	})
 
 	t.Run("directory_without_recursive", func(t *testing.T) {