From d94749ac220cb8f35d063a02579b9eadf73b5cf2 Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Tue, 10 Oct 2023 21:24:36 +0200 Subject: refactor(tvix/store/protos): have Export accept root node We don't need the full PathInfo message, only the root node. Change-Id: I667045ed766875dfbf8ac126a50b02baa2df67a4 Reviewed-on: https://cl.tvl.fyi/c/depot/+/9604 Tested-by: BuildkiteCI Reviewed-by: edef --- tvix/store/protos/export.go | 28 +++++++++++------------ tvix/store/protos/export_test.go | 49 +++++++++++++++++----------------------- 2 files changed, 34 insertions(+), 43 deletions(-) (limited to 'tvix') diff --git a/tvix/store/protos/export.go b/tvix/store/protos/export.go index 8c4a27522977..889311efe70b 100644 --- a/tvix/store/protos/export.go +++ b/tvix/store/protos/export.go @@ -12,12 +12,12 @@ import ( type DirectoryLookupFn func([]byte) (*castorev1pb.Directory, error) type BlobLookupFn func([]byte) (io.ReadCloser, error) -// Export will traverse a given PathInfo structure, and write the contents -// in NAR format to the passed Writer. +// Export will traverse a given root node, and write the contents in NAR format +// to the passed Writer. // It uses directoryLookupFn and blobLookupFn to resolve references. func Export( w io.Writer, - pathInfo *PathInfo, + rootNode *castorev1pb.Node, directoryLookupFn DirectoryLookupFn, blobLookupFn BlobLookupFn, ) error { @@ -43,18 +43,17 @@ func Export( // peek at the pathInfo root and assemble the root node and write to writer // in the case of a regular file, we retrieve and write the contents, close and exit // in the case of a symlink, we write the symlink, close and exit - switch v := (pathInfo.GetNode().GetNode()).(type) { - case *castorev1pb.Node_File: + if fileNode := rootNode.GetFile(); fileNode != nil { rootHeader.Type = nar.TypeRegular - rootHeader.Size = int64(v.File.GetSize()) - rootHeader.Executable = v.File.GetExecutable() + rootHeader.Size = int64(fileNode.GetSize()) + rootHeader.Executable = fileNode.GetExecutable() err := narWriter.WriteHeader(rootHeader) if err != nil { return fmt.Errorf("unable to write root header: %w", err) } // if it's a regular file, retrieve and write the contents - blobReader, err := blobLookupFn(v.File.GetDigest()) + blobReader, err := blobLookupFn(fileNode.GetDigest()) if err != nil { return fmt.Errorf("unable to lookup blob: %w", err) } @@ -76,10 +75,9 @@ func Export( } return nil - - case *castorev1pb.Node_Symlink: + } else if symlinkNode := rootNode.GetSymlink(); symlinkNode != nil { rootHeader.Type = nar.TypeSymlink - rootHeader.LinkTarget = string(v.Symlink.GetTarget()) + rootHeader.LinkTarget = string(symlinkNode.GetTarget()) err := narWriter.WriteHeader(rootHeader) if err != nil { return fmt.Errorf("unable to write root header: %w", err) @@ -89,11 +87,9 @@ func Export( if err != nil { return fmt.Errorf("unable to close nar reader: %w", err) } - - return nil - case *castorev1pb.Node_Directory: + } else if directoryNode := rootNode.GetDirectory(); directoryNode != nil { // We have a directory at the root, look it up and put in on the stack. - directory, err := directoryLookupFn(v.Directory.Digest) + directory, err := directoryLookupFn(directoryNode.GetDigest()) if err != nil { return fmt.Errorf("unable to lookup directory: %w", err) } @@ -108,6 +104,8 @@ func Export( if err != nil { return fmt.Errorf("error writing header: %w", err) } + } else { + panic("invalid type") // unreachable } // as long as the stack is not empty, we keep running. diff --git a/tvix/store/protos/export_test.go b/tvix/store/protos/export_test.go index d45dd1ee7cf3..9b1901758f83 100644 --- a/tvix/store/protos/export_test.go +++ b/tvix/store/protos/export_test.go @@ -30,21 +30,18 @@ func mustBlobDigest(r io.Reader) []byte { } func TestSymlink(t *testing.T) { - pathInfo := &storev1pb.PathInfo{ - - Node: &castorev1pb.Node{ - Node: &castorev1pb.Node_Symlink{ - Symlink: &castorev1pb.SymlinkNode{ - Name: []byte("doesntmatter"), - Target: []byte("/nix/store/somewhereelse"), - }, + node := &castorev1pb.Node{ + Node: &castorev1pb.Node_Symlink{ + Symlink: &castorev1pb.SymlinkNode{ + Name: []byte("doesntmatter"), + Target: []byte("/nix/store/somewhereelse"), }, }, } var buf bytes.Buffer - err := storev1pb.Export(&buf, pathInfo, func([]byte) (*castorev1pb.Directory, error) { + err := storev1pb.Export(&buf, node, func([]byte) (*castorev1pb.Directory, error) { panic("no directories expected") }, func([]byte) (io.ReadCloser, error) { panic("no files expected") @@ -70,22 +67,20 @@ func TestRegular(t *testing.T) { 0x65, 0x2b, } - pathInfo := &storev1pb.PathInfo{ - Node: &castorev1pb.Node{ - Node: &castorev1pb.Node_File{ - File: &castorev1pb.FileNode{ - Name: []byte("doesntmatter"), - Digest: BLAKE3_DIGEST_0X01, - Size: 1, - Executable: false, - }, + node := &castorev1pb.Node{ + Node: &castorev1pb.Node_File{ + File: &castorev1pb.FileNode{ + Name: []byte("doesntmatter"), + Digest: BLAKE3_DIGEST_0X01, + Size: 1, + Executable: false, }, }, } var buf bytes.Buffer - err := storev1pb.Export(&buf, pathInfo, func([]byte) (*castorev1pb.Directory, error) { + err := storev1pb.Export(&buf, node, func([]byte) (*castorev1pb.Directory, error) { panic("no directories expected") }, func(blobRef []byte) (io.ReadCloser, error) { if !bytes.Equal(blobRef, BLAKE3_DIGEST_0X01) { @@ -115,21 +110,19 @@ func TestEmptyDirectory(t *testing.T) { } emptyDirectoryDigest := mustDirectoryDigest(emptyDirectory) - pathInfo := &storev1pb.PathInfo{ - Node: &castorev1pb.Node{ - Node: &castorev1pb.Node_Directory{ - Directory: &castorev1pb.DirectoryNode{ - Name: []byte("doesntmatter"), - Digest: emptyDirectoryDigest, - Size: 0, - }, + node := &castorev1pb.Node{ + Node: &castorev1pb.Node_Directory{ + Directory: &castorev1pb.DirectoryNode{ + Name: []byte("doesntmatter"), + Digest: emptyDirectoryDigest, + Size: 0, }, }, } var buf bytes.Buffer - err := storev1pb.Export(&buf, pathInfo, func(directoryRef []byte) (*castorev1pb.Directory, error) { + err := storev1pb.Export(&buf, node, func(directoryRef []byte) (*castorev1pb.Directory, error) { if !bytes.Equal(directoryRef, emptyDirectoryDigest) { panic("unexpected directoryRef") } -- cgit 1.4.1