Show copy progress for CopyFileToMachine on Shell Service (#4918)

Author: michaellee1019

cli/app.go

@@ -69,6 +69,7 @@ const (
 	generalFlagTags              = "tags"
 	generalFlagStart             = "start"
 	generalFlagEnd               = "end"
+	generalFlagNoProgress        = "no-progress"
 
 	moduleFlagLanguage        = "language"
 	moduleFlagPublicNamespace = "public-namespace"
@@ -2349,6 +2350,11 @@ Copy multiple files from the machine to a local destination with recursion and k
 									// Note(erd): maybe support access time in the future if needed
 									Usage: "preserve modification times and file mode bits from the source files",
 								},
+								&cli.BoolFlag{
+									Name:    generalFlagNoProgress,
+									Aliases: []string{"n"},
+									Usage:   "hide progress of the file transfer",
+								},
 							}...),
 							Action: createCommandWithT[machinesPartCopyFilesArgs](MachinesPartCopyFilesAction),
 						},
@@ -2821,6 +2827,10 @@ This won't work unless you have an existing installation of our GitHub app on yo
 							Name:  moduleFlagLocal,
 							Usage: "if the target machine is localhost, run the entrypoint directly rather than transferring a bundle",
 						},
+						&cli.BoolFlag{
+							Name:  generalFlagNoProgress,
+							Usage: "hide progress of the file transfer",
+						},
 						&cli.StringFlag{
 							Name:  moduleFlagHomeDir,
 							Usage: "remote user's home directory. only necessary if you're targeting a remote machine where $HOME is not /root",

cli/client.go

@@ -7,6 +7,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"io/fs"
+	"math"
 	"net"
 	"net/http"
 	"net/url"
@@ -1289,6 +1291,7 @@ type machinesPartCopyFilesArgs struct {
 	Part         string
 	Recursive    bool
 	Preserve     bool
+	NoProgress   bool
 }
 
 // MachinesPartCopyFilesAction is the corresponding Action for 'machines part cp'.
@@ -1395,6 +1398,7 @@ func (c *viamClient) machinesPartCopyFilesAction(
 			paths,
 			destination,
 			logger,
+			flagArgs.NoProgress,
 		)
 	}
 	if err := doCopy(); err != nil {
@@ -2471,12 +2475,13 @@ func (c *viamClient) copyFilesToMachine(
 	paths []string,
 	destination string,
 	logger logging.Logger,
+	noProgress bool,
 ) error {
 	shellSvc, closeClient, err := c.connectToShellService(orgStr, locStr, robotStr, partStr, debug, logger)
 	if err != nil {
 		return err
 	}
-	return c.copyFilesToMachineInner(shellSvc, closeClient, allowRecursion, preserve, paths, destination)
+	return c.copyFilesToMachineInner(shellSvc, closeClient, allowRecursion, preserve, paths, destination, noProgress)
 }
 
 // copyFilesToFqdn is a copyFilesToMachine variant that makes use of pre-fetched part FQDN.
@@ -2488,12 +2493,13 @@ func (c *viamClient) copyFilesToFqdn(
 	paths []string,
 	destination string,
 	logger logging.Logger,
+	noProgress bool,
 ) error {
 	shellSvc, closeClient, err := c.connectToShellServiceFqdn(fqdn, debug, logger)
 	if err != nil {
 		return err
 	}
-	return c.copyFilesToMachineInner(shellSvc, closeClient, allowRecursion, preserve, paths, destination)
+	return c.copyFilesToMachineInner(shellSvc, closeClient, allowRecursion, preserve, paths, destination, noProgress)
 }
 
 // copyFilesToMachineInner is the common logic for both copyFiles variants.
@@ -2504,16 +2510,81 @@ func (c *viamClient) copyFilesToMachineInner(
 	preserve bool,
 	paths []string,
 	destination string,
+	noProgress bool,
 ) error {
 	defer func() {
 		utils.UncheckedError(closeClient(c.c.Context))
 	}()
 
-	// prepare a factory that understands the file copying service (RPC or not).
-	copyFactory := shell.NewCopyFileToMachineFactory(destination, preserve, shellSvc)
-	// make a reader copier that just does the traversal and copy work for us. Think of
-	// this as a tee reader.
-	readCopier, err := shell.NewLocalFileReadCopier(paths, allowRecursion, false, copyFactory)
+	if noProgress {
+		// prepare a factory that understands the file copying service (RPC or not).
+		copyFactory := shell.NewCopyFileToMachineFactory(destination, preserve, shellSvc)
+		// make a reader copier that just does the traversal and copy work for us. Think of
+		// this as a tee reader.
+		readCopier, err := shell.NewLocalFileReadCopier(paths, allowRecursion, false, copyFactory)
+		if err != nil {
+			return err
+		}
+		defer func() {
+			if err := readCopier.Close(c.c.Context); err != nil {
+				utils.UncheckedError(err)
+			}
+		}()
+
+		// ReadAll the files into the copier.
+		return readCopier.ReadAll(c.c.Context)
+	}
+
+	// Calculate total size of all files to be copied
+	var totalSize int64
+	for _, path := range paths {
+		info, err := os.Stat(path)
+		if err != nil {
+			return err
+		}
+		if info.IsDir() && allowRecursion {
+			err := filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
+				if err != nil {
+					return err
+				}
+				if !info.IsDir() {
+					totalSize += info.Size()
+				}
+				return nil
+			})
+			if err != nil {
+				return err
+			}
+		} else if !info.IsDir() {
+			totalSize += info.Size()
+		}
+	}
+
+	// Create a progress tracking function
+	var currentFile string
+	progressFunc := func(bytes int64, file string, fileSize int64) {
+		if file != currentFile {
+			if currentFile != "" {
+				//nolint:errcheck // progress display is non-critical
+				_, _ = os.Stdout.WriteString("\n")
+			}
+			currentFile = file
+			//nolint:errcheck // progress display is non-critical
+			_, _ = os.Stdout.WriteString(fmt.Sprintf("Copying %s...\n", file))
+		}
+		uploadPercent := int(math.Ceil(100 * float64(bytes) / float64(fileSize)))
+		//nolint:errcheck // progress display is non-critical
+		_, _ = os.Stdout.WriteString(fmt.Sprintf("\rProgress: %d%% (%d/%d bytes)", uploadPercent, bytes, fileSize))
+	}
+
+	// Wrap the copy factory to track progress
+	progressFactory := &progressTrackingFactory{
+		factory:    shell.NewCopyFileToMachineFactory(destination, preserve, shellSvc),
+		onProgress: progressFunc,
+	}
+
+	// Create a new read copier with the progress tracking factory
+	readCopier, err := shell.NewLocalFileReadCopier(paths, allowRecursion, false, progressFactory)
 	if err != nil {
 		return err
 	}
