From fd3e04b63ead7371a30eb2dc307c25c5cbb948fd Mon Sep 17 00:00:00 2001
From: Chris Cotter <cjcotter@google.com>
Date: Fri, 7 Jun 2024 18:28:20 -0400
Subject: [PATCH] feat(storage/transfermanager): add Range to output (#10347)

DownloadOutput should include the range so if multiple ranges were requested for the same file and the download failed, the caller can identify which offset the error occurred on.
---
 storage/transfermanager/downloader.go       |  3 ++-
 storage/transfermanager/integration_test.go | 10 +++++++++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/storage/transfermanager/downloader.go b/storage/transfermanager/downloader.go
index 88ecc2c6e769..484c5334e783 100644
--- a/storage/transfermanager/downloader.go
+++ b/storage/transfermanager/downloader.go
@@ -243,7 +243,7 @@ type DownloadObjectInput struct {
 // If timeout is less than 0, no timeout is set.
 // TODO: download a single shard instead of the entire object.
 func (in *DownloadObjectInput) downloadShard(client *storage.Client, timeout time.Duration) (out *DownloadOutput) {
-	out = &DownloadOutput{Bucket: in.Bucket, Object: in.Object}
+	out = &DownloadOutput{Bucket: in.Bucket, Object: in.Object, Range: in.Range}
 
 	// Set timeout.
 	ctx := in.ctx
@@ -303,6 +303,7 @@ func (in *DownloadObjectInput) downloadShard(client *storage.Client, timeout tim
 type DownloadOutput struct {
 	Bucket string
 	Object string
+	Range  *DownloadRange             // requested range, if it was specified
 	Err    error                      // error occurring during download
 	Attrs  *storage.ReaderObjectAttrs // attributes of downloaded object, if successful
 }
diff --git a/storage/transfermanager/integration_test.go b/storage/transfermanager/integration_test.go
index ae5eb873221c..c4a5cb45b5c8 100644
--- a/storage/transfermanager/integration_test.go
+++ b/storage/transfermanager/integration_test.go
@@ -484,6 +484,10 @@ func TestIntegration_DownloadShard(t *testing.T) {
 				want: &DownloadOutput{
 					Bucket: tb.bucket,
 					Object: objectName,
+					Range: &DownloadRange{
+						Offset: tb.objectSize - 5,
+						Length: -1,
+					},
 					Attrs: &storage.ReaderObjectAttrs{
 						Size:            tb.objectSize,
 						StartOffset:     tb.objectSize - 5,
@@ -582,8 +586,12 @@ func TestIntegration_DownloadShard(t *testing.T) {
 					t.Errorf("wanted bucket %q, object %q, got: %q, %q", test.want.Bucket, test.want.Object, got.Bucket, got.Object)
 				}
 
+				if diff := cmp.Diff(got.Range, test.want.Range); diff != "" {
+					t.Errorf("DownloadOutput.Range: got(-) vs. want(+): %v", diff)
+				}
+
 				if diff := cmp.Diff(got.Attrs, test.want.Attrs); diff != "" {
-					t.Errorf("diff got(-) vs. want(+): %v", diff)
+					t.Errorf("DownloadOutput.Attrs: got(-) vs. want(+): %v", diff)
 				}
 
 				if !errorIs(got.Err, test.want.Err) {