Add simple byte-counting export progress-bar

core/commands/dag/dag.go

@@ -1,10 +1,13 @@
 package dagcmd
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"math"
+	"os"
 	"strings"
+	"time"
 
 	"github.com/ipfs/go-ipfs/core/commands/cmdenv"
 	"github.com/ipfs/go-ipfs/core/coredag"
@@ -22,6 +25,12 @@ import (
 	//gipfree "github.com/ipld/go-ipld-prime/impl/free"
 	//gipselector "github.com/ipld/go-ipld-prime/traversal/selector"
 	//gipselectorbuilder "github.com/ipld/go-ipld-prime/traversal/selector/builder"
+
+	"gopkg.in/cheggaaa/pb.v1"
+)
+
+const (
+	progressOptionName = "progress"
 )
 
 var DagCmd = &cmds.Command{
@@ -260,6 +269,9 @@ The output of blocks happens in strict DAG-traversal, first-seen, order.
 	Arguments: []cmds.Argument{
 		cmds.StringArg("root", true, false, "CID of a root to recursively export").EnableStdin(),
 	},
+	Options: []cmds.Option{
+		cmds.BoolOption(progressOptionName, "p", "Display progress on CLI. Defaults to true when STDERR is a TTY."),
+	},
 	Run: func(req *cmds.Request, res cmds.ResponseEmitter, env cmds.Environment) error {
 
 		c, err := cid.Decode(req.Arguments[0])
@@ -343,4 +355,53 @@ The output of blocks happens in strict DAG-traversal, first-seen, order.
 
 		return err
 	},
+	PostRun: cmds.PostRunMap{
+		cmds.CLI: func(res cmds.Response, re cmds.ResponseEmitter) error {
+
+			var showProgress bool
+			val, specified := res.Request().Options[progressOptionName]
+			if !specified {
+				// default based on TTY availability
+				errStat, _ := os.Stderr.Stat()
+				if 0 != (errStat.Mode() & os.ModeCharDevice) {
+					showProgress = true
+				}
+			} else if val.(bool) {
+				showProgress = true
+			}
+
+			// simple passthrough
+			if !showProgress {
+				return cmds.Copy(re, res)
+			}
+
+			// We will be supplied one or more readers, simply wrap them with a progress bar
+			for {
+				v, err := res.Next()
+				if err == io.EOF {
+					return re.Close()
+				} else if err != nil {
+					return re.CloseWithError(err)
+				}
+
+				r, ok := v.(io.Reader)
+				if !ok {
+					// some sort of encoded response, this should not be happening
+					return errors.New("unexpected non-stream passed to PostRun: can't happen")
+				}
+
+				if err := func() error {
+					bar := pb.New64(0).SetUnits(pb.U_BYTES)
+					bar.Output = os.Stderr
+					bar.ShowSpeed = true
+					bar.ShowElapsedTime = true
+					bar.RefreshRate = 500 * time.Millisecond
+					defer bar.Finish()
+					return re.Emit(bar.Start().NewProxyReader(r))
+				}(); err != nil {
+					return err
+				}
+			}
+		},
+	},
 }