fix(aws): render variable inside SNS and SQS Publish tasks (#152)

kestra-io · May 20, 2023 · c63c7ab · c63c7ab
1 parent ab552e5
commit c63c7ab
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 21 deletions.
diff --git a/src/main/java/io/kestra/plugin/aws/sns/Publish.java b/src/main/java/io/kestra/plugin/aws/sns/Publish.java
@@ -69,7 +69,7 @@ public Publish.Output run(RunContext runContext) throws Exception {
 
                 try (BufferedReader inputStream = new BufferedReader(new InputStreamReader(runContext.uriToInputStream(from)))) {
                     flowable = Flowable.create(FileSerde.reader(inputStream, Message.class), BackpressureStrategy.BUFFER);
-                    resultFlowable = this.buildFlowable(flowable, snsClient);
+                    resultFlowable = this.buildFlowable(flowable, snsClient, runContext);
 
                     count = resultFlowable.reduce(Integer::sum).blockingGet();
                 }
@@ -79,12 +79,12 @@ public Publish.Output run(RunContext runContext) throws Exception {
                     .fromArray(((List<Message>) this.from).toArray())
                     .cast(Message.class);
 
-                resultFlowable = this.buildFlowable(flowable, snsClient);
+                resultFlowable = this.buildFlowable(flowable, snsClient, runContext);
 
                 count = resultFlowable.reduce(Integer::sum).blockingGet();
             } else {
                 var msg = JacksonMapper.toMap(this.from, Message.class);
-                snsClient.publish(msg.to(PublishRequest.builder().topicArn(getTopicArn())));
+                snsClient.publish(msg.to(PublishRequest.builder().topicArn(getTopicArn()), runContext));
 
                 count = 1;
             }
@@ -98,10 +98,10 @@ public Publish.Output run(RunContext runContext) throws Exception {
         }
     }
 
-    private Flowable<Integer> buildFlowable(Flowable<Message> flowable, SnsClient snsClient) {
+    private Flowable<Integer> buildFlowable(Flowable<Message> flowable, SnsClient snsClient, RunContext runContext) {
         return flowable
             .map(message -> {
-                snsClient.publish(message.to(PublishRequest.builder().topicArn(getTopicArn())));
+                snsClient.publish(message.to(PublishRequest.builder().topicArn(getTopicArn()), runContext));
                 return 1;
             });
     }

diff --git a/src/main/java/io/kestra/plugin/aws/sns/model/Message.java b/src/main/java/io/kestra/plugin/aws/sns/model/Message.java
@@ -1,6 +1,8 @@
 package io.kestra.plugin.aws.sns.model;
 
+import io.kestra.core.exceptions.IllegalVariableEvaluationException;
 import io.kestra.core.models.annotations.PluginProperty;
+import io.kestra.core.runners.RunContext;
 import io.swagger.v3.oas.annotations.media.Schema;
 import lombok.Builder;
 import lombok.Getter;
@@ -39,14 +41,14 @@ public class Message {
     private String structure;
 
 
-    public PublishRequest to(PublishRequest.Builder builder) {
+    public PublishRequest to(PublishRequest.Builder builder, RunContext runContext) throws IllegalVariableEvaluationException {
         return builder
-            .message(data)
-            .messageGroupId(groupId)
-            .messageDeduplicationId(deduplicationId)
-            .subject(subject)
-            .phoneNumber(phoneNumber)
-            .messageStructure(structure)
+            .message(runContext.render(data))
+            .messageGroupId(runContext.render(groupId))
+            .messageDeduplicationId(runContext.render(deduplicationId))
+            .subject(runContext.render(subject))
+            .phoneNumber(runContext.render(phoneNumber))
+            .messageStructure(runContext.render(structure))
             .build();
     }
 }
diff --git a/src/main/java/io/kestra/plugin/aws/sqs/Publish.java b/src/main/java/io/kestra/plugin/aws/sqs/Publish.java
@@ -69,7 +69,7 @@ public Output run(RunContext runContext) throws Exception {
 
                 try (BufferedReader inputStream = new BufferedReader(new InputStreamReader(runContext.uriToInputStream(from)))) {
                     flowable = Flowable.create(FileSerde.reader(inputStream, Message.class), BackpressureStrategy.BUFFER);
-                    resultFlowable = this.buildFlowable(flowable, sqsClient);
+                    resultFlowable = this.buildFlowable(flowable, sqsClient, runContext);
 
                     count = resultFlowable.reduce(Integer::sum).blockingGet();
                 }
@@ -79,12 +79,12 @@ public Output run(RunContext runContext) throws Exception {
                     .fromArray(((List<Message>) this.from).toArray())
                     .cast(Message.class);
 
-                resultFlowable = this.buildFlowable(flowable, sqsClient);
+                resultFlowable = this.buildFlowable(flowable, sqsClient, runContext);
 
                 count = resultFlowable.reduce(Integer::sum).blockingGet();
             } else {
                 var msg = JacksonMapper.toMap(this.from, Message.class);
-                sqsClient.sendMessage(msg.to(SendMessageRequest.builder().queueUrl(getQueueUrl())));
+                sqsClient.sendMessage(msg.to(SendMessageRequest.builder().queueUrl(getQueueUrl()), runContext));
 
                 count = 1;
             }
@@ -98,10 +98,10 @@ public Output run(RunContext runContext) throws Exception {
         }
     }
 
-    private Flowable<Integer> buildFlowable(Flowable<Message> flowable, SqsClient sqsClient) {
+    private Flowable<Integer> buildFlowable(Flowable<Message> flowable, SqsClient sqsClient, RunContext runContext) {
         return flowable
             .map(message -> {
-                sqsClient.sendMessage(message.to(SendMessageRequest.builder().queueUrl(getQueueUrl())));
+                sqsClient.sendMessage(message.to(SendMessageRequest.builder().queueUrl(getQueueUrl()), runContext));
                 return 1;
             });
     }

diff --git a/src/main/java/io/kestra/plugin/aws/sqs/model/Message.java b/src/main/java/io/kestra/plugin/aws/sqs/model/Message.java
@@ -1,6 +1,8 @@
 package io.kestra.plugin.aws.sqs.model;
 
+import io.kestra.core.exceptions.IllegalVariableEvaluationException;
 import io.kestra.core.models.annotations.PluginProperty;
+import io.kestra.core.runners.RunContext;
 import io.swagger.v3.oas.annotations.media.Schema;
 import lombok.Builder;
 import lombok.Getter;
@@ -30,11 +32,11 @@ public class Message {
     @PluginProperty
     private Integer delaySeconds;
 
-    public SendMessageRequest to(SendMessageRequest.Builder builder) {
+    public SendMessageRequest to(SendMessageRequest.Builder builder, RunContext runContext) throws IllegalVariableEvaluationException {
         return builder
-            .messageBody(data)
-            .messageGroupId(groupId)
-            .messageDeduplicationId(deduplicationId)
+            .messageBody(runContext.render(data))
+            .messageGroupId(runContext.render(groupId))
+            .messageDeduplicationId(runContext.render(deduplicationId))
             .delaySeconds(delaySeconds)
             .build();
     }