Add multipart presigned URL download support to MultipartS3AsyncClient

services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/MultipartAsyncPresignedUrlExtension.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.services.s3.internal.multipart;
+
+import java.util.concurrent.CompletableFuture;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.async.AsyncResponseTransformer;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
+import software.amazon.awssdk.services.s3.model.GetObjectResponse;
+import software.amazon.awssdk.services.s3.presignedurl.AsyncPresignedUrlExtension;
+import software.amazon.awssdk.services.s3.presignedurl.model.PresignedUrlDownloadRequest;
+import software.amazon.awssdk.utils.Validate;
+
+/**
+ * An {@link AsyncPresignedUrlExtension} that automatically converts presigned URL downloads
+ * to multipart downloads.
+ */
+@SdkInternalApi
+public class MultipartAsyncPresignedUrlExtension implements AsyncPresignedUrlExtension {
+    private final PresignedUrlDownloadHelper downloadHelper;
+
+    public MultipartAsyncPresignedUrlExtension(
+            S3AsyncClient s3AsyncClient,
+            AsyncPresignedUrlExtension asyncPresignedUrlExtension,
+            long bufferSizeInBytes,
+            long partSizeInBytes) {
+        Validate.paramNotNull(s3AsyncClient, "s3AsyncClient");
+        Validate.paramNotNull(asyncPresignedUrlExtension, "asyncPresignedUrlExtension");
+        this.downloadHelper = new PresignedUrlDownloadHelper(
+                s3AsyncClient,
+                asyncPresignedUrlExtension,
+                bufferSizeInBytes,
+                partSizeInBytes);
+    }
+
+    @Override
+    public <ReturnT> CompletableFuture<ReturnT> getObject(
+            PresignedUrlDownloadRequest presignedUrlDownloadRequest,
+            AsyncResponseTransformer<GetObjectResponse, ReturnT> asyncResponseTransformer) {
+        Validate.paramNotNull(presignedUrlDownloadRequest, "presignedUrlDownloadRequest");
+        Validate.paramNotNull(asyncResponseTransformer, "asyncResponseTransformer");
+        return downloadHelper.downloadObject(presignedUrlDownloadRequest, asyncResponseTransformer);
+    }
+}

services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/MultipartS3AsyncClient.java

