Skip to content

Commit 24a74d8

Browse files
Fred1155Fred1155
authored
S3 transfer manager client constructor and method recipe (#6151)
* recipe added * delete upload recipe * comments addressed
1 parent cef97ad commit 24a74d8

File tree

6 files changed

+158
-5
lines changed

6 files changed

+158
-5
lines changed

‎test/v2-migration-tests/src/test/resources/software/amazon/awssdk/v2migrationtests/maven-tm/after/src/main/java/foo/bar/TransferManagerS3.java‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,25 @@
1515

1616
package foo.bar;
1717

18+
import software.amazon.awssdk.auth.credentials.AwsCredentials;
19+
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
20+
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
1821
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
1922
import software.amazon.awssdk.core.async.AsyncRequestBody;
23+
import software.amazon.awssdk.services.s3.S3AsyncClient;
2024
import software.amazon.awssdk.services.s3.model.CopyObjectRequest;
2125
import software.amazon.awssdk.services.s3.model.GetObjectRequest;
2226
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
2327
import software.amazon.awssdk.transfer.s3.S3TransferManager;
2428
import software.amazon.awssdk.transfer.s3.model.Copy;
2529
import software.amazon.awssdk.transfer.s3.model.CopyRequest;
30+
import software.amazon.awssdk.transfer.s3.model.DirectoryDownload;
31+
import software.amazon.awssdk.transfer.s3.model.DownloadDirectoryRequest;
2632
import software.amazon.awssdk.transfer.s3.model.DownloadFileRequest;
2733
import software.amazon.awssdk.transfer.s3.model.FileDownload;
34+
import software.amazon.awssdk.transfer.s3.model.FileUpload;
35+
import software.amazon.awssdk.transfer.s3.model.ResumableFileDownload;
36+
import software.amazon.awssdk.transfer.s3.model.ResumableFileUpload;
2837
import software.amazon.awssdk.transfer.s3.model.UploadFileRequest;
2938
import software.amazon.awssdk.transfer.s3.model.UploadRequest;
3039

@@ -35,11 +44,13 @@ public class TransferManagerS3 {
3544

3645
File file = new File("path/to/file.txt");
3746

38-
void tmConstructor() {
47+
void tmConstructor(AwsCredentials credentials, AwsCredentialsProvider credentialsProvider) {
3948
S3TransferManager tm = S3TransferManager.builder()
4049
.build();
4150
S3TransferManager tmBuilderDefault = S3TransferManager.create();
4251
S3TransferManager tmBuilderWithS3 = S3TransferManager.builder().build();
52+
S3TransferManager tmConstructorWithCred = S3TransferManager.builder().s3Client(S3AsyncClient.builder().credentialsProvider(StaticCredentialsProvider.create(credentials)).build()).build();
53+
S3TransferManager tmConstructorWithCredProvider = S3TransferManager.builder().s3Client(S3AsyncClient.builder().credentialsProvider(credentialsProvider).build()).build();
4354
}
4455

4556
void download(S3TransferManager tm, String bucket, String key) {
@@ -76,4 +87,14 @@ void copy(S3TransferManager tm, String sourceBucket, String sourceKey, String de
7687
.build();
7788
Copy copy2 = tm.copy(CopyRequest.builder().copyObjectRequest(copyRequest).build());
7889
}
90+
91+
void downloadDirectory(S3TransferManager tm, File destination) {
92+
DirectoryDownload fileDownload = tm.downloadDirectory(DownloadDirectoryRequest.builder().bucket("bucket").listObjectsV2RequestTransformer(builder -> builder.prefix("key")).destination(destination.toPath()).build());
93+
tm.close();
94+
}
95+
96+
void resume(S3TransferManager tm, ResumableFileDownload persistableDownload, ResumableFileUpload persistableUpload) {
97+
FileDownload download = tm.resumeDownloadFile(persistableDownload);
98+
FileUpload upload = tm.resumeUploadFile(persistableUpload);
99+
}
79100
}

‎test/v2-migration-tests/src/test/resources/software/amazon/awssdk/v2migrationtests/maven-tm/before/src/main/java/foo/bar/TransferManagerS3.java‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,23 +15,31 @@
1515

1616
package foo.bar;
1717

18+
import com.amazonaws.auth.AWSCredentials;
19+
import com.amazonaws.auth.AWSCredentialsProvider;
1820
import com.amazonaws.services.s3.model.CopyObjectRequest;
1921
import com.amazonaws.services.s3.model.GetObjectRequest;
2022
import com.amazonaws.services.s3.model.PutObjectRequest;
2123
import com.amazonaws.services.s3.transfer.Copy;
2224
import com.amazonaws.services.s3.transfer.Download;
25+
import com.amazonaws.services.s3.transfer.MultipleFileDownload;
26+
import com.amazonaws.services.s3.transfer.PersistableDownload;
27+
import com.amazonaws.services.s3.transfer.PersistableUpload;
2328
import com.amazonaws.services.s3.transfer.TransferManager;
2429
import com.amazonaws.services.s3.transfer.TransferManagerBuilder;
30+
import com.amazonaws.services.s3.transfer.Upload;
2531
import java.io.File;
2632

2733
public class TransferManagerS3 {
2834

2935
File file = new File("path/to/file.txt");
3036

31-
void tmConstructor() {
37+
void tmConstructor(AWSCredentials credentials, AWSCredentialsProvider credentialsProvider) {
3238
TransferManager tm = new TransferManager();
3339
TransferManager tmBuilderDefault = TransferManagerBuilder.defaultTransferManager();
3440
TransferManager tmBuilderWithS3 = TransferManagerBuilder.standard().build();
41+
TransferManager tmConstructorWithCred = new TransferManager(credentials);
42+
TransferManager tmConstructorWithCredProvider = new TransferManager(credentialsProvider);
3543
}
3644

3745
void download(TransferManager tm, String bucket, String key) {
@@ -64,4 +72,14 @@ void copy(TransferManager tm, String sourceBucket, String sourceKey, String dest
6472
CopyObjectRequest copyRequest = new CopyObjectRequest(sourceBucket, sourceKey, destinationBucket, destinationKey);
6573
Copy copy2 = tm.copy(copyRequest);
6674
}
75+
76+
void downloadDirectory(TransferManager tm, File destination) {
77+
MultipleFileDownload fileDownload = tm.downloadDirectory("bucket", "key", destination);
78+
tm.shutdownNow();
79+
}
80+
81+
void resume(TransferManager tm, PersistableDownload persistableDownload, PersistableUpload persistableUpload) {
82+
Download download = tm.resumeDownload(persistableDownload);
83+
Upload upload = tm.resumeUpload(persistableUpload);
84+
}
6785
}

‎v2-migration/src/main/java/software/amazon/awssdk/v2migration/TransferManagerMethodsToV2.java‎

Lines changed: 82 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,17 +16,21 @@
1616
package software.amazon.awssdk.v2migration;
1717

1818
import static software.amazon.awssdk.v2migration.internal.utils.S3TransformUtils.V2_S3_MODEL_PKG;
19+
import static software.amazon.awssdk.v2migration.internal.utils.S3TransformUtils.V2_TM_CLIENT;
1920
import static software.amazon.awssdk.v2migration.internal.utils.S3TransformUtils.V2_TM_MODEL_PKG;
2021
import static software.amazon.awssdk.v2migration.internal.utils.S3TransformUtils.v2TmMethodMatcher;
2122

23+
import java.util.regex.Pattern;
2224
import org.openrewrite.ExecutionContext;
2325
import org.openrewrite.Recipe;
2426
import org.openrewrite.TreeVisitor;
2527
import org.openrewrite.java.AddImport;
26-
import org.openrewrite.java.JavaIsoVisitor;
2728
import org.openrewrite.java.JavaTemplate;
29+
import org.openrewrite.java.JavaVisitor;
2830
import org.openrewrite.java.MethodMatcher;
31+
import org.openrewrite.java.tree.Expression;
2932
import org.openrewrite.java.tree.J;
33+
import org.openrewrite.java.tree.JavaType;
3034
import software.amazon.awssdk.annotations.SdkInternalApi;
3135

3236
@SdkInternalApi
@@ -47,6 +51,13 @@ public class TransferManagerMethodsToV2 extends Recipe {
4751
private static final MethodMatcher COPY_BUCKET_KEY =
4852
v2TmMethodMatcher("copy(String, String, String, String");
4953

54+
private static final MethodMatcher DOWNLOAD_DIR = v2TmMethodMatcher("downloadDirectory(String, String, java.io.File)");
55+
56+
private static final Pattern S3_TM_CREDENTIAL = Pattern.compile(V2_TM_CLIENT);
57+
private static final Pattern V2_AWSCREDENTAIL = Pattern.compile("software.amazon.awssdk.auth.credentials.AwsCredentials");
58+
private static final Pattern V2_CREDENTIAL_PROVIDER = Pattern.compile("software.amazon.awssdk.auth.credentials"
59+
+ ".AwsCredentialsProvider");
60+
5061
@Override
5162
public String getDisplayName() {
5263
return "Transfer Manager Methods to V2";
@@ -62,10 +73,10 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
6273
return new Visitor();
6374
}
6475

65-
private static final class Visitor extends JavaIsoVisitor<ExecutionContext> {
76+
private static final class Visitor extends JavaVisitor<ExecutionContext> {
6677

6778
@Override
68-
public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, ExecutionContext executionContext) {
79+
public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext executionContext) {
6980

7081
if (DOWNLOAD_BUCKET_KEY_FILE.matches(method, false)) {
7182
method = transformDownloadWithBucketKeyFile(method);
@@ -95,10 +106,70 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
95106
method = transformUploadWithBucketKeyFile(method);
96107
return super.visitMethodInvocation(method, executionContext);
97108
}
109+
if (DOWNLOAD_DIR.matches(method, false)) {
110+
method = transformDownloadDirectory(method);
111+
return super.visitMethodInvocation(method, executionContext);
112+
}
98113

99114
return super.visitMethodInvocation(method, executionContext);
100115
}
101116

117+
@Override
118+
public J visitNewClass(J.NewClass newClass, ExecutionContext executionContext) {
119+
JavaType type = newClass.getType();
120+
if (!(type instanceof JavaType.FullyQualified)) {
121+
return newClass;
122+
}
123+
124+
if (type.isAssignableFrom(S3_TM_CREDENTIAL) &&
125+
newClass.getArguments().size() == 1 &&
126+
newClass.getArguments().get(0).getType() != null) {
127+
Expression arg = newClass.getArguments().get(0);
128+
if (arg.getType().isAssignableFrom(V2_AWSCREDENTAIL)) {
129+
addS3AsyncClientImport();
130+
addStaticCredentialsProviderImport();
131+
132+
return JavaTemplate
133+
.builder("S3TransferManager.builder()" +
134+
".s3Client(S3AsyncClient.builder()" +
135+
".credentialsProvider(StaticCredentialsProvider.create(#{any()}))" +
136+
".build())" +
137+
".build()")
138+
.build()
139+
.apply(getCursor(), newClass.getCoordinates().replace(), arg);
140+
}
141+
if (arg.getType().isAssignableFrom(V2_CREDENTIAL_PROVIDER)) {
142+
addS3AsyncClientImport();
143+
144+
return JavaTemplate
145+
.builder("S3TransferManager.builder()" +
146+
".s3Client(S3AsyncClient.builder()" +
147+
".credentialsProvider(#{any()})" +
148+
".build())" +
149+
".build()")
150+
.build()
151+
.apply(getCursor(), newClass.getCoordinates().replace(), arg);
152+
}
153+
}
154+
155+
return super.visitNewClass(newClass, executionContext);
156+
}
157+
158+
private J.MethodInvocation transformDownloadDirectory(J.MethodInvocation method) {
159+
String v2Method = "#{any()}.downloadDirectory(DownloadDirectoryRequest.builder()"
160+
+ ".bucket(#{any()}).listObjectsV2RequestTransformer(builder -> builder.prefix(#{any()}))"
161+
+ ".destination(#{any()}.toPath()).build())";
162+
163+
method = JavaTemplate.builder(v2Method).build()
164+
.apply(getCursor(), method.getCoordinates().replace(), method.getSelect(),
165+
method.getArguments().get(0), method.getArguments().get(1),
166+
method.getArguments().get(2));
167+
168+
addTmImport("DirectoryDownload");
169+
addTmImport("DownloadDirectoryRequest");
170+
return method;
171+
}
172+
102173
private J.MethodInvocation transformUploadWithBucketKeyFile(J.MethodInvocation method) {
103174
String v2Method = "#{any()}.uploadFile(UploadFileRequest.builder()"
104175
+ ".putObjectRequest(PutObjectRequest.builder().bucket(#{any()}).key(#{any()}).build())"
@@ -220,5 +291,13 @@ private void addDurationImport() {
220291
private void addRequestOverrideConfigImport() {
221292
doAfterVisit(new AddImport<>("software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration", null, false));
222293
}
294+
295+
private void addS3AsyncClientImport() {
296+
doAfterVisit(new AddImport<>("software.amazon.awssdk.services.s3.S3AsyncClient", null, false));
297+
}
298+
299+
private void addStaticCredentialsProviderImport() {
300+
doAfterVisit(new AddImport<>("software.amazon.awssdk.auth.credentials.StaticCredentialsProvider", null, false));
301+
}
223302
}
224303
}

‎v2-migration/src/main/resources/META-INF/rewrite/aws-sdk-java-v1-to-v2-with-tm.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,4 +53,5 @@ recipeList:
5353
- software.amazon.awssdk.v2migration.S3NonStreamingRequestToV2Complex
5454
- software.amazon.awssdk.v2migration.S3PutObjectRequestToV2
5555
- software.amazon.awssdk.v2migration.SettersToBuilderV2
56+
- software.amazon.awssdk.v2migration.ChangeTransferManagerSimpleMethods
5657
- software.amazon.awssdk.v2migration.TransferManagerMethodsToV2

‎v2-migration/src/main/resources/META-INF/rewrite/change-transfer-manager-simple-methods.yml‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
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+
type: specs.openrewrite.org/v1beta/recipe
16+
name: software.amazon.awssdk.v2migration.ChangeTransferManagerSimpleMethods
17+
displayName: Change TransferManager simple methods to v2.
18+
description: Change TransferManager simple methods to v2.
19+
recipeList:
20+
- org.openrewrite.java.ChangeMethodName:
21+
methodPattern: software.amazon.awssdk.transfer.s3.S3TransferManager resumeDownload(..)
22+
newMethodName: resumeDownloadFile
23+
- org.openrewrite.java.ChangeMethodName:
24+
methodPattern: software.amazon.awssdk.transfer.s3.S3TransferManager resumeUpload(..)
25+
newMethodName: resumeUploadFile
26+
- org.openrewrite.java.ChangeMethodName:
27+
methodPattern: software.amazon.awssdk.transfer.s3.S3TransferManager shutdownNow()
28+
newMethodName: close

‎v2-migration/src/main/resources/META-INF/rewrite/change-transfer-manager-types.yml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,12 @@ name: software.amazon.awssdk.v2migration.ChangeTransferManagerTypes
1818
displayName: Change SDK TransferManager types from v1 to v2
1919
description: Change SDK TransferManager types from v1 to v2.
2020
recipeList:
21+
- software.amazon.awssdk.v2migration.openrewrite.ChangeMethodInvocationReturnType:
22+
methodPattern: com.amazonaws.services.s3.transfer.TransferManager resumeDownload(..)
23+
newReturnType: software.amazon.awssdk.transfer.s3.model.FileDownload
24+
- software.amazon.awssdk.v2migration.openrewrite.ChangeMethodInvocationReturnType:
25+
methodPattern: com.amazonaws.services.s3.transfer.TransferManager resumeUpload(..)
26+
newReturnType: software.amazon.awssdk.transfer.s3.model.FileUpload
2127
- org.openrewrite.java.ChangeType:
2228
oldFullyQualifiedTypeName: com.amazonaws.services.s3.transfer.TransferManager
2329
newFullyQualifiedTypeName: software.amazon.awssdk.transfer.s3.S3TransferManager

0 commit comments

Comments
 (0)