gitpipe: Add WithDiskSize option for CatfileInfo

The %(objectsize:disk) formatting directive is useful for calculating
the amount of disk space that an object takes up. Add this as an option.

2 files changed, 76 insertions, 1 deletions

diff --git a/internal/git/gitpipe/catfile_info.go b/internal/git/gitpipe/catfile_info.go
index 3aec77cca..07e29525b 100644
--- a/internal/git/gitpipe/catfile_info.go
+++ b/internal/git/gitpipe/catfile_info.go
@@ -27,6 +27,7 @@ type CatfileInfoResult struct {
 
 type catfileInfoConfig struct {
 	skipResult func(*catfile.ObjectInfo) bool
+	diskUsage  bool
 }
 
 // CatfileInfoOption is an option for the CatfileInfo and CatfileInfoAllObjects pipeline steps.
@@ -41,6 +42,14 @@ func WithSkipCatfileInfoResult(skipResult func(*catfile.ObjectInfo) bool) Catfil
 	}
 }
 
+// WithDiskUsageSize will cause the size of the object to be returned to be the
+// size it takes up on disk. This value will override the existing size field.
+func WithDiskUsageSize() CatfileInfoOption {
+	return func(cfg *catfileInfoConfig) {
+		cfg.diskUsage = true
+	}
+}
+
 type catfileInfoRequest struct {
 	objectID   git.ObjectID
 	objectName []byte
@@ -193,12 +202,19 @@ func CatfileInfoAllObjects(
 			return
 		}
 
+		batchCheckOption := git.Flag{Name: "--batch-check"}
+
+		if cfg.diskUsage {
+			batchCheckOption.Name = batchCheckOption.Name +
+				fmt.Sprintf("=%%(objectname) %%(objecttype) %%(objectsize:disk)")
+		}
+
 		var stderr bytes.Buffer
 		cmd, err := repo.Exec(ctx, git.SubCmd{
 			Name: "cat-file",
 			Flags: []git.Option{
+				batchCheckOption,
 				git.Flag{Name: "--batch-all-objects"},
-				git.Flag{Name: "--batch-check"},
 				git.Flag{Name: "--buffer"},
 				git.Flag{Name: "--unordered"},
 			},
diff --git a/internal/git/gitpipe/catfile_info_test.go b/internal/git/gitpipe/catfile_info_test.go
index 94cf6a78d..317c93ba3 100644
--- a/internal/git/gitpipe/catfile_info_test.go
+++ b/internal/git/gitpipe/catfile_info_test.go
@@ -335,3 +335,62 @@ func TestCatfileInfoAllObjects(t *testing.T) {
 		}, it.Result())
 	})
 }
+
+func TestCatfileInfo_WithDiskUsageSize(t *testing.T) {
+	t.Parallel()
+
+	cfg := testcfg.Build(t)
+	ctx := testhelper.Context(t)
+
+	repoProto, repoPath := gittest.CreateRepository(ctx, t, cfg, gittest.CreateRepositoryConfig{
+		SkipCreationViaService: true,
+	})
+	repo := localrepo.NewTestRepo(t, cfg, repoProto)
+
+	tree1 := gittest.WriteTree(t, cfg, repoPath, []gittest.TreeEntry{
+		{
+			Path: "foobar",
+			Mode: "100644",
+			OID:  gittest.WriteBlob(t, cfg, repoPath, bytes.Repeat([]byte("a"), 100)),
+		},
+	})
+	initialCommitID := gittest.WriteCommit(t, cfg, repoPath, gittest.WithTree(tree1))
+
+	tree2 := gittest.WriteTree(t, cfg, repoPath, []gittest.TreeEntry{
+		{
+			Path: "foobar",
+			Mode: "100644",
+			// take advantage of compression
+			OID: gittest.WriteBlob(t, cfg, repoPath, append(bytes.Repeat([]byte("a"), 100),
+				'\n',
+				'b',
+			)),
+		},
+	})
+	gittest.WriteCommit(
+		t,
+		cfg,
+		repoPath,
+		gittest.WithTree(tree2),
+		gittest.WithParents(initialCommitID),
+		gittest.WithBranch("master"),
+	)
+
+	gittest.Exec(t, cfg, "-C", repoPath, "gc")
+
+	itWithoutDiskSize := CatfileInfoAllObjects(ctx, repo)
+	itWithDiskSize := CatfileInfoAllObjects(ctx, repo, WithDiskUsageSize())
+
+	var totalWithoutDiskSize, totalWithDiskSize int64
+	for itWithoutDiskSize.Next() {
+		totalWithoutDiskSize += itWithoutDiskSize.Result().ObjectSize()
+	}
+	require.NoError(t, itWithoutDiskSize.Err())
+
+	for itWithDiskSize.Next() {
+		totalWithDiskSize += itWithDiskSize.Result().ObjectSize()
+	}
+	require.NoError(t, itWithDiskSize.Err())
+
+	require.Less(t, totalWithDiskSize, totalWithoutDiskSize)
+}