mdebray/tvl-depot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor(tvix/store/protos): have Export accept root node

Browse source 
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 <edef@edef.eu>
This commit is contained in:
Florian Klink 2023-10-10 21:24:36 +02:00 committed by flokli
parent fe963ae0a3
commit d94749ac22
2 changed files with 34 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
tvix/store/protos/export.go
View file

@ -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.

49
tvix/store/protos/export_test.go
View file

@ -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")
}