io.kestra.plugin.aws.sqs.Publish Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of plugin-aws Show documentation
AWS plugin for Kestra
The newest version!
package io.kestra.plugin.aws.sqs;

import io.kestra.core.exceptions.IllegalVariableEvaluationException;
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.executions.metrics.Counter;
import io.kestra.core.models.tasks.RunnableTask;
import io.kestra.core.runners.RunContext;
import io.kestra.core.serializers.FileSerde;
import io.kestra.core.serializers.JacksonMapper;
import io.kestra.plugin.aws.sqs.model.Message;
import io.swagger.v3.oas.annotations.media.Schema;
import lombok.*;
import lombok.experimental.SuperBuilder;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;
import software.amazon.awssdk.services.sqs.SqsClient;
import software.amazon.awssdk.services.sqs.model.SendMessageRequest;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.net.URI;
import java.util.List;
import java.util.Map;

import jakarta.validation.constraints.NotNull;

import static io.kestra.core.utils.Rethrow.throwFunction;

@SuperBuilder
@ToString
@EqualsAndHashCode
@Getter
@NoArgsConstructor
@Schema(
    title = "Publish a message to an SQS queue."
)
@Plugin(
    examples = {
        @Example(
            full = true,
            code = """
                id: aws_sqs_publish
                namespace: company.team

                tasks:
                  - id: publish
                    type: io.kestra.plugin.aws.sqs.Publish
                    accessKeyId: ""
                    secretKeyId: ""
                    region: "eu-central-1"
                    queueUrl: "https://sqs.eu-central-1.amazonaws.com/000000000000/test-queue"
                    from:
                    - data: Hello World
                    - data: Hello Kestra
                      delaySeconds: 5
                """
        )
    }
)
public class Publish extends AbstractSqs implements RunnableTask {
    @PluginProperty(dynamic = true)
    @NotNull
    @Schema(
        title = "The source of the published data.",
        description = "Can be an internal storage URI, a list of SQS messages, or a single SQS message.",
        anyOf = {String.class, List.class, Message.class}
    )
    private Object from;

    @SuppressWarnings("unchecked")
    @Override
    public Output run(RunContext runContext) throws Exception {
        var queueUrl = runContext.render(getQueueUrl());
        try (var sqsClient = this.client(runContext)) {
            Integer count;
            Flux flowable;
            Flux resultFlowable;

            if (this.from instanceof String) {
                URI from = new URI(runContext.render((String) this.from));
                if (!from.getScheme().equals("kestra")) {
                    throw new Exception("Invalid from parameter, must be a Kestra internal storage URI");
                }


                try (BufferedReader inputStream = new BufferedReader(new InputStreamReader(runContext.storage().getFile(from)))) {
                    flowable = FileSerde.readAll(inputStream, Message.class);
                    resultFlowable = this.buildFlowable(flowable, sqsClient, queueUrl, runContext);

                    count = resultFlowable.reduce(Integer::sum).blockOptional().orElse(0);
                }

            } else if (this.from instanceof List) {
                flowable = Flux
                    .fromIterable((List) this.from)
                    .map(map -> JacksonMapper.toMap(map, Message.class));

                resultFlowable = this.buildFlowable(flowable, sqsClient, queueUrl, runContext);

                count = resultFlowable.reduce(Integer::sum).blockOptional().orElse(0);
            } else {
                var msg = JacksonMapper.toMap(this.from, Message.class);
                sqsClient.sendMessage(msg.to(SendMessageRequest.builder().queueUrl(queueUrl), runContext));

                count = 1;
            }

            // metrics
            runContext.metric(Counter.of("records", count, "queue", queueUrl));

            return Output.builder()
                .messagesCount(count)
                .build();
        }
    }

    private Flux buildFlowable(Flux flowable, SqsClient sqsClient, String queueUrl, RunContext runContext) throws IllegalVariableEvaluationException {
        return flowable
            .map(throwFunction(message -> {
                sqsClient.sendMessage(message.to(SendMessageRequest.builder().queueUrl(queueUrl), runContext));
                return 1;
            }));
    }

    @Builder
    @Getter
    public static class Output implements io.kestra.core.models.tasks.Output {
        @Schema(title = "Number of published messages.")
        private final Integer messagesCount;
    }
}