From 607534cdd5edf2d15d3de891cf6a0b6cbaa7d545 Mon Sep 17 00:00:00 2001
From: Brenna N Epp <brennae@google.com>
Date: Tue, 24 Sep 2024 09:53:08 -0600
Subject: [PATCH] feat(storage/transfermanager): add option to StripPrefix on
 directory download (#10894)

---
 storage/transfermanager/downloader.go       | 16 ++++++++++++++--
 storage/transfermanager/integration_test.go |  9 +++++++--
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/storage/transfermanager/downloader.go b/storage/transfermanager/downloader.go
index a15c4a418477..d7b7d5a7f5aa 100644
--- a/storage/transfermanager/downloader.go
+++ b/storage/transfermanager/downloader.go
@@ -25,6 +25,7 @@ import (
 	"math"
 	"os"
 	"path/filepath"
+	"strings"
 	"sync"
 	"time"
 
@@ -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)
@@ -750,6 +756,12 @@ 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 the local file on
+	// download from GCS. For example, if you have an object in GCS called
+	// "mydirectory/a/file", and StripPrefix is set to "mydirectory/", the file
+	// will be downloaded to "{LocalDirectory}/a/file". Optional.
+	StripPrefix string
+
 	// Prefix is the prefix filter to download objects whose names begin with this.
 	// Optional.
 	Prefix string
diff --git a/storage/transfermanager/integration_test.go b/storage/transfermanager/integration_test.go
index 1d06b5fc67bb..484798a92525 100644
--- a/storage/transfermanager/integration_test.go
+++ b/storage/transfermanager/integration_test.go
@@ -28,6 +28,7 @@ import (
 	"math/rand"
 	"os"
 	"path/filepath"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -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++
@@ -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)
@@ -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)