Update snapshot summary on rewrite

Signed-off-by: Alex Johnson <hello@alex-johnson.net>

changelog/unreleased/issue-4902

@@ -0,0 +1,6 @@
+Enhancement: Update snapshot summary on rewrite
+
+Restic now recalculates the total number of files and bytes processed when files are excluded during rewrite.
+
+https://github.com/restic/restic/issues/4902
+https://github.com/restic/restic/pull/4905

cmd/restic/cmd_rewrite.go

@@ -134,20 +134,29 @@ func rewriteSnapshot(ctx context.Context, repo *repository.Repository, sn *resti
 			return true
 		}
 
-		rewriter := walker.NewTreeRewriter(walker.RewriteOpts{
-			RewriteNode: func(node *restic.Node, path string) *restic.Node {
-				if selectByName(path) {
-					return node
-				}
-				Verbosef(fmt.Sprintf("excluding %s\n", path))
-				return nil
-			},
-			DisableNodeCache: true,
-		})
+		rewriteNode := func(node *restic.Node, path string) *restic.Node {
+			if selectByName(path) {
+				return node
+			}
+			Verbosef(fmt.Sprintf("excluding %s\n", path))
+			return nil
+		}
+
+		rewriter, querySize := walker.NewSnapshotSizeRewriter(rewriteNode)
 
 		filter = func(ctx context.Context, sn *restic.Snapshot) (restic.ID, error) {
-			return rewriter.RewriteTree(ctx, repo, "/", *sn.Tree)
+			id, err := rewriter.RewriteTree(ctx, repo, "/", *sn.Tree)
+			if err != nil {
+				return restic.ID{}, err
+			}
+			ss := querySize()
+			if sn.Summary != nil {
+				sn.Summary.TotalFilesProcessed = ss.FileCount
+				sn.Summary.TotalBytesProcessed = ss.FileSize
+			}
+			return id, err
 		}
