Skip to content

Commit

Permalink
feat(storage/transfermanager): add option to StripPrefix on directory…
Browse files Browse the repository at this point in the history
… download
  • Loading branch information
BrennaEpp committed Sep 20, 2024
1 parent ac256c5 commit 06fe327
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 4 deletions.
14 changes: 12 additions & 2 deletions storage/transfermanager/downloader.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import (
"math"
"os"
"path/filepath"
"strings"
"sync"
"time"

Expand Down Expand Up @@ -175,8 +176,13 @@ func (d *Downloader) DownloadDirectory(ctx context.Context, input *DownloadDirec
inputs := make([]DownloadObjectInput, 0, len(objectsToQueue))

for _, object := range objectsToQueue {
objDirectory := filepath.Join(input.LocalDirectory, filepath.Dir(object))
filePath := filepath.Join(input.LocalDirectory, object)
objectWithoutPrefix := object
if len(input.StripPrefix) > 0 {
objectWithoutPrefix, _ = strings.CutPrefix(object, input.StripPrefix)
}

objDirectory := filepath.Join(input.LocalDirectory, filepath.Dir(objectWithoutPrefix))
filePath := filepath.Join(input.LocalDirectory, objectWithoutPrefix)

// Make sure all directories in the object path exist.
err := os.MkdirAll(objDirectory, fs.ModeDir|fs.ModePerm)
Expand Down Expand Up @@ -750,6 +756,10 @@ type DownloadDirectoryInput struct {
// The directory will be created if it does not already exist. Required.
LocalDirectory string

// StripPrefix is a prefix to be removed from the path of an object on
// download from GCS. Optional.
StripPrefix string

// Prefix is the prefix filter to download objects whose names begin with this.
// Optional.
Prefix string
Expand Down
9 changes: 7 additions & 2 deletions storage/transfermanager/integration_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ import (
"math/rand"
"os"
"path/filepath"
"strings"
"sync"
"testing"
"time"
Expand Down Expand Up @@ -100,6 +101,7 @@ func TestIntegration_DownloadDirectory(t *testing.T) {
if err := d.DownloadDirectory(ctx, &DownloadDirectoryInput{
Bucket: tb.bucket,
LocalDirectory: localDir,
StripPrefix: "dir/",
OnObjectDownload: func(got *DownloadOutput) {
callbackMu.Lock()
numCallbacks++
Expand All @@ -113,7 +115,9 @@ func TestIntegration_DownloadDirectory(t *testing.T) {
t.Errorf("expected object size %d, got %d", want, got)
}

path := filepath.Join(localDir, got.Object)
objectWithoutPrefix, _ := strings.CutPrefix(got.Object, "dir/")
path := filepath.Join(localDir, objectWithoutPrefix)

f, err := os.Open(path)
if err != nil {
t.Errorf("os.Open(%q): %v", path, err)
Expand Down Expand Up @@ -156,7 +160,8 @@ func TestIntegration_DownloadDirectory(t *testing.T) {
t.Errorf("expected object size %d, got %d", want, got)
}

path := filepath.Join(localDir, got.Object)
objectWithoutPrefix, _ := strings.CutPrefix(got.Object, "dir/")
path := filepath.Join(localDir, objectWithoutPrefix)
f, err := os.Open(path)
if err != nil {
t.Errorf("os.Open(%q): %v", path, err)
Expand Down

0 comments on commit 06fe327

Please sign in to comment.