Add support for request-level configuration override in AsyncClient multipart upload (#6900)

* Add support for request-level configuration override in S3 AsyncClient multipart upload

When using the S3 AsyncClient for multipart uploads and copies, request-level override configuration set on the original PutObjectRequest or CopyObjectRequest was not
being propagated to the sub-requests (CreateMultipartUpload, UploadPart, CompleteMultipartUpload, AbortMultipartUpload, ListParts, HeadObject, UploadPartCopy).

This change ensures overrideConfiguration is forwarded from the original request to all derived multipart sub-requests in SdkPojoConversionUtils and CopyObjectHelper.

* revision 2

---------

Co-authored-by: Dhananjay Badaya <badayd@amazon.com>
Co-authored-by: Alex Woods <alexwoo@amazon.com>

Author: 3 people

.changes/next-release/bugfix-AWSS3-52d62ca.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "Amazon S3",
+    "contributor": "dbadaya1",
+    "description": "Fixed request-level override configuration not being propagated to sub-requests during multipart uploads and copies in the S3 AsyncClient."
+}

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

@@ -162,8 +162,7 @@ private CompletableFuture<CompleteMultipartUploadResponse> completeMultipartUplo
             IntStream.range(0, completedParts.length())
                      .mapToObj(completedParts::get)
                      .toArray(CompletedPart[]::new);
-        CompleteMultipartUploadRequest completeMultipartUploadRequest =
-            CompleteMultipartUploadRequest.builder()
+        CompleteMultipartUploadRequest.Builder builder = CompleteMultipartUploadRequest.builder()
                                           .bucket(copyObjectRequest.destinationBucket())
                                           .key(copyObjectRequest.destinationKey())
                                           .uploadId(uploadId)
@@ -172,8 +171,9 @@ private CompletableFuture<CompleteMultipartUploadResponse> completeMultipartUplo
                                                                                    .build())
                                           .sseCustomerAlgorithm(copyObjectRequest.sseCustomerAlgorithm())
                                           .sseCustomerKey(copyObjectRequest.sseCustomerKey())
-                                          .sseCustomerKeyMD5(copyObjectRequest.sseCustomerKeyMD5())
-                                          .build();
+                                          .sseCustomerKeyMD5(copyObjectRequest.sseCustomerKeyMD5());
+        copyObjectRequest.overrideConfiguration().ifPresent(builder::overrideConfiguration);
+        CompleteMultipartUploadRequest completeMultipartUploadRequest = builder.build();
         return s3AsyncClient.completeMultipartUpload(completeMultipartUploadRequest);
     }
 

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

@@ -23,6 +23,7 @@
 import java.util.Objects;
 import java.util.Set;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.awscore.AwsRequest;
 import software.amazon.awssdk.core.SdkField;
 import software.amazon.awssdk.core.SdkPojo;
 import software.amazon.awssdk.core.exception.SdkClientException;
@@ -164,6 +165,7 @@ public static UploadPartRequest toUploadPartRequest(PutObjectRequest putObjectRe
         UploadPartRequest.Builder builder = UploadPartRequest.builder();
         validateRequestFields(putObjectRequest, builder.build(), PUT_OBJECT_TO_UPLOAD_PART_ALLOWED_FIELDS);
         setSdkFields(builder, putObjectRequest, PUT_OBJECT_REQUEST_TO_UPLOAD_PART_FIELDS_TO_IGNORE);
+        propagateOverrideConfig(builder, putObjectRequest);
         return builder.uploadId(uploadId).partNumber(partNumber).build();
     }
 
