feat(storage/transfermanager): add SkipIfExists option

storage/transfermanager/downloader.go

@@ -160,15 +160,22 @@ func (d *Downloader) DownloadDirectory(ctx context.Context, input *DownloadDirec
  }
 
  // Check if the file exists.
- // TODO: add skip option.
+ fileExists := false
+
  filePath := filepath.Join(input.LocalDirectory, attrs.Name)
  if _, err := os.Stat(filePath); err == nil {
- return fmt.Errorf("transfermanager: failed to create file(%q): %w", filePath, os.ErrExist)
+ fileExists = true
+ if !d.config.skipIfExists {
+ return fmt.Errorf("transfermanager: failed to create file(%q): %w", filePath, os.ErrExist)
+ }
  } else if !errors.Is(err, os.ErrNotExist) {
+ // Encountered an error other than file does not exist.
  return fmt.Errorf("transfermanager: failed to create file(%q): %w", filePath, err)
  }
 
- objectsToQueue = append(objectsToQueue, attrs.Name)
+ if !(d.config.skipIfExists && fileExists) {
+ objectsToQueue = append(objectsToQueue, attrs.Name)
+ }
  }
 
  outs := make(chan DownloadOutput, len(objectsToQueue))

storage/transfermanager/integration_test.go

@@ -28,6 +28,7 @@ import (
  "math/rand"
  "os"
  "path/filepath"
+ "strings"
  "sync"
  "testing"
  "time"
@@ -295,6 +296,75 @@ func TestIntegration_DownloadDirectoryAsync(t *testing.T) {
  t.Errorf("unexpected file %q in dir", entry.Name())
  }
  }
+
+ // Now attempt to download the entire directory.
+ // The existing files should be skipped.
+ callbacks := make(chan bool)
+
+ d, err = NewDownloader(c, WithWorkers(2), SkipIfExists())
+ if err != nil {
+ t.Fatalf("NewDownloader: %v", err)
+ }
+
+ if err := d.DownloadDirectory(ctx, &DownloadDirectoryInput{
+ Bucket: tb.bucket,
+ LocalDirectory: localDir,
+ OnObjectDownload: func(got *DownloadOutput) {
+ callbacks <- true
+
+ if got.Err != nil {
+ t.Errorf("result.Err: %v", got.Err)
+ }
+
+ if strings.EqualFold(got.Object, "dir/objA") {
+ t.Errorf("should have skipped download of object %s", got.Object)
+ }
+
+ if got, want := got.Attrs.Size, tb.objectSizes[got.Object]; want != got {
+ t.Errorf("expected object size %d, got %d", want, got)
+ }
+
+ path := filepath.Join(localDir, got.Object)
+ f, err := os.Open(path)
+ if err != nil {
+ t.Errorf("os.Open(%q): %v", path, err)
+ }
+ defer f.Close()
+
+ b := bytes.NewBuffer(make([]byte, 0, got.Attrs.Size))
+ if _, err := io.Copy(b, f); err != nil {
+ t.Errorf("io.Copy: %v", err)
+ }
+
+ if wantCRC, gotCRC := tb.contentHashes[got.Object], crc32c(b.Bytes()); gotCRC != wantCRC {
+ t.Errorf("object(%q) at filepath(%q): content crc32c does not match; got: %v, expected: %v", got.Object, path, gotCRC, wantCRC)
+ }
+ },
+ }); err != nil {
+ t.Errorf("d.DownloadDirectory: %v", err)
+ }
+
+ gotCallbacks := 0
+ done := make(chan bool)
+ go func() {
+ for {
+ select {
+ case <-done:
+ break
+ case <-callbacks:
+ gotCallbacks++
+ }
+ }
+ }()
+
+ if _, err := d.WaitAndClose(); err != nil {
+ t.Errorf("d.WaitAndClose: %v", err)
+ }
+ done <- true
+
+ if want, got := len(tb.objects)-wantObjs, gotCallbacks; want != got {
+ t.Errorf("expected to receive %d callbacks, got %d", want, got)
+ }
  })
 }
 

storage/transfermanager/option.go

@@ -91,6 +91,21 @@ func (wps withPartSize) apply(tm *transferManagerConfig) {
  tm.partSize = wps.partSize
 }
 
+// SkipIfExists returns a TransferManagerOption that will not download files
+// that already exist in the local directory.
+//
+// By default, if a file already exists the operation will abort and return an error.
+func SkipIfExists() Option {
+ return &skipIfExists{}
+}
+
+type skipIfExists struct {
+}
+
+func (sie skipIfExists) apply(tm *transferManagerConfig) {
+ tm.skipIfExists = true
+}
+
 type transferManagerConfig struct {
  // Workers in thread pool; default numCPU/2 based on previous benchmarks?
  numWorkers int
@@ -107,6 +122,10 @@ type transferManagerConfig struct {
  // If true, callbacks are used instead of returning results synchronously
  // in a slice at the end.
  asynchronous bool
+
+ // If true, files that already exist in the local directory will not be
+ // downloaded.
+ skipIfExists bool
 }
 
 func defaultTransferManagerConfig() *transferManagerConfig {