@@ -2524,7 +2595,88 @@ func (c *viamClient) copyFilesToMachineInner(
 	}()
 
 	// ReadAll the files into the copier.
-	return readCopier.ReadAll(c.c.Context)
+	err = readCopier.ReadAll(c.c.Context)
+	return err
+}
+
+// progressTrackingFactory wraps a copy factory to track progress.
+type progressTrackingFactory struct {
+	factory    shell.FileCopyFactory
+	onProgress func(int64, string, int64)
+}
+
+func (ptf *progressTrackingFactory) MakeFileCopier(ctx context.Context, sourceType shell.CopyFilesSourceType) (shell.FileCopier, error) {
+	copier, err := ptf.factory.MakeFileCopier(ctx, sourceType)
+	if err != nil {
+		return nil, err
+	}
+	return &progressTrackingCopier{
+		copier:     copier,
+		onProgress: ptf.onProgress,
+	}, nil
+}
+
+// progressTrackingCopier wraps a file copier to track progress.
+type progressTrackingCopier struct {
+	copier     shell.FileCopier
+	onProgress func(int64, string, int64)
+}
+
+func (ptc *progressTrackingCopier) Copy(ctx context.Context, file shell.File) error {
+	// Get file size
+	info, err := file.Data.Stat()
+	if err != nil {
+		return err
+	}
+	fileSize := info.Size()
+
+	// Create a progress tracking reader
+	progressReader := &progressReader{
+		reader:     file.Data,
+		onProgress: ptc.onProgress,
+		fileName:   file.RelativeName,
+		fileSize:   fileSize,
+	}
+
+	// Create a new file with the progress tracking reader
+	progressFile := shell.File{
+		RelativeName: file.RelativeName,
+		Data:         progressReader,
+	}
+
+	return ptc.copier.Copy(ctx, progressFile)
+}
+
+func (ptc *progressTrackingCopier) Close(ctx context.Context) error {
+	//nolint:errcheck // progress display is non-critical
+	_, _ = os.Stdout.WriteString("\n")
+	return ptc.copier.Close(ctx)
+}
+
+// progressReader wraps a reader to track progress.
+type progressReader struct {
+	reader     fs.File
+	onProgress func(int64, string, int64)
+	copied     int64
+	fileName   string
+	fileSize   int64
+}
+
+func (pr *progressReader) Read(p []byte) (int, error) {
+	n, err := pr.reader.Read(p)
+	if n > 0 {
+		pr.copied += int64(n)
+		pr.onProgress(pr.copied, pr.fileName, pr.fileSize)
+	}
+	return n, err
+}
+
+func (pr *progressReader) Stat() (fs.FileInfo, error) {
+	return pr.reader.Stat()
+}
+
+func (pr *progressReader) Close() error {
+	return pr.reader.Close()
 }
 
 func (c *viamClient) copyFilesFromMachine(

cli/module_build.go

@@ -488,6 +488,7 @@ type reloadModuleArgs struct {
 	RestartOnly bool
 	NoBuild     bool
 	Local       bool
+	NoProgress  bool
 }
 
 // ReloadModuleAction builds a module, configures it on a robot, and starts or restarts it.
@@ -558,14 +559,14 @@ func reloadModuleAction(c *cli.Context, vc *viamClient, args reloadModuleArgs, l
 				return err
 			}
 			infof(c.App.Writer, "Copying %s to part %s", manifest.Build.Path, part.Part.Id)
-			args, err := getGlobalArgs(c)
+			globalArgs, err := getGlobalArgs(c)
 			if err != nil {
 				return err
 			}
 			dest := reloadingDestination(c, manifest)
 			err = vc.copyFilesToFqdn(
-				part.Part.Fqdn, args.Debug, false, false, []string{manifest.Build.Path},
-				dest, logging.NewLogger("reload"))
+				part.Part.Fqdn, globalArgs.Debug, false, false, []string{manifest.Build.Path},
+				dest, logging.NewLogger("reload"), args.NoProgress)
 			if err != nil {
 				if s, ok := status.FromError(err); ok && s.Code() == codes.PermissionDenied {
 					warningf(c.App.ErrWriter, "RDK couldn't write to the default file copy destination. "+