Added UploadDirectoryRequest.Builder#uploadRequestTransformer that allows modifying the UploadRequests generated during directory uploads.

Author: millems

services-custom/s3-transfer-manager/src/it/java/software/amazon/awssdk/transfer/s3/S3TransferManagerUploadDirectoryIntegrationTest.java

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.transfer.s3;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static software.amazon.awssdk.testutils.service.S3BucketUtils.temporaryBucketName;
 import static software.amazon.awssdk.utils.IoUtils.closeQuietly;
 
@@ -25,6 +26,9 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.stream.Collectors;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.junit.AfterClass;
@@ -124,6 +128,23 @@ public void uploadDirectory_withDelimiter_filesSentCorrectly() {
         });
     }
 
+    @Test
+    public void uploadDirectory_withRequestTransformer_usesRequestTransformer() throws Exception {
+        Path newSourceForEachUpload = Paths.get(directory.toString(), "bar.txt");
+
+        CompletedUploadDirectory result =
+            tm.uploadDirectory(r -> r.sourceDirectory(directory)
+                                     .bucket(TEST_BUCKET)
+                                     .uploadRequestTransformer(u -> u.toBuilder().source(newSourceForEachUpload).build()))
+              .completionFuture()
+              .get(10, TimeUnit.SECONDS);
+        assertThat(result.failedUploads()).isEmpty();
+
+        s3Client.listObjectsV2Paginator(b -> b.bucket(TEST_BUCKET)).contents().forEach(object -> {
+            verifyContent(object.key(), "bar.txt" + randomString);
+        });
+    }
+
     private static Path createLocalTestDirectory() throws IOException {
         Path directory = Files.createTempDirectory("test");
 

services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/UploadDirectoryRequest.java

@@ -22,6 +22,7 @@
 import java.util.function.Consumer;
 import software.amazon.awssdk.annotations.SdkPreviewApi;
 import software.amazon.awssdk.annotations.SdkPublicApi;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.utils.Validate;
 import software.amazon.awssdk.utils.builder.CopyableBuilder;
 import software.amazon.awssdk.utils.builder.ToCopyableBuilder;
@@ -41,13 +42,15 @@ public final class UploadDirectoryRequest implements TransferRequest, ToCopyable
     private final String prefix;
     private final UploadDirectoryOverrideConfiguration overrideConfiguration;
     private final String delimiter;
+    private final UploadRequestTransformer uploadRequestTransformer;
 
     public UploadDirectoryRequest(DefaultBuilder builder) {
         this.sourceDirectory = Validate.paramNotNull(builder.sourceDirectory, "sourceDirectory");
         this.bucket = Validate.paramNotNull(builder.bucket, "bucket");
         this.prefix = builder.prefix;
         this.overrideConfiguration = builder.configuration;
         this.delimiter = builder.delimiter;
+        this.uploadRequestTransformer = builder.uploadRequestTransformer;
     }
 
     /**
@@ -86,6 +89,14 @@ public Optional<String> delimiter() {
         return Optional.ofNullable(delimiter);
     }
 
+    /**
+     * @return the optional upload request transformer
+     * @see Builder#uploadRequestTransformer(UploadRequestTransformer)
+     */
+    public Optional<UploadRequestTransformer> uploadRequestTransformer() {
+        return Optional.ofNullable(uploadRequestTransformer);
+    }
+
     /**
      * @return the optional override configuration
      * @see Builder#overrideConfiguration(UploadDirectoryOverrideConfiguration)
@@ -130,6 +141,9 @@ public boolean equals(Object o) {
         if (!Objects.equals(delimiter, that.delimiter)) {
             return false;
         }
+        if (!Objects.equals(uploadRequestTransformer, that.uploadRequestTransformer)) {
+            return false;
+        }
         return Objects.equals(overrideConfiguration, that.overrideConfiguration);
     }
 
@@ -140,6 +154,7 @@ public int hashCode() {
         result = 31 * result + (prefix != null ? prefix.hashCode() : 0);
         result = 31 * result + (delimiter != null ? delimiter.hashCode() : 0);
         result = 31 * result + (overrideConfiguration != null ? overrideConfiguration.hashCode() : 0);
+        result = 31 * result + (uploadRequestTransformer != null ? uploadRequestTransformer.hashCode() : 0);
         return result;
     }
 
@@ -220,6 +235,51 @@ public interface Builder extends CopyableBuilder<Builder, UploadDirectoryRequest
          */
         Builder delimiter(String delimiter);
 
+        /**
+         * Specify a function used to transform the {@link UploadRequest}s generated by this {@link UploadDirectoryRequest}. The
+         * provided function is called once for each file that is uploaded, allowing you to modify the paths resolved by
+         * TransferManager on a per-file basis, modify the created {@link PutObjectRequest} before it is passed to S3, or
+         * configure a {@link TransferRequestOverrideConfiguration}.
+         *
+         * <p>The factory receives the {@link UploadRequest}s created by Transfer Manager for each file in the directory being
+         * uploaded, and returns a (potentially modified) {@link UploadRequest}.
+         *
+         * <p>
+         * <b>Usage Example:</b>
+         * <pre>
+         * {@code
+         * // Add a LoggingTransferListener to every transfer within the upload directory request
+         * TransferRequestOverrideConfiguration transferConfiguration =
+         *     TransferRequestOverrideConfiguration.builder()
+         *         .addListener(LoggingTransferListener.create())
+         *         .build();
+         *
+         * UploadDirectoryRequest request =
+         *     UploadDirectoryRequest.builder()
+         *         .sourceDirectory(Paths.get("."))
+         *         .bucket("bucket")
+         *         .prefix("prefix")
+         *         .uploadRequestTransformer(uploadRequest -> uploadRequest.toBuilder()
+         *                                                                 .overrideConfiguration(transferConfiguration)
+         *                                                                 .build())
+         *         .build()
+         *
+         * UploadDirectoryTransfer uploadDirectory = transferManager.uploadDirectory(request);
+         *
+         * // Wait for the transfer to complete
+         * CompletedUploadDirectory completedUploadDirectory = uploadDirectory.completionFuture().join();
+         *
+         * // Print out the failed uploads
+         * completedUploadDirectory.failedUploads().forEach(System.out::println);
+         * }
+         * </pre>
+         *
+         * @param uploadRequestTransformer A transformer to use for modifying the file-level upload requests before execution
+         * @return This builder for method chaining
+         * @see UploadRequestTransformer
+         */
+        Builder uploadRequestTransformer(UploadRequestTransformer uploadRequestTransformer);
+
         /**
          * Add an optional request override configuration.
          *
@@ -249,13 +309,15 @@ default Builder overrideConfiguration(Consumer<UploadDirectoryOverrideConfigurat
         UploadDirectoryRequest build();
     }
 
+
     private static final class DefaultBuilder implements Builder {
 
         private Path sourceDirectory;
         private String bucket;
         private String prefix;
         private UploadDirectoryOverrideConfiguration configuration;
         private String delimiter;
+        private UploadRequestTransformer uploadRequestTransformer;
 
         private DefaultBuilder() {
         }
@@ -265,6 +327,8 @@ private DefaultBuilder(UploadDirectoryRequest request) {
             this.bucket = request.bucket;
             this.prefix = request.prefix;
             this.configuration = request.overrideConfiguration;
+            this.delimiter = request.delimiter;
+            this.uploadRequestTransformer = request.uploadRequestTransformer;
         }
 
         @Override
@@ -315,6 +379,12 @@ public Builder delimiter(String delimiter) {
             return this;
         }
 
+        @Override
+        public Builder uploadRequestTransformer(UploadRequestTransformer uploadRequestTransformer) {
+            this.uploadRequestTransformer = uploadRequestTransformer;
+            return this;
+        }
+
         public void setDelimiter(String delimiter) {
             delimiter(delimiter);
         }

services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/UploadRequest.java

@@ -87,7 +87,7 @@ public static Class<? extends Builder> serializableBuilderClass() {
 
     @Override
     public Builder toBuilder() {
-        return new BuilderImpl();
+        return new BuilderImpl(this);
     }
 
     @Override
@@ -219,6 +219,15 @@ private static class BuilderImpl implements Builder {
         private Path source;
         private TransferRequestOverrideConfiguration configuration;
 
+        private BuilderImpl() {
+        }
+
+        private BuilderImpl(UploadRequest uploadRequest) {
+            this.source = uploadRequest.source;
+            this.putObjectRequest = uploadRequest.putObjectRequest;
+            this.configuration = uploadRequest.overrideConfiguration;
+        }
+
         @Override
         public Builder source(Path source) {
             this.source = source;

services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/UploadRequestTransformer.java

@@ -0,0 +1,45 @@
+/*
+ * 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.transfer.s3;
+
+import software.amazon.awssdk.annotations.SdkPreviewApi;
+import software.amazon.awssdk.annotations.SdkPublicApi;
+import software.amazon.awssdk.annotations.ThreadSafe;
+
+/**
+ * A transformer for modifying {@link UploadRequest}s. This is commonly used as a functional interface in an upload request via
+ * {@link UploadDirectoryRequest.Builder#uploadRequestTransformer(UploadRequestTransformer)}.
+ */
+@FunctionalInterface
+@ThreadSafe
+@SdkPublicApi
+@SdkPreviewApi
+public interface UploadRequestTransformer {
+    /**
+     * Transform an {@link UploadRequest}.
+     *
+     * @param originalRequest The request to transform
+     * @return The transformed request.
+     */
+    UploadRequest transform(UploadRequest originalRequest);
+
+    /**
+     * Return an implementation of {@link UploadRequestTransformer} that returns the input request without modification.
+     */
+    static UploadRequestTransformer identity() {
+        return r -> r;
+    }
+}

services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/UploadDirectoryHelper.java

@@ -43,6 +43,7 @@
 import software.amazon.awssdk.transfer.s3.UploadDirectoryRequest;
 import software.amazon.awssdk.transfer.s3.UploadDirectoryTransfer;
 import software.amazon.awssdk.transfer.s3.UploadRequest;
+import software.amazon.awssdk.transfer.s3.UploadRequestTransformer;
 import software.amazon.awssdk.utils.CompletableFutureUtils;
 import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.StringUtils;
@@ -230,9 +231,13 @@ private UploadRequest constructUploadRequest(UploadDirectoryRequest uploadDirect
                                                             .bucket(uploadDirectoryRequest.bucket())
                                                             .key(key)
                                                             .build();
-        return UploadRequest.builder()
-                            .source(path)
-                            .putObjectRequest(putObjectRequest)
-                            .build();
+
+        UploadRequestTransformer uploadRequestTransformer = uploadDirectoryRequest.uploadRequestTransformer()
+                                                                                  .orElse(UploadRequestTransformer.identity());
+
+        return uploadRequestTransformer.transform(UploadRequest.builder()
+                                                               .source(path)
+                                                               .putObjectRequest(putObjectRequest)
+                                                               .build());
     }
 }