@@ -53,6 +53,8 @@ public final class MultipartS3AsyncClient extends DelegatingS3AsyncClient {
     private final CopyObjectHelper copyObjectHelper;
     private final DownloadObjectHelper downloadObjectHelper;
     private final boolean checksumEnabled;
+    private final long apiCallBufferSize;
+    private final long minPartSizeInBytes;
 
     private MultipartS3AsyncClient(S3AsyncClient delegate, MultipartConfiguration multipartConfiguration,
                                    boolean checksumEnabled) {
@@ -63,6 +65,8 @@ private MultipartS3AsyncClient(S3AsyncClient delegate, MultipartConfiguration mu
         long minPartSizeInBytes = resolver.minimalPartSizeInBytes();
         long threshold = resolver.thresholdInBytes();
         long apiCallBufferSize = resolver.apiCallBufferSize();
+        this.apiCallBufferSize = apiCallBufferSize;
+        this.minPartSizeInBytes = minPartSizeInBytes;
         mpuHelper = new UploadObjectHelper(delegate, resolver);
         copyObjectHelper = new CopyObjectHelper(delegate, minPartSizeInBytes, threshold);
         downloadObjectHelper = new DownloadObjectHelper(delegate, apiCallBufferSize);
@@ -114,7 +118,11 @@ protected <T extends S3Request, ReturnT> CompletableFuture<ReturnT> invokeOperat
 
     @Override
     public AsyncPresignedUrlExtension presignedUrlExtension() {
-        // TODO: Implement presigned URL extension support for multipart client
-        throw new UnsupportedOperationException("Presigned URL extension is not supported for multipart client");
+        AsyncPresignedUrlExtension delegateExtension = ((S3AsyncClient) delegate()).presignedUrlExtension();
+        return new MultipartAsyncPresignedUrlExtension(
+            (S3AsyncClient) delegate(),
+            delegateExtension,
+            apiCallBufferSize,
+            minPartSizeInBytes);
     }
 }

services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/PresignedUrlDownloadHelper.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.services.s3.internal.multipart;
+
+import java.util.concurrent.CompletableFuture;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.SplittingTransformerConfiguration;
+import software.amazon.awssdk.core.async.AsyncResponseTransformer;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
+import software.amazon.awssdk.services.s3.model.GetObjectResponse;
+import software.amazon.awssdk.services.s3.presignedurl.AsyncPresignedUrlExtension;
+import software.amazon.awssdk.services.s3.presignedurl.model.PresignedUrlDownloadRequest;
+import software.amazon.awssdk.utils.Logger;
+import software.amazon.awssdk.utils.Validate;
+
+@SdkInternalApi
+public class PresignedUrlDownloadHelper {
+    private static final Logger log = Logger.loggerFor(PresignedUrlDownloadHelper.class);
+
+    private final S3AsyncClient s3AsyncClient;
+    private final AsyncPresignedUrlExtension asyncPresignedUrlExtension;
+    private final long bufferSizeInBytes;
+    private final long configuredPartSizeInBytes;
+
+    public PresignedUrlDownloadHelper(S3AsyncClient s3AsyncClient,
+                                      AsyncPresignedUrlExtension asyncPresignedUrlExtension,
+                                      long bufferSizeInBytes,
+                                      long configuredPartSizeInBytes) {
+        this.s3AsyncClient = Validate.paramNotNull(s3AsyncClient, "s3AsyncClient");
+        this.asyncPresignedUrlExtension = Validate.paramNotNull(asyncPresignedUrlExtension, "asyncPresignedUrlExtension");
+        this.bufferSizeInBytes = Validate.isPositive(bufferSizeInBytes, "bufferSizeInBytes");
+        this.configuredPartSizeInBytes = Validate.isPositive(configuredPartSizeInBytes, "configuredPartSizeInBytes");
+    }
+
+    public <T> CompletableFuture<T> downloadObject(
+        PresignedUrlDownloadRequest presignedRequest,
+        AsyncResponseTransformer<GetObjectResponse, T> asyncResponseTransformer) {
+
+        Validate.paramNotNull(presignedRequest, "presignedRequest");
+        Validate.paramNotNull(asyncResponseTransformer, "asyncResponseTransformer");
+
+        if (presignedRequest.range() != null) {
+            log.debug(() -> "Using single part download because presigned URL request range is included in the request. range = "
+                            + presignedRequest.range());
+            return asyncPresignedUrlExtension.getObject(presignedRequest, asyncResponseTransformer);
+        }
+
+        SplittingTransformerConfiguration splittingConfig = SplittingTransformerConfiguration.builder()
+                                                                                             .bufferSizeInBytes(bufferSizeInBytes)
+                                                                                             .build();
+        AsyncResponseTransformer.SplitResult<GetObjectResponse, T> split =
+            asyncResponseTransformer.split(splittingConfig);
+        // TODO: PresignedUrlMultipartDownloaderSubscriber needs to be implemented in next PR
+        // PresignedUrlMultipartDownloaderSubscriber subscriber =
+        //     new PresignedUrlMultipartDownloaderSubscriber(
+        //         s3AsyncClient,
+        //         presignedRequest,
+        //         configuredPartSizeInBytes);
+        //
+        // split.publisher().subscribe(subscriber);
+        // return split.resultFuture();
+        throw new UnsupportedOperationException("Multipart presigned URL download not yet implemented - TODO in next PR");
+    }
+}

services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/MultipartS3AsyncClientTest.java

@@ -15,13 +15,17 @@
 
 package software.amazon.awssdk.services.s3.internal.multipart;
 
+import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+import java.net.MalformedURLException;
+import java.net.URL;
 import org.junit.jupiter.api.Test;
 import software.amazon.awssdk.core.ResponseBytes;
 import software.amazon.awssdk.core.SplittingTransformerConfiguration;
@@ -30,6 +34,8 @@
 import software.amazon.awssdk.services.s3.model.GetObjectRequest;
 import software.amazon.awssdk.services.s3.model.GetObjectResponse;
 import software.amazon.awssdk.services.s3.multipart.MultipartConfiguration;
+import software.amazon.awssdk.services.s3.presignedurl.AsyncPresignedUrlExtension;
+import software.amazon.awssdk.services.s3.presignedurl.model.PresignedUrlDownloadRequest;
 
 class MultipartS3AsyncClientTest {
 
@@ -64,4 +70,52 @@ void partManuallySpecified_shouldBypassMultipart() {
         verify(mockTransformer, never()).split(any(SplittingTransformerConfiguration.class));
         verify(mockDelegate, times(1)).getObject(any(GetObjectRequest.class), eq(mockTransformer));
     }
+
+    @Test
+    void presignedUrlExtension_rangeSpecified_shouldBypassMultipart() throws MalformedURLException {
+        S3AsyncClient mockDelegate = mock(S3AsyncClient.class);
+        AsyncPresignedUrlExtension mockDelegateExtension = mock(AsyncPresignedUrlExtension.class);
+        AsyncResponseTransformer<GetObjectResponse, String> mockTransformer = mock(AsyncResponseTransformer.class);
+        PresignedUrlDownloadRequest req = PresignedUrlDownloadRequest.builder()
+                                                                     .presignedUrl(new URL("https://s3.amazonaws.com/bucket/key?signature=abc"))
+                                                                     .range("bytes=0-1023")
+                                                                     .build();
+        when(mockDelegate.presignedUrlExtension()).thenReturn(mockDelegateExtension);
+        S3AsyncClient s3AsyncClient = MultipartS3AsyncClient.create(mockDelegate, MultipartConfiguration.builder().build(), true);
+        s3AsyncClient.presignedUrlExtension().getObject(req, mockTransformer);
+        verify(mockTransformer, never()).split(any(SplittingTransformerConfiguration.class));
+        verify(mockDelegateExtension, times(1)).getObject(eq(req), eq(mockTransformer));
+    }
+
+    // TODO: Enable this test once PresignedUrlMultipartDownloaderSubscriber is implemented
+    // // Currently fails because multipart presigned URL download throws UnsupportedOperationException
+    // @Test
+    // void presignedUrlExtension_noRange_shouldUseMultipart() throws MalformedURLException {
+    //     S3AsyncClient mockDelegate = mock(S3AsyncClient.class);
+    //     AsyncPresignedUrlExtension mockDelegateExtension = mock(AsyncPresignedUrlExtension.class);
+    //     AsyncResponseTransformer<GetObjectResponse, String> mockTransformer = mock(AsyncResponseTransformer.class);
+    //     AsyncResponseTransformer.SplitResult<GetObjectResponse, String> mockSplitResult = mock(AsyncResponseTransformer.SplitResult.class);
+    //     PresignedUrlDownloadRequest req = PresignedUrlDownloadRequest.builder()
+    //                                                                  .presignedUrl(new URL("https://s3.amazonaws.com/bucket/key?signature=abc"))
+    //                                                                  .build();
+    //     when(mockDelegate.presignedUrlExtension()).thenReturn(mockDelegateExtension);
+    //     when(mockTransformer.split(any(SplittingTransformerConfiguration.class))).thenReturn(mockSplitResult);
+    //     when(mockSplitResult.publisher()).thenReturn(mock(software.amazon.awssdk.core.async.SdkPublisher.class));
+    //     S3AsyncClient s3AsyncClient = MultipartS3AsyncClient.create(mockDelegate, MultipartConfiguration.builder().build(), true);
+    //     s3AsyncClient.presignedUrlExtension().getObject(req, mockTransformer);
+    //     verify(mockTransformer, times(1)).split(any(SplittingTransformerConfiguration.class));
+    //     verify(mockDelegateExtension, never()).getObject(any(PresignedUrlDownloadRequest.class), any(AsyncResponseTransformer.class));
+    // }
+
+    @Test
+    void presignedUrlExtension_shouldReturnMultipartExtension() {
+        S3AsyncClient mockDelegate = mock(S3AsyncClient.class);
+        AsyncPresignedUrlExtension mockDelegateExtension = mock(AsyncPresignedUrlExtension.class);
+        when(mockDelegate.presignedUrlExtension()).thenReturn(mockDelegateExtension);
+
+        S3AsyncClient s3AsyncClient = MultipartS3AsyncClient.create(mockDelegate, MultipartConfiguration.builder().build(), true);
+        AsyncPresignedUrlExtension extension = s3AsyncClient.presignedUrlExtension();
+
+        assertThat(extension).isInstanceOf(MultipartAsyncPresignedUrlExtension.class);
+    }
 }