Skip to content

Commit 994a820

Browse files
zoewanggzoewangg
committed
Implement download directory in transfer manager
1 parent 6aaec97 commit 994a820

18 files changed

+1172
-35
lines changed

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

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.transfer.s3;
17+
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
import static software.amazon.awssdk.testutils.service.S3BucketUtils.temporaryBucketName;
20+
import static software.amazon.awssdk.utils.IoUtils.closeQuietly;
21+
22+
import java.io.IOException;
23+
import java.io.UncheckedIOException;
24+
import java.nio.charset.StandardCharsets;
25+
import java.nio.file.FileVisitResult;
26+
import java.nio.file.Files;
27+
import java.nio.file.Path;
28+
import java.nio.file.Paths;
29+
import java.nio.file.SimpleFileVisitor;
30+
import java.nio.file.attribute.BasicFileAttributes;
31+
import org.apache.commons.lang3.RandomStringUtils;
32+
import org.junit.After;
33+
import org.junit.AfterClass;
34+
import org.junit.Before;
35+
import org.junit.BeforeClass;
36+
import org.junit.Test;
37+
import software.amazon.awssdk.testutils.FileUtils;
38+
import software.amazon.awssdk.utils.Logger;
39+
40+
public class S3TransferManagerDownloadDirectoryIntegrationTest extends S3IntegrationTestBase {
41+
private static final Logger log = Logger.loggerFor(S3TransferManagerDownloadDirectoryIntegrationTest.class);
42+
private static final String TEST_BUCKET = temporaryBucketName(S3TransferManagerUploadIntegrationTest.class);
43+
private static final String TEST_BUCKET_SPECIAL_DELIMITER = temporaryBucketName("S3TransferManagerUploadIntegrationTest"
44+
+ "-delimiter");
45+
private static final String SPECIAL_DELIMITER = "-";
46+
47+
private static S3TransferManager tm;
48+
private static Path sourceDirectory;
49+
private Path destinationDirectory;
50+
51+
@BeforeClass
52+
public static void setUp() throws Exception {
53+
S3IntegrationTestBase.setUp();
54+
createBucket(TEST_BUCKET);
55+
createBucket(TEST_BUCKET_SPECIAL_DELIMITER);
56+
sourceDirectory = createLocalTestDirectory();
57+
58+
tm = S3TransferManager.builder()
59+
.s3ClientConfiguration(b -> b.credentialsProvider(CREDENTIALS_PROVIDER_CHAIN)
60+
.region(DEFAULT_REGION)
61+
.maxConcurrency(100))
62+
.build();
63+
64+
tm.uploadDirectory(u -> u.sourceDirectory(sourceDirectory).bucket(TEST_BUCKET)).completionFuture().join();
65+
66+
tm.uploadDirectory(u -> u.sourceDirectory(sourceDirectory)
67+
.delimiter(SPECIAL_DELIMITER)
68+
.bucket(TEST_BUCKET_SPECIAL_DELIMITER))
69+
.completionFuture().join();
70+
}
71+
72+
@Before
73+
public void setUpPerTest() throws IOException {
74+
destinationDirectory = Files.createTempDirectory("destination");
75+
}
76+
77+
@After
78+
public void cleanup() {
79+
FileUtils.cleanUpTestDirectory(destinationDirectory);
80+
}
81+
82+
@AfterClass
83+
public static void teardown() {
84+
try {
85+
FileUtils.cleanUpTestDirectory(sourceDirectory);
86+
} catch (Exception exception) {
87+
log.warn(() -> "Failed to clean up test directory " + sourceDirectory, exception);
88+
}
89+
90+
try {
91+
deleteBucketAndAllContents(TEST_BUCKET);
92+
} catch (Exception exception) {
93+
log.warn(() -> "Failed to delete s3 bucket " + TEST_BUCKET, exception);
94+
}
95+
96+
try {
97+
deleteBucketAndAllContents(TEST_BUCKET_SPECIAL_DELIMITER);
98+
} catch (Exception exception) {
99+
log.warn(() -> "Failed to delete s3 bucket " + TEST_BUCKET_SPECIAL_DELIMITER, exception);
100+
}
101+
102+
closeQuietly(tm, log.logger());
103+
S3IntegrationTestBase.cleanUp();
104+
}
105+
106+
@Test
107+
public void downloadDirectory() {
108+
DirectoryDownload downloadDirectory = tm.downloadDirectory(u -> u.destinationDirectory(destinationDirectory)
109+
.bucket(TEST_BUCKET));
110+
CompletedDirectoryDownload completedDirectoryDownload = downloadDirectory.completionFuture().join();
111+
assertThat(completedDirectoryDownload.failedTransfers()).isEmpty();
112+
assertTwoDirectoriesHaveSameStructure(sourceDirectory, destinationDirectory);
113+
}
114+
115+
@Test
116+
public void downloadDirectory_withPrefix() {
117+
String prefix = "notes";
118+
DirectoryDownload downloadDirectory = tm.downloadDirectory(u -> u.destinationDirectory(destinationDirectory)
119+
.prefix(prefix)
120+
.bucket(TEST_BUCKET));
121+
CompletedDirectoryDownload completedDirectoryDownload = downloadDirectory.completionFuture().join();
122+
assertThat(completedDirectoryDownload.failedTransfers()).isEmpty();
123+
124+
assertTwoDirectoriesHaveSameStructure(sourceDirectory.resolve(prefix), destinationDirectory.resolve(prefix));
125+
}
126+
127+
@Test
128+
public void downloadDirectory_withDelimiter() {
129+
String prefix = "notes";
130+
DirectoryDownload downloadDirectory = tm.downloadDirectory(u -> u.destinationDirectory(destinationDirectory)
131+
.delimiter(SPECIAL_DELIMITER)
132+
.prefix(prefix)
133+
.bucket(TEST_BUCKET_SPECIAL_DELIMITER));
134+
CompletedDirectoryDownload completedDirectoryDownload = downloadDirectory.completionFuture().join();
135+
assertThat(completedDirectoryDownload.failedTransfers()).isEmpty();
136+
assertTwoDirectoriesHaveSameStructure(sourceDirectory.resolve(prefix), destinationDirectory.resolve(prefix));
137+
}
138+
139+
private static void assertTwoDirectoriesHaveSameStructure(Path path, Path otherPath) {
140+
try {
141+
Files.walkFileTree(path, new SimpleFileVisitor<Path>() {
142+
@Override
143+
public FileVisitResult visitFile(Path file,
144+
BasicFileAttributes attrs)
145+
throws IOException {
146+
FileVisitResult result = super.visitFile(file, attrs);
147+
148+
Path relativePath = path.relativize(file);
149+
Path otherFile = otherPath.resolve(relativePath);
150+
log.info(() -> String.format("Comparing %s with %s", file, otherFile));
151+
assertThat(file).hasSameBinaryContentAs(otherFile);
152+
return result;
153+
}
154+
});
155+
} catch (IOException e) {
156+
throw new UncheckedIOException(String.format("Failed to compare %s with %s", path, otherPath), e);
157+
}
158+
}
159+
160+
/**
161+
* Create a test directory with the following structure
162+
* <pre>
163+
* {@code
164+
* - source
165+
* - README.md
166+
* - notes
167+
* - 2021
168+
* - 1.txt
169+
* - 2.txt
170+
* - 2022
171+
* - 1.txt
172+
* - important.txt
173+
* }
174+
* </pre>
175+
*/
176+
private static Path createLocalTestDirectory() throws IOException {
177+
Path directory = Files.createTempDirectory("source");
178+
179+
String directoryName = directory.toString();
180+
181+
Files.createDirectory(Paths.get(directoryName, "notes"));
182+
Files.createDirectory(Paths.get(directoryName, "notes", "2021"));
183+
Files.createDirectory(Paths.get(directoryName, "notes", "2022"));
184+
Files.write(Paths.get(directoryName, "README.md"), RandomStringUtils.random(100).getBytes(StandardCharsets.UTF_8));
185+
Files.write(Paths.get(directoryName, "notes", "2021", "1.txt"),
186+
RandomStringUtils.random(100).getBytes(StandardCharsets.UTF_8));
187+
Files.write(Paths.get(directoryName, "notes", "2021", "2.txt"),
188+
RandomStringUtils.random(100).getBytes(StandardCharsets.UTF_8));
189+
Files.write(Paths.get(directoryName, "notes", "2022", "1.txt"),
190+
RandomStringUtils.random(100).getBytes(StandardCharsets.UTF_8));
191+
Files.write(Paths.get(directoryName, "notes", "important.txt"),
192+
RandomStringUtils.random(100).getBytes(StandardCharsets.UTF_8));
193+
return directory;
194+
}
195+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -166,7 +166,7 @@ public interface Builder extends CopyableBuilder<Builder, DownloadDirectoryReque
166166
Builder destinationDirectory(Path destinationDirectory);
167167

168168
/**
169-
* The name of the bucket to download objects to.
169+
* The name of the bucket to download objects from.
170170
*
171171
* @param bucket the bucket name
172172
* @return This builder for method chaining.

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -281,8 +281,11 @@ default DirectoryUpload uploadDirectory(Consumer<UploadDirectoryRequest.Builder>
281281
* bucket will be downloaded.
282282
*
283283
* <p>
284+
* The SDK will create the destination directory if it does not already exist.
285+
*
286+
* <p>
284287
* The returned {@link CompletableFuture} only completes exceptionally if the request cannot be attempted as a whole (the
285-
* destination directory provided does not exist for example). The future completes successfully for partial successful
288+
* downloadDirectoryRequest is invalid for example). The future completes successfully for partial successful
286289
* requests, i.e., there might be failed downloads in a successfully completed response. As a result, you should check for
287290
* errors in the response via {@link CompletedDirectoryDownload#failedTransfers()} even when the future completes
288291
* successfully.

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

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.transfer.s3.internal;
17+
18+
import java.util.Objects;
19+
import java.util.concurrent.CompletableFuture;
20+
import software.amazon.awssdk.annotations.SdkInternalApi;
21+
import software.amazon.awssdk.transfer.s3.CompletedDirectoryDownload;
22+
import software.amazon.awssdk.transfer.s3.DirectoryDownload;
23+
import software.amazon.awssdk.utils.ToString;
24+
import software.amazon.awssdk.utils.Validate;
25+
26+
@SdkInternalApi
27+
public final class DefaultDirectoryDownload implements DirectoryDownload {
28+
29+
private final CompletableFuture<CompletedDirectoryDownload> completionFuture;
30+
31+
DefaultDirectoryDownload(CompletableFuture<CompletedDirectoryDownload> completionFuture) {
32+
this.completionFuture = Validate.paramNotNull(completionFuture, "completionFuture");
33+
}
34+
35+
@Override
36+
public CompletableFuture<CompletedDirectoryDownload> completionFuture() {
37+
return completionFuture;
38+
}
39+
40+
@Override
41+
public boolean equals(Object o) {
42+
if (this == o) {
43+
return true;
44+
}
45+
if (o == null || getClass() != o.getClass()) {
46+
return false;
47+
}
48+
49+
DefaultDirectoryDownload that = (DefaultDirectoryDownload) o;
50+
51+
return Objects.equals(completionFuture, that.completionFuture);
52+
}
53+
54+
@Override
55+
public int hashCode() {
56+
return completionFuture != null ? completionFuture.hashCode() : 0;
57+
}
58+
59+
@Override
60+
public String toString() {
61+
return ToString.builder("DefaultDirectoryDownload")
62+
.add("completionFuture", completionFuture)
63+
.build();
64+
}
65+
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,15 @@
2121
import software.amazon.awssdk.transfer.s3.CompletedDirectoryUpload;
2222
import software.amazon.awssdk.transfer.s3.DirectoryUpload;
2323
import software.amazon.awssdk.utils.ToString;
24+
import software.amazon.awssdk.utils.Validate;
2425

2526
@SdkInternalApi
2627
public final class DefaultDirectoryUpload implements DirectoryUpload {
2728

2829
private final CompletableFuture<CompletedDirectoryUpload> completionFuture;
2930

3031
DefaultDirectoryUpload(CompletableFuture<CompletedDirectoryUpload> completionFuture) {
31-
this.completionFuture = completionFuture;
32+
this.completionFuture = Validate.paramNotNull(completionFuture, "completionFuture");
3233
}
3334

3435
@Override

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@
4040
import software.amazon.awssdk.services.s3.S3Configuration;
4141
import software.amazon.awssdk.services.s3.model.GetObjectRequest;
4242
import software.amazon.awssdk.services.s3.model.GetObjectResponse;
43+
import software.amazon.awssdk.services.s3.model.ListObjectsV2Request;
44+
import software.amazon.awssdk.services.s3.model.ListObjectsV2Response;
4345
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
4446
import software.amazon.awssdk.services.s3.model.PutObjectResponse;
4547

@@ -97,6 +99,11 @@ public CompletableFuture<PutObjectResponse> putObject(PutObjectRequest putObject
9799
return s3AsyncClient.putObject(putObjectRequest, requestBody);
98100
}
99101

102+
@Override
103+
public CompletableFuture<ListObjectsV2Response> listObjectsV2(ListObjectsV2Request listObjectsV2Request) {
104+
return s3AsyncClient.listObjectsV2(listObjectsV2Request);
105+
}
106+
100107
@Override
101108
public String serviceName() {
102109
return SERVICE_NAME;

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

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,10 @@
3030
import software.amazon.awssdk.transfer.s3.CompletedFileDownload;
3131
import software.amazon.awssdk.transfer.s3.CompletedFileUpload;
3232
import software.amazon.awssdk.transfer.s3.CompletedUpload;
33+
import software.amazon.awssdk.transfer.s3.DirectoryDownload;
3334
import software.amazon.awssdk.transfer.s3.DirectoryUpload;
3435
import software.amazon.awssdk.transfer.s3.Download;
36+
import software.amazon.awssdk.transfer.s3.DownloadDirectoryRequest;
3537
import software.amazon.awssdk.transfer.s3.DownloadFileRequest;
3638
import software.amazon.awssdk.transfer.s3.DownloadRequest;
3739
import software.amazon.awssdk.transfer.s3.FileDownload;
@@ -52,20 +54,26 @@ public final class DefaultS3TransferManager implements S3TransferManager {
5254
private final S3CrtAsyncClient s3CrtAsyncClient;
5355
private final TransferManagerConfiguration transferConfiguration;
5456
private final UploadDirectoryHelper uploadDirectoryManager;
57+
private final DownloadDirectoryHelper downloadDirectoryHelper;
5558

5659
public DefaultS3TransferManager(DefaultBuilder tmBuilder) {
5760
transferConfiguration = resolveTransferManagerConfiguration(tmBuilder);
5861
s3CrtAsyncClient = initializeS3CrtClient(tmBuilder);
5962
uploadDirectoryManager = new UploadDirectoryHelper(transferConfiguration, this::uploadFile);
63+
downloadDirectoryHelper = new DownloadDirectoryHelper(transferConfiguration,
64+
s3CrtAsyncClient::listObjectsV2,
65+
this::downloadFile);
6066
}
6167

6268
@SdkTestInternalApi
6369
DefaultS3TransferManager(S3CrtAsyncClient s3CrtAsyncClient,
6470
UploadDirectoryHelper uploadDirectoryManager,
65-
TransferManagerConfiguration configuration) {
71+
TransferManagerConfiguration configuration,
72+
DownloadDirectoryHelper downloadDirectoryHelper) {
6673
this.s3CrtAsyncClient = s3CrtAsyncClient;
6774
this.transferConfiguration = configuration;
6875
this.uploadDirectoryManager = uploadDirectoryManager;
76+
this.downloadDirectoryHelper = downloadDirectoryHelper;
6977
}
7078

7179
private static TransferManagerConfiguration resolveTransferManagerConfiguration(DefaultBuilder tmBuilder) {
@@ -234,6 +242,19 @@ public FileDownload downloadFile(DownloadFileRequest downloadRequest) {
234242
return new DefaultFileDownload(downloadFuture, progressUpdater.progress());
235243
}
236244

245+
@Override
246+
public DirectoryDownload downloadDirectory(DownloadDirectoryRequest downloadDirectoryRequest) {
247+
Validate.paramNotNull(downloadDirectoryRequest, "downloadDirectoryRequest");
248+
249+
try {
250+
assertNotUnsupportedArn(downloadDirectoryRequest.bucket(), "downloadDirectoryRequest");
251+
252+
return downloadDirectoryHelper.downloadDirectory(downloadDirectoryRequest);
253+
} catch (Throwable throwable) {
254+
return new DefaultDirectoryDownload(CompletableFutureUtils.failedFuture(throwable));
255+
}
256+
}
257+
237258
@Override
238259
public void close() {
239260
s3CrtAsyncClient.close();

0 commit comments

Comments
 (0)