services-custom/s3-transfer-manager/src/test/java/software/amazon/awssdk/transfer/s3/internal/UploadDirectoryHelperTest.java

@@ -26,6 +26,8 @@
 import java.nio.file.FileSystem;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.List;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -36,10 +38,13 @@
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectResponse;
 import software.amazon.awssdk.transfer.s3.CompletedUpload;
 import software.amazon.awssdk.transfer.s3.CompletedUploadDirectory;
+import software.amazon.awssdk.transfer.s3.TransferRequestOverrideConfiguration;
 import software.amazon.awssdk.transfer.s3.Upload;
 import software.amazon.awssdk.transfer.s3.UploadDirectoryRequest;
 import software.amazon.awssdk.transfer.s3.UploadDirectoryTransfer;
@@ -156,6 +161,56 @@ public void uploadDirectory_partialSuccess_shouldProvideFailedUploads() throws E
         assertThat(completedUploadDirectory.failedUploads().iterator().next().request().source().toString()).isEqualTo("test/2");
     }
 
+    @Test
+    public void uploadDirectory_withRequestTransformer_usesRequestTransformer() throws Exception {
+        PutObjectResponse putObjectResponse = PutObjectResponse.builder().eTag("1234").build();
+        CompletedUpload completedUpload = CompletedUpload.builder().response(putObjectResponse).build();
+        CompletableFuture<CompletedUpload> successfulFuture = new CompletableFuture<>();
+
+        Upload upload = newUpload(successfulFuture);
+        successfulFuture.complete(completedUpload);
+
+        PutObjectResponse putObjectResponse2 = PutObjectResponse.builder().eTag("5678").build();
+        CompletedUpload completedUpload2 = CompletedUpload.builder().response(putObjectResponse2).build();
+        CompletableFuture<CompletedUpload> failedFuture = new CompletableFuture<>();
+        Upload upload2 = newUpload(failedFuture);
+        failedFuture.complete(completedUpload2);
+
+        ArgumentCaptor<UploadRequest> uploadRequestCaptor = ArgumentCaptor.forClass(UploadRequest.class);
+
+        when(singleUploadFunction.apply(uploadRequestCaptor.capture())).thenReturn(upload, upload2);
+
+        Path newSource = Paths.get("/new/path");
+        PutObjectRequest newPutObjectRequest = PutObjectRequest.builder().build();
+        TransferRequestOverrideConfiguration newOverrideConfig = TransferRequestOverrideConfiguration.builder()
+                                                                                                     .build();
+
+        uploadDirectoryHelper.uploadDirectory(UploadDirectoryRequest.builder()
+                                                                    .sourceDirectory(directory)
+                                                                    .bucket("bucket")
+                                                                    .uploadRequestTransformer(r -> r.toBuilder()
+                                                                                                    .source(newSource)
+                                                                                                    .putObjectRequest(newPutObjectRequest)
+                                                                                                    .overrideConfiguration(newOverrideConfig)
+                                                                                                    .build())
+                                                                    .build())
+                             .completionFuture()
+                             .get(5, TimeUnit.SECONDS);
+
+        List<UploadRequest> uploadRequests = uploadRequestCaptor.getAllValues();
+        assertThat(uploadRequests).hasSize(2);
+        assertThat(uploadRequests).element(0).satisfies(r -> {
+            assertThat(r.source()).isEqualTo(newSource);
+            assertThat(r.putObjectRequest()).isEqualTo(newPutObjectRequest);
+            assertThat(r.overrideConfiguration()).hasValue(newOverrideConfig);
+        });
+        assertThat(uploadRequests).element(1).satisfies(r -> {
+            assertThat(r.source()).isEqualTo(newSource);
+            assertThat(r.putObjectRequest()).isEqualTo(newPutObjectRequest);
+            assertThat(r.overrideConfiguration()).hasValue(newOverrideConfig);
+        });
+    }
+
     private Upload newUpload(CompletableFuture<CompletedUpload> future) {
         return new DefaultUpload(future,
                                  new DefaultTransferProgress(DefaultTransferProgressSnapshot.builder().build())