io.kestra.plugin.mongodb.Load Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of plugin-mongodb Show documentation
Access and manipulate MongoDB data within Kestra workflows.
There is a newer version: 0.21.0
package io.kestra.plugin.mongodb;

import com.mongodb.client.model.InsertOneModel;
import com.mongodb.client.model.WriteModel;
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.runners.RunContext;
import io.kestra.core.serializers.FileSerde;
import io.kestra.core.serializers.JacksonMapper;
import io.swagger.v3.oas.annotations.media.Schema;
import lombok.*;
import lombok.experimental.SuperBuilder;
import org.bson.BsonDocument;
import org.bson.BsonObjectId;
import org.bson.conversions.Bson;
import org.bson.types.ObjectId;
import reactor.core.publisher.Flux;

import java.io.BufferedReader;
import java.util.Map;

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

@SuperBuilder
@ToString
@EqualsAndHashCode
@Getter
@NoArgsConstructor
@Schema(
    title = "Bulk load documents in MongoDB using Kestra internal storage file."
)
@Plugin(
    examples = {
        @Example(
            full = true,
            code = """
                id: mongodb_load
                namespace: company.team

                inputs:
                  - id: file
                    type: FILE
                
                tasks:
                  - id: load
                    type: io.kestra.plugin.mongodb.Load
                    connection:
                      uri: "mongodb://root:example@localhost:27017/?authSource=admin"
                    database: "my_database"
                    collection: "my_collection"
                    from: "{{ inputs.file }}"
                """
        )
    }
)
public class Load extends AbstractLoad {
    @Schema(
        title = "Use this key as ID."
    )
    @PluginProperty(dynamic = true)
    private String idKey;

    @Schema(
        title = "Whether to remove idKey from the final document."
    )
    @PluginProperty(dynamic = true)
    @Builder.Default
    private Boolean removeIdKey = true;

    @SuppressWarnings("unchecked")
    @Override
    protected Flux> source(RunContext runContext, BufferedReader inputStream) throws Exception {
        return FileSerde.readAll(inputStream)
            .map(throwFunction(o -> {
                Map values = (Map) o;

                if (this.idKey != null) {
                    String idKey = runContext.render(this.idKey);

                    values.put(
                        "_id",
                        new BsonObjectId(new ObjectId(values.get(idKey).toString()))
                    );

                    if (this.removeIdKey) {
                        values.remove(idKey);
                    }
                }

                return new InsertOneModel<>(
                    BsonDocument.parse(JacksonMapper.ofJson().writeValueAsString(values))
                );
            }));
    }
}