io.kestra.plugin.nats.Produce Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of plugin-nats Show documentation
Kestra's NATS Plugin enables high-performance communication in distributed systems for enhanced workflow management and seamless messaging.
There is a newer version: 0.20.0
Show newest version
package io.kestra.plugin.nats;

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.tasks.RunnableTask;
import io.kestra.core.runners.RunContext;
import io.kestra.core.serializers.FileSerde;
import io.nats.client.Connection;
import io.nats.client.Message;
import io.nats.client.impl.Headers;
import io.nats.client.impl.NatsMessage;
import io.swagger.v3.oas.annotations.media.Schema;
import lombok.*;
import lombok.experimental.SuperBuilder;

import jakarta.validation.constraints.NotBlank;
import jakarta.validation.constraints.NotNull;
import reactor.core.publisher.Flux;

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

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

@SuperBuilder
@ToString
@EqualsAndHashCode
@Getter
@NoArgsConstructor
@Schema(
    title = "Produce messages to a NATS subject on a NATS server."
)
@Plugin(
    examples = {
        @Example(
            title = "Produce a single message to kestra.publish subject, using user password authentication.",
            full = true,
            code = """
                id: nats_produce_single_message
                namespace: company.team
                
                tasks:
                  - id: produce
                    type: io.kestra.plugin.nats.Produce
                    url: nats://localhost:4222
                    username: nats_user
                    password: nats_password
                    subject: kestra.publish
                    from:
                      headers:
                        someHeaderKey: someHeaderValue
                      data: Some message
                """
        ),
        @Example(
            title = "Produce 2 messages to kestra.publish subject, using user password authentication.",
            full = true,
            code = """
                id: nats_produce_two_messages
                namespace: company.team
                
                tasks:
                  - id: produce
                    type: io.kestra.plugin.nats.Produce
                    url: nats://localhost:4222
                    username: nats_user
                    password: nats_password
                    subject: kestra.publish
                    from:
                      - headers:
                          someHeaderKey: someHeaderValue
                        data: Some message
                      - data: Another message
                """
        ),
        @Example(
            title = "Produce messages (1 / row) from an internal storage file to kestra.publish subject, using user password authentication.",
            full = true,
            code = """
                id: nats_produce_messages_from_file
                namespace: company.team
                
                tasks:
                  - id: produce
                    type: io.kestra.plugin.nats.Produce
                    url: nats://localhost:4222
                    username: nats_user
                    password: nats_password
                    subject: kestra.publish
                    from: "{{ outputs.some_task_with_output_file.uri }}"
                """
        ),
    }
)
public class Produce extends NatsConnection implements RunnableTask {
    @Schema(
        title = "Subject to produce message to"
    )
    @PluginProperty(dynamic = true)
    @NotBlank
    @NotNull
    private String subject;
    @Schema(
        title = "Source of message(s) to send",
        description = "Can be an internal storage uri, a map or a list." +
            "with the following format: headers, data",
        anyOf = {String.class, List.class, Map.class}
    )
    @NotNull
    @PluginProperty(dynamic = true)
    private Object from;

    public Output run(RunContext runContext) throws Exception {
        Connection connection = connect(runContext);

        int messagesCount;

        if (this.from instanceof String || this.from instanceof List) {
            if (this.from instanceof String fromStr) {
                URI from = new URI(runContext.render(fromStr));
                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)))) {
                    messagesCount = publish(runContext, connection, FileSerde.readAll(inputStream));
                }
            } else {
                messagesCount = publish(runContext, connection, Flux.fromIterable(((List) this.from)));
            }

        } else {
            connection.publish(this.producerMessage(runContext.render(this.subject), runContext.render((Map) this.from)));
            messagesCount = 1;
        }

        connection.flushBuffer();
        connection.close();

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

    private Integer publish(RunContext runContext, Connection connection, Flux