+
 	} else {
 		filter = func(_ context.Context, sn *restic.Snapshot) (restic.ID, error) {
 			return *sn.Tree, nil

cmd/restic/cmd_rewrite_integration_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/restic/restic/internal/restic"
 	rtest "github.com/restic/restic/internal/test"
+	"github.com/restic/restic/internal/ui"
 )
 
 func testRunRewriteExclude(t testing.TB, gopts GlobalOptions, excludes []string, forget bool, metadata snapshotMetadataArgs) {
@@ -33,6 +34,24 @@ func createBasicRewriteRepo(t testing.TB, env *testEnvironment) restic.ID {
 	return snapshotIDs[0]
 }
 
+func getSnapshot(t testing.TB, snapshotID restic.ID, env *testEnvironment) *restic.Snapshot {
+	t.Helper()
+
+	ctx, repo, unlock, err := openWithReadLock(context.TODO(), env.gopts, false)
+	rtest.OK(t, err)
+	defer unlock()
+
+	snapshots, err := restic.TestLoadAllSnapshots(ctx, repo, nil)
+	rtest.OK(t, err)
+
+	for _, s := range snapshots {
+		if *s.ID() == snapshotID {
+			return s
+		}
+	}
+	return nil
+}
+
 func TestRewrite(t *testing.T) {
 	env, cleanup := withTestEnvironment(t)
 	defer cleanup()
@@ -63,10 +82,21 @@ func TestRewriteReplace(t *testing.T) {
 	defer cleanup()
 	snapshotID := createBasicRewriteRepo(t, env)
 
+	snapshot := getSnapshot(t, snapshotID, env)
+
 	// exclude some data
 	testRunRewriteExclude(t, env.gopts, []string{"3"}, true, snapshotMetadataArgs{Hostname: "", Time: ""})
+	bytesExcluded, err := ui.ParseBytes("16K")
+	rtest.OK(t, err)
+
 	newSnapshotIDs := testListSnapshots(t, env.gopts, 1)
 	rtest.Assert(t, snapshotID != newSnapshotIDs[0], "snapshot id should have changed")
+
+	newSnapshot := getSnapshot(t, newSnapshotIDs[0], env)
+
+	rtest.Equals(t, snapshot.Summary.TotalFilesProcessed-1, newSnapshot.Summary.TotalFilesProcessed, "snapshot file count should have changed")
+	rtest.Equals(t, snapshot.Summary.TotalBytesProcessed-uint64(bytesExcluded), newSnapshot.Summary.TotalBytesProcessed, "snapshot size should have changed")
+
 	// check forbids unused blobs, thus remove them first
 	testRunPrune(t, env.gopts, PruneOptions{MaxUnused: "0"})
 	testRunCheck(t, env.gopts)

internal/walker/rewriter.go

@@ -11,6 +11,12 @@ import (
 
 type NodeRewriteFunc func(node *restic.Node, path string) *restic.Node
 type FailedTreeRewriteFunc func(nodeID restic.ID, path string, err error) (restic.ID, error)
+type QueryRewrittenSizeFunc func() SnapshotSize
+
+type SnapshotSize struct {
+	FileCount uint
+	FileSize  uint64
+}
 
 type RewriteOpts struct {
 	// return nil to remove the node
@@ -52,6 +58,29 @@ func NewTreeRewriter(opts RewriteOpts) *TreeRewriter {
 	return rw
 }
 
+func NewSnapshotSizeRewriter(rewriteNode NodeRewriteFunc) (*TreeRewriter, QueryRewrittenSizeFunc) {
+	var count uint
+	var size uint64
+
+	t := NewTreeRewriter(RewriteOpts{
+		RewriteNode: func(node *restic.Node, path string) *restic.Node {
+			node = rewriteNode(node, path)
+			if node != nil && node.Type == "file" {
+				count++
+				size += node.Size
+			}
+			return node
+		},
+		DisableNodeCache: true,
+	})
+
+	ss := func() SnapshotSize {
+		return SnapshotSize{count, size}
+	}
+
+	return t, ss
+}
+
 type BlobLoadSaver interface {
 	restic.BlobSaver
 	restic.BlobLoader

internal/walker/rewriter_test.go

@@ -303,6 +303,60 @@ func TestRewriter(t *testing.T) {
 	}
 }
 
+func TestSnapshotSizeQuery(t *testing.T) {
+	tree := TestTree{
+		"foo": TestFile{Size: 21},
+		"bar": TestFile{Size: 21},
+		"subdir": TestTree{
+			"subfile": TestFile{Size: 21},
+		},
+	}
+	newTree := TestTree{
+		"foo": TestFile{Size: 42},
+		"subdir": TestTree{
+			"subfile": TestFile{Size: 42},
+		},
+	}
+	t.Run("", func(t *testing.T) {
+		repo, root := BuildTreeMap(tree)
+		expRepo, expRoot := BuildTreeMap(newTree)
+		modrepo := WritableTreeMap{repo}
+
+		ctx, cancel := context.WithCancel(context.TODO())
+		defer cancel()
+
+		rewriteNode := func(node *restic.Node, path string) *restic.Node {
+			if path == "/bar" {
+				return nil
+			}
+			if node.Type == "file" {
+				node.Size += 21
+			}
+			return node
+		}
+		rewriter, querySize := NewSnapshotSizeRewriter(rewriteNode)
+		newRoot, err := rewriter.RewriteTree(ctx, modrepo, "/", root)
+		if err != nil {
+			t.Error(err)
+		}
+
+		ss := querySize()
+
+		test.Equals(t, uint(2), ss.FileCount, "snapshot file count mismatch")
+		test.Equals(t, uint64(84), ss.FileSize, "snapshot size mismatch")
+
+		// verifying against the expected tree root also implicitly checks the structural integrity
+		if newRoot != expRoot {
+			t.Error("hash mismatch")
+			fmt.Println("Got")
+			modrepo.Dump()
+			fmt.Println("Expected")
+			WritableTreeMap{expRepo}.Dump()
+		}
+	})
+
+}
+
 func TestRewriterFailOnUnknownFields(t *testing.T) {
 	tm := WritableTreeMap{TreeMap{}}
 	node := []byte(`{"nodes":[{"name":"subfile","type":"file","mtime":"0001-01-01T00:00:00Z","atime":"0001-01-01T00:00:00Z","ctime":"0001-01-01T00:00:00Z","uid":0,"gid":0,"content":null,"unknown_field":42}]}`)