Allow to use DefaultAsyncClient (#237)

* feat: add configuration to use DefaultAsyncClient * feat: update configuration name --------- Co-authored-by: Sébastien ALLEMAND <[email protected]>
kestra-io · Sep 1, 2023 · 22f6485 · 22f6485
1 parent cddd445
commit 22f6485
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 18 deletions.
diff --git a/src/main/java/io/kestra/plugin/aws/AbstractConnection.java b/src/main/java/io/kestra/plugin/aws/AbstractConnection.java
@@ -26,6 +26,8 @@ public abstract class AbstractConnection extends Task implements AbstractConnect
 
     protected String endpointOverride;
 
+    private Boolean compatibilityMode;
+
     protected AwsCredentialsProvider credentials(RunContext runContext) throws IllegalVariableEvaluationException {
         String accessKeyId = runContext.render(this.accessKeyId);
         String secretKeyId = runContext.render(this.secretKeyId);

diff --git a/src/main/java/io/kestra/plugin/aws/s3/AbstractS3.java b/src/main/java/io/kestra/plugin/aws/s3/AbstractS3.java
@@ -19,7 +19,7 @@
 @EqualsAndHashCode
 @Getter
 @NoArgsConstructor
-public abstract class AbstractS3 extends AbstractConnection  {
+public abstract class AbstractS3 extends AbstractConnection {
 
     protected S3Client client(RunContext runContext) throws IllegalVariableEvaluationException {
         S3ClientBuilder s3ClientBuilder = S3Client.builder()
@@ -36,18 +36,36 @@ protected S3Client client(RunContext runContext) throws IllegalVariableEvaluatio
 
         return s3ClientBuilder.build();
     }
+
     protected S3AsyncClient asyncClient(RunContext runContext) throws IllegalVariableEvaluationException {
-        S3CrtAsyncClientBuilder s3ClientBuilder = S3AsyncClient.crtBuilder()
-            .credentialsProvider(this.credentials(runContext));
 
-        if (this.region != null) {
-            s3ClientBuilder.region(Region.of(runContext.render(this.region)));
-        }
+        if (this.getCompatibilityMode()) {
+            S3AsyncClientBuilder s3ClientBuilder = S3AsyncClient.builder()
+                .credentialsProvider(this.credentials(runContext));
 
-        if (this.endpointOverride != null) {
-            s3ClientBuilder.endpointOverride(URI.create(runContext.render(this.endpointOverride)));
+            if (this.region != null) {
+                s3ClientBuilder.region(Region.of(runContext.render(this.region)));
+            }
+
+            if (this.endpointOverride != null) {
+                s3ClientBuilder.endpointOverride(URI.create(runContext.render(this.endpointOverride)));
+            }
+            return s3ClientBuilder.build();
+
+        } else {
+            S3CrtAsyncClientBuilder s3ClientBuilder = S3AsyncClient.crtBuilder()
+                .credentialsProvider(this.credentials(runContext));
+
+            if (this.region != null) {
+                s3ClientBuilder.region(Region.of(runContext.render(this.region)));
+            }
+
+            if (this.endpointOverride != null) {
+                s3ClientBuilder.endpointOverride(URI.create(runContext.render(this.endpointOverride)));
+            }
+
+            return s3ClientBuilder.build();
         }
 
-        return s3ClientBuilder.build();
     }
 }
diff --git a/src/main/java/io/kestra/plugin/aws/s3/Download.java b/src/main/java/io/kestra/plugin/aws/s3/Download.java
@@ -6,10 +6,7 @@
 import io.kestra.core.models.tasks.RunnableTask;
 import io.kestra.core.runners.RunContext;
 import io.swagger.v3.oas.annotations.media.Schema;
-import lombok.EqualsAndHashCode;
-import lombok.Getter;
-import lombok.NoArgsConstructor;
-import lombok.ToString;
+import lombok.*;
 import lombok.experimental.SuperBuilder;
 import org.apache.commons.lang3.tuple.Pair;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
@@ -54,6 +51,13 @@ public class Download extends AbstractS3Object implements RunnableTask<Download.
     @PluginProperty(dynamic = true)
     protected String versionId;
 
+    @Schema(
+        title = "this property will use the AsynS3Client instead of the S3CrtAsynClient which maximize compatibility with S3-compatible services but restrict uploads and downloads to 2GB"
+    )
+    @PluginProperty
+    @Builder.Default
+    private Boolean compatibilityMode = false;
+
     @Override
     public Output run(RunContext runContext) throws Exception {
         String bucket = runContext.render(this.bucket);

diff --git a/src/main/java/io/kestra/plugin/aws/s3/Downloads.java b/src/main/java/io/kestra/plugin/aws/s3/Downloads.java
@@ -2,6 +2,7 @@
 
 import io.kestra.core.models.annotations.Example;
 import io.kestra.core.models.annotations.Plugin;
+import io.kestra.core.models.annotations.PluginProperty;
 import io.kestra.core.models.tasks.RunnableTask;
 import io.kestra.core.runners.RunContext;
 import io.kestra.plugin.aws.s3.models.S3Object;
@@ -51,6 +52,14 @@ public class Downloads extends AbstractS3Object implements RunnableTask<List.Out
     @Builder.Default
     private Integer maxKeys = 1000;
 
+    @Schema(
+        title = "this property will use the AsynS3Client instead of the S3CrtAsynClient which maximize compatibility with S3-compatible services but restrict uploads and downloads to 2GB"
+    )
+    @PluginProperty
+    @Builder.Default
+    private Boolean compatibilityMode = false;
+
+
     private String expectedBucketOwner;
 
     protected String regexp;

diff --git a/src/main/java/io/kestra/plugin/aws/s3/Upload.java b/src/main/java/io/kestra/plugin/aws/s3/Upload.java
@@ -7,16 +7,12 @@
 import io.kestra.core.models.tasks.RunnableTask;
 import io.kestra.core.runners.RunContext;
 import io.swagger.v3.oas.annotations.media.Schema;
-import lombok.EqualsAndHashCode;
-import lombok.Getter;
-import lombok.NoArgsConstructor;
-import lombok.ToString;
+import lombok.*;
 import lombok.experimental.SuperBuilder;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectResponse;
 import software.amazon.awssdk.transfer.s3.S3TransferManager;
-import software.amazon.awssdk.transfer.s3.model.CompletedFileUpload;
 import software.amazon.awssdk.transfer.s3.model.FileUpload;
 import software.amazon.awssdk.transfer.s3.model.UploadFileRequest;
 
@@ -67,6 +63,13 @@ public class Upload extends AbstractS3Object implements RunnableTask<Upload.Outp
     @PluginProperty
     private Map<String, String> metadata;
 
+    @Schema(
+        title = "this property will use the AsynS3Client instead of the S3CrtAsynClient which maximize compatibility with S3-compatible services but restrict uploads and downloads to 2GB"
+    )
+    @PluginProperty
+    @Builder.Default
+    private Boolean compatibilityMode = false;
+
     @Schema(
         title = "If you don't specify, S3 Standard is the default storage class. Amazon S3 supports other storage classes."
     )
@@ -79,6 +82,7 @@ public Output run(RunContext runContext) throws Exception {
         String key = runContext.render(this.key);
 
         try (S3AsyncClient client = this.asyncClient(runContext)) {
+
             File tempFile = runContext.tempFile().toFile();
             URI from = new URI(runContext.render(this.from));
             Files.copy(runContext.uriToInputStream(from), tempFile.toPath(), StandardCopyOption.REPLACE_EXISTING);