@@ -175,6 +177,7 @@ public static CompleteMultipartUploadRequest toCompleteMultipartUploadRequest(Pu
         setSdkFields(builder, putObjectRequest);
 
         builder.mpuObjectSize(contentLength);
+        propagateOverrideConfig(builder, putObjectRequest);
 
         if (S3ChecksumUtils.checksumValueSpecified(putObjectRequest)) {
             builder.checksumType(ChecksumType.FULL_OBJECT);
@@ -188,6 +191,7 @@ public static CreateMultipartUploadRequest toCreateMultipartUploadRequest(PutObj
         CreateMultipartUploadRequest.Builder builder = CreateMultipartUploadRequest.builder();
         validateRequestFields(putObjectRequest, builder.build(), PUT_OBJECT_TO_UPLOAD_PART_ALLOWED_FIELDS);
         setSdkFields(builder, putObjectRequest);
+        propagateOverrideConfig(builder, putObjectRequest);
 
         if (S3ChecksumUtils.checksumValueSpecified(putObjectRequest)) {
             builder.checksumType(ChecksumType.FULL_OBJECT);
@@ -200,7 +204,7 @@ public static HeadObjectRequest toHeadObjectRequest(CopyObjectRequest copyObject
 
         // We can't set SdkFields directly because the fields in CopyObjectRequest do not match 100% with the ones in
         // HeadObjectRequest
-        return HeadObjectRequest.builder()
+        HeadObjectRequest.Builder builder = HeadObjectRequest.builder()
                                 .bucket(copyObjectRequest.sourceBucket())
                                 .key(copyObjectRequest.sourceKey())
                                 .versionId(copyObjectRequest.sourceVersionId())
@@ -211,8 +215,9 @@ public static HeadObjectRequest toHeadObjectRequest(CopyObjectRequest copyObject
                                 .expectedBucketOwner(copyObjectRequest.expectedSourceBucketOwner())
                                 .sseCustomerAlgorithm(copyObjectRequest.copySourceSSECustomerAlgorithm())
                                 .sseCustomerKey(copyObjectRequest.copySourceSSECustomerKey())
-                                .sseCustomerKeyMD5(copyObjectRequest.copySourceSSECustomerKeyMD5())
-                                .build();
+                                .sseCustomerKeyMD5(copyObjectRequest.copySourceSSECustomerKeyMD5());
+        propagateOverrideConfig(builder, copyObjectRequest);
+        return builder.build();
     }
 
     public static CompletedPart toCompletedPart(CopyPartResult copyPartResult, int partNumber) {
@@ -238,6 +243,7 @@ public static ListPartsRequest toListPartsRequest(String uploadId, PutObjectRequ
         ListPartsRequest.Builder builder = ListPartsRequest.builder();
         validateRequestFields(putObjectRequest, builder.build(), PUT_OBJECT_TO_UPLOAD_PART_ALLOWED_FIELDS);
         setSdkFields(builder, putObjectRequest);
+        propagateOverrideConfig(builder, putObjectRequest);
         return builder.uploadId(uploadId).build();
     }
 
@@ -247,6 +253,7 @@ public static CreateMultipartUploadRequest toCreateMultipartUploadRequest(CopyOb
         setSdkFields(builder, copyObjectRequest);
         builder.bucket(copyObjectRequest.destinationBucket());
         builder.key(copyObjectRequest.destinationKey());
+        propagateOverrideConfig(builder, copyObjectRequest);
         return builder.build();
     }
 
@@ -275,13 +282,15 @@ public static AbortMultipartUploadRequest.Builder toAbortMultipartUploadRequest(
         setSdkFields(builder, copyObjectRequest);
         builder.bucket(copyObjectRequest.destinationBucket());
         builder.key(copyObjectRequest.destinationKey());
+        propagateOverrideConfig(builder, copyObjectRequest);
         return builder;
     }
 
     public static AbortMultipartUploadRequest.Builder toAbortMultipartUploadRequest(PutObjectRequest putObjectRequest) {
         AbortMultipartUploadRequest.Builder builder = AbortMultipartUploadRequest.builder();
         validateRequestFields(putObjectRequest, builder.build(), PUT_OBJECT_TO_UPLOAD_PART_ALLOWED_FIELDS);
         setSdkFields(builder, putObjectRequest);
+        propagateOverrideConfig(builder, putObjectRequest);
         return builder;
     }
 
@@ -296,6 +305,7 @@ public static UploadPartCopyRequest toUploadPartCopyRequest(CopyObjectRequest co
                       .uploadId(uploadId)
                       .bucket(copyObjectRequest.destinationBucket())
                       .key(copyObjectRequest.destinationKey())
+                      .overrideConfiguration(copyObjectRequest.overrideConfiguration().orElse(null))
                       .build();
     }
 
@@ -311,6 +321,10 @@ public static PutObjectResponse toPutObjectResponse(CompleteMultipartUploadRespo
         return builder.build();
     }
 
+    private static void propagateOverrideConfig(AwsRequest.Builder builder, AwsRequest source) {
+        source.overrideConfiguration().ifPresent(builder::overrideConfiguration);
+    }
+
     private static Map<String, Object> retrieveSdkFields(SdkPojo sourceObject, List<SdkField<?>> sdkFields) {
         return sdkFields.stream().collect(
             HashMap::new,

services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/crt/CopyObjectHelperTest.java

@@ -31,6 +31,7 @@
 import org.mockito.Mockito;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.internal.multipart.CopyObjectHelper;
@@ -281,6 +282,31 @@ void multiPartCopy_contentSizeExceeds10000Parts_shouldAdjustPartSize() {
     }
 
 
+    @Test
+    public void multiPartCopy_overrideConfigSetInOriginalRequest_includedInCompleteMultipart() {
+        AwsRequestOverrideConfiguration overrideConfig = AwsRequestOverrideConfiguration.builder()
+                                                                                        .putHeader("x-custom", "value")
+                                                                                        .build();
+        CopyObjectRequest copyRequest = copyObjectRequest().copy(r -> r.overrideConfiguration(overrideConfig));
+
+        stubSuccessfulHeadObjectCall(3 * PART_SIZE_BYTES);
+        stubSuccessfulCreateMulipartCall();
+        stubSuccessfulUploadPartCopyCalls();
+        stubSuccessfulCompleteMultipartCall();
+
+        copyHelper.copyObject(copyRequest).join();
+
+        ArgumentCaptor<CompleteMultipartUploadRequest> completeMultipartCaptor =
+            ArgumentCaptor.forClass(CompleteMultipartUploadRequest.class);
+
+        verify(s3AsyncClient).completeMultipartUpload(completeMultipartCaptor.capture());
+
+        CompleteMultipartUploadRequest completeRequest = completeMultipartCaptor.getValue();
+
+        assertThat(completeRequest.overrideConfiguration()).isPresent();
+        assertThat(completeRequest.overrideConfiguration().get()).isEqualTo(overrideConfig);
+    }
+
     @Test
     public void multiPartCopy_sseCHeadersSetInOriginalRequest_includedInCompleteMultipart() {
         String customerAlgorithm = "algorithm";