io.kestra.plugin.compress.ArchiveDecompress Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of plugin-compress Show documentation
Manage file compression and decompression within Kestra workflows.
There is a newer version: 0.20.0
package io.kestra.plugin.compress;

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.swagger.v3.oas.annotations.media.Schema;
import lombok.*;
import lombok.experimental.SuperBuilder;
import org.apache.commons.compress.archivers.ArchiveEntry;
import org.apache.commons.compress.archivers.ArchiveInputStream;
import org.apache.commons.compress.compressors.CompressorInputStream;
import org.apache.commons.io.IOUtils;

import java.io.BufferedInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URI;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.HashMap;
import java.util.Map;
import javax.validation.constraints.NotNull;

@SuperBuilder
@ToString
@EqualsAndHashCode
@Getter
@NoArgsConstructor
@Schema(
    title = "Decompress an archive file"
)
@Plugin(
    examples = {
        @Example(
            code = {
                "from: \"{{ inputs.files }} \"",
                "algorithm: ZIP"
            }
        )
    }
)
public class ArchiveDecompress extends AbstractArchive implements RunnableTask {
    @NotNull
    @Schema(
        title = "The file internal storage uri"
    )
    @PluginProperty(dynamic = true)
    private String from;

    public Output run(RunContext runContext) throws Exception {
        Map files;

        URI from = new URI(runContext.render(this.from));
        try (
            InputStream fromInputStream = runContext.uriToInputStream(from);
            InputStream fromInputStreamBuffered = new BufferedInputStream(fromInputStream);
        ) {
            if (this.compression != null) {
                try (
                    CompressorInputStream compressorInputStream = this.compressorInputStream(this.compression, fromInputStreamBuffered);
                    ArchiveInputStream archiveInputStream = this.archiveInputStream(compressorInputStream);
                ) {
                    files = this.readArchive(runContext, archiveInputStream);
                }
            } else {
                try (ArchiveInputStream archiveInputStream = this.archiveInputStream(fromInputStreamBuffered)) {
                    files = this.readArchive(runContext, archiveInputStream);
                }
            }
        }

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

    private Map readArchive(RunContext runContext, ArchiveInputStream archiveInputStream) throws IOException {
        HashMap files = new HashMap<>();

        long size = 0;
        ArchiveEntry entry;
        while ((entry = archiveInputStream.getNextEntry()) != null) {
            if (!archiveInputStream.canReadEntryData(entry)) {
                throw new IOException("Unable to read entry '" + entry.getName() + "'");
            }

            if (!entry.isDirectory()) {
                Path path = runContext.tempFile();
                try (OutputStream o = Files.newOutputStream(path)) {
                    IOUtils.copy(archiveInputStream, o);
                }

                size = size + entry.getSize();

                files.put(entry.getName(), runContext.putTempFile(path.toFile()));
            }
        }

        runContext.metric(Counter.of("size", size));
        runContext.metric(Counter.of("count", files.size()));

        return files;
    }


    @Builder
    @Getter
    public static class Output implements io.kestra.core.models.tasks.Output {
        @Schema(
            title = "The url of the zip file on kestra storage"
        )
        @PluginProperty(additionalProperties = URI.class)
        private final Map files;
    }
}