All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.kestra.plugin.aws.dynamodb.AbstractDynamoDb Maven / Gradle / Ivy

There is a newer version: 0.20.0
Show newest version
package io.kestra.plugin.aws.dynamodb;

import io.kestra.core.exceptions.IllegalVariableEvaluationException;
import io.kestra.core.models.annotations.PluginProperty;
import io.kestra.core.models.executions.metrics.Counter;
import io.kestra.core.models.tasks.common.FetchOutput;
import io.kestra.core.models.tasks.common.FetchType;
import io.kestra.core.runners.RunContext;
import io.kestra.core.serializers.FileSerde;
import io.kestra.plugin.aws.AbstractConnection;
import io.kestra.plugin.aws.ConnectionUtils;
import io.swagger.v3.oas.annotations.media.Schema;
import jakarta.validation.constraints.NotNull;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.ToString;
import lombok.experimental.SuperBuilder;
import org.apache.commons.lang3.tuple.Pair;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;

import java.io.*;
import java.net.URI;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Stream;

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

@SuperBuilder
@ToString
@EqualsAndHashCode
@Getter
@NoArgsConstructor
public abstract class AbstractDynamoDb extends AbstractConnection {
    @Schema(title = "The DynamoDB table name.")
    @PluginProperty(dynamic = true)
    @NotNull
    protected String tableName;

    protected DynamoDbClient client(final RunContext runContext) throws IllegalVariableEvaluationException {
        final AwsClientConfig clientConfig = awsClientConfig(runContext);
        return ConnectionUtils.configureSyncClient(clientConfig, DynamoDbClient.builder()).build();
    }

    protected Map objectMapFrom(Map fields) {
        Map row = new HashMap<>();
        for(var field : fields.entrySet()) {
            var key = field.getKey();
            var value = field.getValue();
            row.put(key, objectFrom(value));
        }
        return row;
    }

    protected Object objectFrom(AttributeValue value) {
        if(value == null || (value.nul() != null && value.hasSs())){
            return null;
        }
        if(value.bool() != null && value.bool()) {
            return true;
        }
        if(value.hasSs()) {
            return value.ss();
        }
        if(value.hasM()) {
            return objectMapFrom(value.m());
        }

        //we may miss some cases, but it should be good for a first implementation.
        return value.s();
    }

    protected Map valueMapFrom(Map fields) {
        Map item = new HashMap<>();
        for(var field : fields.entrySet()) {
            var key = field.getKey();
            var value = field.getValue();
            item.put(key, objectFrom(value));
        }
        return item;
    }

    @SuppressWarnings("unchecked")
    protected AttributeValue objectFrom(Object value) {
        if(value == null){
            return AttributeValue.fromNul(true);
        }
        if(value instanceof String) {
            return AttributeValue.fromS((String) value);
        }
        if(value instanceof Boolean) {
            return AttributeValue.fromBool((Boolean) value);
        }
        if(value instanceof List) {
            return AttributeValue.fromSs((List) value);
        }
        if(value instanceof Map) {
            return AttributeValue.fromM(valueMapFrom((Map) value));
        }

        // in case we don't have any class we can handle, we call toString()
        return AttributeValue.fromS(value.toString());
    }

    protected FetchOutput fetchOutputs(List> items, FetchType fetchType, RunContext runContext) throws IOException {
        var outputBuilder = FetchOutput.builder();
        switch (fetchType) {
            case FETCH:
                Pair, Long> fetch = this.fetch(items);
                outputBuilder
                    .rows(fetch.getLeft())
                    .size(fetch.getRight());
                break;

            case FETCH_ONE:
                var o = this.fetchOne(items);

                outputBuilder
                    .row(o)
                    .size(o != null ? 1L : 0L);
                break;

            case STORE:
                Pair store = this.store(runContext, items);
                outputBuilder
                    .uri(store.getLeft())
                    .size(store.getRight());
                break;
        }

        var output = outputBuilder.build();

        runContext.metric(Counter.of(
            "records", output.getSize(),
            "tableName", getTableName()
        ));

        return output;
    }

    private Pair store(RunContext runContext, List> items) throws IOException {
        File tempFile = runContext.workingDir().createTempFile(".ion").toFile();

        try (var output = new BufferedWriter(new FileWriter(tempFile), FileSerde.BUFFER_SIZE)) {
            var flux = Flux.fromIterable(items).map(attributes -> objectMapFrom(attributes));
            Long count = FileSerde.writeAll(output, flux).block();
            return Pair.of(
                runContext.storage().putFile(tempFile),
                count
            );
        }
    }

    private Pair, Long> fetch(List> items) {
        List result = new ArrayList<>();
        AtomicLong count = new AtomicLong();

        items.forEach(throwConsumer(attributes -> {
            count.incrementAndGet();
            result.add(objectMapFrom(attributes));
        }));

        return Pair.of(result, count.get());
    }


    private Map fetchOne(List> items) {
        return items.stream()
            .findFirst()
            .map(this::objectMapFrom)
            .orElse(null);
    }
}