Added UploadDirectoryOverrideConfiguration.Builder#uploadFileRequestTransformer that allows modifying the UploadFileRequests generated during directory uploads. (#2799)

millems · web-flow · commit 821a34232a16 · 2021-11-10T10:30:10.000-08:00
diff --git a/services-custom/s3-transfer-manager/src/it/java/software/amazon/awssdk/transfer/s3/S3TransferManagerUploadDirectoryIntegrationTest.java b/services-custom/s3-transfer-manager/src/it/java/software/amazon/awssdk/transfer/s3/S3TransferManagerUploadDirectoryIntegrationTest.java
@@ -25,6 +25,7 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.junit.AfterClass;
@@ -124,6 +125,26 @@ public void uploadDirectory_withDelimiter_filesSentCorrectly() {
         });
     }
 
+    @Test
+    public void uploadDirectory_withRequestTransformer_usesRequestTransformer() throws Exception {
+        String prefix = "requestTransformerTest";
+        Path newSourceForEachUpload = Paths.get(directory.toString(), "bar.txt");
+
+        CompletedDirectoryUpload result =
+            tm.uploadDirectory(r -> r.sourceDirectory(directory)
+                                     .bucket(TEST_BUCKET)
+                                     .prefix(prefix)
+                                     .overrideConfiguration(c -> c.uploadFileRequestTransformer(f -> f.source(newSourceForEachUpload))
+                                                                  .recursive(true)))
+              .completionFuture()
+              .get(10, TimeUnit.SECONDS);
+        assertThat(result.failedTransfers()).isEmpty();
+
+        s3Client.listObjectsV2Paginator(b -> b.bucket(TEST_BUCKET).prefix(prefix)).contents().forEach(object -> {
+            verifyContent(object.key(), "bar.txt" + randomString);
+        });
+    }
+
     private static Path createLocalTestDirectory() throws IOException {
         Path directory = Files.createTempDirectory("test");
 
diff --git a/services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/UploadDirectoryOverrideConfiguration.java b/services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/UploadDirectoryOverrideConfiguration.java
@@ -17,8 +17,10 @@
 
 import java.util.Objects;
 import java.util.Optional;
+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.ToString;
 import software.amazon.awssdk.utils.Validate;
 import software.amazon.awssdk.utils.builder.CopyableBuilder;
@@ -39,11 +41,13 @@ public final class UploadDirectoryOverrideConfiguration implements ToCopyableBui
     private final Boolean followSymbolicLinks;
     private final Integer maxDepth;
     private final Boolean recursive;
+    private final Consumer<UploadFileRequest.Builder> uploadFileRequestTransformer;
 
     public UploadDirectoryOverrideConfiguration(DefaultBuilder builder) {
         this.followSymbolicLinks = builder.followSymbolicLinks;
         this.maxDepth = Validate.isPositiveOrNull(builder.maxDepth, "maxDepth");
         this.recursive = builder.recursive;
+        this.uploadFileRequestTransformer = builder.uploadFileRequestTransformer;
     }
 
     /**
@@ -70,6 +74,14 @@ public Optional<Boolean> recursive() {
         return Optional.ofNullable(recursive);
     }
 
+    /**
+     * @return the optional upload request transformer
+     * @see UploadDirectoryOverrideConfiguration.Builder#uploadFileRequestTransformer(Consumer)
+     */
+    public Optional<Consumer<UploadFileRequest.Builder>> uploadFileRequestTransformer() {
+        return Optional.ofNullable(uploadFileRequestTransformer);
+    }
+
     @Override
     public Builder toBuilder() {
         return new DefaultBuilder(this);
@@ -92,6 +104,9 @@ public boolean equals(Object o) {
         if (!Objects.equals(maxDepth, that.maxDepth)) {
             return false;
         }
+        if (!Objects.equals(uploadFileRequestTransformer, that.uploadFileRequestTransformer)) {
+            return false;
+        }
         return Objects.equals(recursive, that.recursive);
     }
 
@@ -100,6 +115,7 @@ public int hashCode() {
         int result = followSymbolicLinks != null ? followSymbolicLinks.hashCode() : 0;
         result = 31 * result + (maxDepth != null ? maxDepth.hashCode() : 0);
         result = 31 * result + (recursive != null ? recursive.hashCode() : 0);
+        result = 31 * result + (uploadFileRequestTransformer != null ? uploadFileRequestTransformer.hashCode() : 0);
         return result;
     }
 
@@ -109,6 +125,7 @@ public String toString() {
                        .add("followSymbolicLinks", followSymbolicLinks)
                        .add("maxDepth", maxDepth)
                        .add("recursive", recursive)
+                       .add("uploadFileRequestTransformer", uploadFileRequestTransformer)
                        .build();
     }
 
@@ -156,6 +173,53 @@ public interface Builder extends CopyableBuilder<Builder, UploadDirectoryOverrid
          */
         Builder maxDepth(Integer maxDepth);
 
+        /**
+         * Specify a function used to transform the {@link UploadFileRequest}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 UploadFileRequest}s created by Transfer Manager for each file in the directory
+         * being uploaded, and returns a (potentially modified) {@code UploadFileRequest}.
+         *
+         * <p>
+         * <b>Usage Example:</b>
+         * <pre>
+         * {@code
+         * // Add a LoggingTransferListener to every transfer within the upload directory request
+         * TransferRequestOverrideConfiguration fileUploadConfiguration =
+         *     TransferRequestOverrideConfiguration.builder()
+         *         .addListener(LoggingTransferListener.create())
+         *         .build();
+         *
+         * UploadDirectoryOverrideConfiguration directoryUploadConfiguration =
+         *     UploadDirectoryOverrideConfiguration.builder()
+         *         .uploadFileRequestTransformer(request -> request.overrideConfiguration(fileUploadConfiguration))
+         *         .build();
+         *
+         * UploadDirectoryRequest request =
+         *     UploadDirectoryRequest.builder()
+         *         .sourceDirectory(Paths.get("."))
+         *         .bucket("bucket")
+         *         .prefix("prefix")
+         *         .overrideConfiguration(directoryUploadConfiguration)
+         *         .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 uploadFileRequestTransformer A transformer to use for modifying the file-level upload requests before execution
+         * @return This builder for method chaining
+         */
+        Builder uploadFileRequestTransformer(Consumer<UploadFileRequest.Builder> uploadFileRequestTransformer);
+
         @Override
         UploadDirectoryOverrideConfiguration build();
     }
@@ -164,11 +228,13 @@ private static final class DefaultBuilder implements Builder {
         private Boolean followSymbolicLinks;
         private Integer maxDepth;
         private Boolean recursive;
+        private Consumer<UploadFileRequest.Builder> uploadFileRequestTransformer;
 
         private DefaultBuilder(UploadDirectoryOverrideConfiguration configuration) {
             this.followSymbolicLinks = configuration.followSymbolicLinks;
             this.maxDepth = configuration.maxDepth;
             this.recursive = configuration.recursive;
+            this.uploadFileRequestTransformer = configuration.uploadFileRequestTransformer;
         }
 
         private DefaultBuilder() {
@@ -216,6 +282,20 @@ public Integer getMaxDepth() {
             return maxDepth;
         }
 
+        @Override
+        public Builder uploadFileRequestTransformer(Consumer<UploadFileRequest.Builder> uploadFileRequestTransformer) {
+            this.uploadFileRequestTransformer = uploadFileRequestTransformer;
+            return this;
+        }
+
+        public Consumer<UploadFileRequest.Builder> getUploadFileRequestTransformer() {
+            return uploadFileRequestTransformer;
+        }
+
+        public void setUploadFileRequestTransformer(Consumer<UploadFileRequest.Builder> uploadFileRequestTransformer) {
+            this.uploadFileRequestTransformer = uploadFileRequestTransformer;
+        }
+
         @Override
         public UploadDirectoryOverrideConfiguration build() {
             return new UploadDirectoryOverrideConfiguration(this);
diff --git a/services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/UploadDirectoryRequest.java b/services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/UploadDirectoryRequest.java
@@ -156,7 +156,6 @@ public String toString() {
     }
 
     public interface Builder extends CopyableBuilder<Builder, UploadDirectoryRequest> {
-
         /**
          * Specify the source directory to upload. The source directory must exist.
          * Fle wildcards are not supported and treated literally. Hidden files/directories are visited.
@@ -261,6 +260,7 @@ default Builder overrideConfiguration(Consumer<UploadDirectoryOverrideConfigurat
         UploadDirectoryRequest build();
     }
 
+
     private static final class DefaultBuilder implements Builder {
 
         private Path sourceDirectory;
@@ -277,6 +277,7 @@ private DefaultBuilder(UploadDirectoryRequest request) {
             this.bucket = request.bucket;
             this.prefix = request.prefix;
             this.configuration = request.overrideConfiguration;
+            this.delimiter = request.delimiter;
         }
 
         @Override
diff --git a/services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/UploadDirectoryHelper.java b/services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/UploadDirectoryHelper.java
@@ -41,6 +41,7 @@
 import software.amazon.awssdk.transfer.s3.FailedFileUpload;
 import software.amazon.awssdk.transfer.s3.FileUpload;
 import software.amazon.awssdk.transfer.s3.S3TransferManager;
+import software.amazon.awssdk.transfer.s3.UploadDirectoryOverrideConfiguration;
 import software.amazon.awssdk.transfer.s3.UploadDirectoryRequest;
 import software.amazon.awssdk.transfer.s3.UploadFileRequest;
 import software.amazon.awssdk.utils.CompletableFutureUtils;
@@ -230,9 +231,16 @@ private UploadFileRequest constructUploadRequest(UploadDirectoryRequest uploadDi
                                                             .bucket(uploadDirectoryRequest.bucket())
                                                             .key(key)
                                                             .build();
-        return UploadFileRequest.builder()
-                                .source(path)
-                                .putObjectRequest(putObjectRequest)
-                                .build();
+
+        UploadFileRequest.Builder requestBuilder = UploadFileRequest.builder()
+                                                                    .source(path)
+                                                                    .putObjectRequest(putObjectRequest);
+
+        uploadDirectoryRequest.overrideConfiguration()
+                              .flatMap(UploadDirectoryOverrideConfiguration::uploadFileRequestTransformer)
+                              .ifPresent(c -> c.accept(requestBuilder));
+
+        return requestBuilder.build();
     }
+
 }
diff --git a/services-custom/s3-transfer-manager/src/test/java/software/amazon/awssdk/transfer/s3/internal/UploadDirectoryHelperTest.java b/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,12 +38,16 @@
 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.CompletedDirectoryUpload;
 import software.amazon.awssdk.transfer.s3.CompletedFileUpload;
 import software.amazon.awssdk.transfer.s3.DirectoryUpload;
 import software.amazon.awssdk.transfer.s3.FileUpload;
+import software.amazon.awssdk.transfer.s3.TransferRequestOverrideConfiguration;
+import software.amazon.awssdk.transfer.s3.UploadDirectoryOverrideConfiguration;
 import software.amazon.awssdk.transfer.s3.UploadDirectoryRequest;
 import software.amazon.awssdk.transfer.s3.UploadFileRequest;
 import software.amazon.awssdk.transfer.s3.internal.progress.DefaultTransferProgress;
@@ -156,6 +162,59 @@ public void uploadDirectory_partialSuccess_shouldProvideFailedUploads() throws E
         assertThat(completedDirectoryUpload.failedTransfers().iterator().next().request().source().toString()).isEqualTo("test/2");
     }
 
+    @Test
+    public void uploadDirectory_withRequestTransformer_usesRequestTransformer() throws Exception {
+        PutObjectResponse putObjectResponse = PutObjectResponse.builder().eTag("1234").build();
+        CompletedFileUpload completedFileUpload = CompletedFileUpload.builder().response(putObjectResponse).build();
+        CompletableFuture<CompletedFileUpload> successfulFuture = new CompletableFuture<>();
+
+        FileUpload upload = newUpload(successfulFuture);
+        successfulFuture.complete(completedFileUpload);
+
+        PutObjectResponse putObjectResponse2 = PutObjectResponse.builder().eTag("5678").build();
+        CompletedFileUpload completedFileUpload2 = CompletedFileUpload.builder().response(putObjectResponse2).build();
+        CompletableFuture<CompletedFileUpload> failedFuture = new CompletableFuture<>();
+        FileUpload upload2 = newUpload(failedFuture);
+        failedFuture.complete(completedFileUpload2);
+
+        ArgumentCaptor<UploadFileRequest> uploadRequestCaptor = ArgumentCaptor.forClass(UploadFileRequest.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();
+
+        UploadDirectoryOverrideConfiguration uploadConfig =
+            UploadDirectoryOverrideConfiguration.builder()
+                                                .uploadFileRequestTransformer(r -> r.source(newSource)
+                                                                                    .putObjectRequest(newPutObjectRequest)
+                                                                                    .overrideConfiguration(newOverrideConfig))
+                                                .build();
+
+        uploadDirectoryHelper.uploadDirectory(UploadDirectoryRequest.builder()
+                                                                    .sourceDirectory(directory)
+                                                                    .bucket("bucket")
+                                                                    .overrideConfiguration(uploadConfig)
+                                                                    .build())
+                             .completionFuture()
+                             .get(5, TimeUnit.SECONDS);
+
+        List<UploadFileRequest> 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 FileUpload newUpload(CompletableFuture<CompletedFileUpload> future) {
         return new DefaultFileUpload(future,
                                  new DefaultTransferProgress(DefaultTransferProgressSnapshot.builder().build())

Original file line number	Diff line number	Diff line change
`@@ -156,7 +156,6 @@ public String toString() {`
`156`	`156`	`}`
`157`	`157`
`158`	`158`	`public interface Builder extends CopyableBuilder<Builder, UploadDirectoryRequest> {`
`159`		`-`
`160`	`159`	`/**`
`161`	`160`	`* Specify the source directory to upload. The source directory must exist.`
`162`	`161`	`* Fle wildcards are not supported and treated literally. Hidden files/directories are visited.`
`@@ -261,6 +260,7 @@ default Builder overrideConfiguration(Consumer<UploadDirectoryOverrideConfigurat`
`261`	`260`	`UploadDirectoryRequest build();`
`262`	`261`	`}`
`263`	`262`
	`263`	`+`
`264`	`264`	`private static final class DefaultBuilder implements Builder {`
`265`	`265`
`266`	`266`	`private Path sourceDirectory;`
`@@ -277,6 +277,7 @@ private DefaultBuilder(UploadDirectoryRequest request) {`
`277`	`277`	`this.bucket = request.bucket;`
`278`	`278`	`this.prefix = request.prefix;`
`279`	`279`	`this.configuration = request.overrideConfiguration;`
	`280`	`+ this.delimiter = request.delimiter;`
`280`	`281`	`}`
`281`	`282`
`282`	`283`	`@Override`