From 3699694fe189244523f013573a8699c2e197b676 Mon Sep 17 00:00:00 2001
From: Adin Schmahmann <adin.schmahmann@gmail.com>
Date: Fri, 10 Jul 2020 19:43:15 -0400
Subject: [PATCH] add ipfs dag stat command

---
 core/commands/dag/dag.go | 134 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 134 insertions(+)

diff --git a/core/commands/dag/dag.go b/core/commands/dag/dag.go
index 85588be2f144..141d9b4d4362 100644
--- a/core/commands/dag/dag.go
+++ b/core/commands/dag/dag.go
@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/ipfs/go-ipfs/core/commands/cmdenv"
+	"github.com/ipfs/go-ipfs/core/commands/e"
 	"github.com/ipfs/go-ipfs/core/coredag"
 	iface "github.com/ipfs/interface-go-ipfs-core"
 
@@ -19,6 +20,7 @@ import (
 	files "github.com/ipfs/go-ipfs-files"
 	ipld "github.com/ipfs/go-ipld-format"
 	mdag "github.com/ipfs/go-merkledag"
+	traverse "github.com/ipfs/go-merkledag/traverse"
 	ipfspath "github.com/ipfs/go-path"
 	"github.com/ipfs/interface-go-ipfs-core/options"
 	path "github.com/ipfs/interface-go-ipfs-core/path"
@@ -54,6 +56,7 @@ to deprecate and replace the existing 'ipfs object' command moving forward.
 		"resolve": DagResolveCmd,
 		"import":  DagImportCmd,
 		"export":  DagExportCmd,
+		"stat":    DagStatCmd,
 	},
 }
 
@@ -668,3 +671,134 @@ The output of blocks happens in strict DAG-traversal, first-seen, order.
 		},
 	},
 }
+
+type DagStat struct {
+	Size      uint64
+	NumBlocks int64
+}
+
+func (s *DagStat) String() string {
+	return fmt.Sprintf("Size: %d, NumBlocks: %d", s.Size, s.NumBlocks)
+}
+
+var DagStatCmd = &cmds.Command{
+	Helptext: cmds.HelpText{
+		Tagline: "Gets stats for a DAG",
+		ShortDescription: `
+'ipfs dag size' fetches a dag and returns various statistics about the DAG.
+Statistics include size and number of blocks.
+
+Note: This command skips duplicate blocks in reporting both size and the number of blocks
+`,
+	},
+	Arguments: []cmds.Argument{
+		cmds.StringArg("root", true, false, "CID of a DAG root to get statistics for").EnableStdin(),
+	},
+	Options: []cmds.Option{
+		cmds.BoolOption(progressOptionName, "p", "Return progressive data while reading through the DAG").WithDefault("true"),
+	},
+	Run: func(req *cmds.Request, res cmds.ResponseEmitter, env cmds.Environment) error {
+		progressive := req.Options[progressOptionName].(bool)
+
+		api, err := cmdenv.GetApi(env, req)
+		if err != nil {
+			return err
+		}
+
+		rp, err := api.ResolvePath(req.Context, path.New(req.Arguments[0]))
+		if err != nil {
+			return err
+		}
+
+		if len(rp.Remainder()) > 0 {
+			return fmt.Errorf("cannot return size for anything other than a DAG with a root CID")
+		}
+
+		obj, err := api.Dag().Get(req.Context, rp.Cid())
+		if err != nil {
+			return err
+		}
+
+		dagstats := &DagStat{}
+		err = traverse.Traverse(obj, traverse.Options{
+			DAG:   api.Dag(),
+			Order: traverse.DFSPre,
+			Func: func(current traverse.State) error {
+				stat, err := current.Node.Stat()
+				if err != nil {
+					return err
+				}
+
+				// stat.BlockSize gives the raw size of the data which is what we want here.
+				// However, node.Stat() is not implemented by any node types other than DagPB, which defines Size() as
+				// a field stored in the data instead of the raw size of the data itself.
+				//
+				// Therefore, we check if stat is defined and if so use `stat.BlockSize` and otherwise just rely on the
+				// Size() function.
+				if len(stat.Hash) == 0 {
+					size, err := current.Node.Size()
+					if err != nil {
+						return err
+					}
+					dagstats.Size += size
+				} else {
+					dagstats.Size += uint64(stat.BlockSize)
+				}
+				dagstats.NumBlocks++
+
+				if progressive {
+					if err := res.Emit(dagstats); err != nil {
+						return err
+					}
+				}
+				return nil
+			},
+			ErrFunc:        nil,
+			SkipDuplicates: true,
+		})
+		if err != nil {
+			return fmt.Errorf("error traversing DAG: %w", err)
+		}
+
+		if !progressive {
+			if err := res.Emit(dagstats); err != nil {
+				return err
+			}
+		}
+
+		return nil
+	},
+	Type: DagStat{},
+	PostRun: cmds.PostRunMap{
+		cmds.CLI: func(res cmds.Response, re cmds.ResponseEmitter) error {
+			var dagStats *DagStat
+			for {
+				v, err := res.Next()
+				if err != nil {
+					if err == io.EOF {
+						break
+					}
+					return err
+				}
+
+				out, ok := v.(*DagStat)
+				if !ok {
+					return e.TypeErr(out, v)
+				}
+				dagStats = out
+				fmt.Fprintf(os.Stderr, "%v\r", out)
+			}
+			return re.Emit(dagStats)
+		},
+	},
+	Encoders: cmds.EncoderMap{
+		cmds.Text: cmds.MakeTypedEncoder(func(req *cmds.Request, w io.Writer, event *DagStat) error {
+			_, err := fmt.Fprintf(
+				w,
+				"%v",
+				event,
+			)
+			return err
+		}),
+	},
+}