embulk-decoder-commons-compress 0.3.0

data/src/main/java/org/embulk/decoder/ArchiveInputStreamIterator.java

@@ -0,0 +1,68 @@
+package org.embulk.decoder;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Iterator;
+
+import org.apache.commons.compress.archivers.ArchiveEntry;
+import org.apache.commons.compress.archivers.ArchiveInputStream;
+
+class ArchiveInputStreamIterator implements Iterator<InputStream> {
+    private ArchiveInputStream ain;
+    private ArchiveEntry entry;
+    private boolean endOfArchive = false;
+
+    ArchiveInputStreamIterator(ArchiveInputStream ain)
+    {
+        this.ain = ain;
+    }
+
+    @Override
+    public boolean hasNext() {
+        try {
+            return checkNext();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public InputStream next() {
+        try {
+            if (checkNext()) {
+                entry = null;
+            } else {
+                return null;
+            }
+
+            return ain;
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void remove() {
+        throw new UnsupportedOperationException();
+    }
+
+    private boolean checkNext() throws IOException {
+        if (endOfArchive) {
+            return false;
+        } else if (entry != null) {
+            return true;
+        }
+
+        while (true) {
+            entry = ain.getNextEntry();
+            if (entry == null) {
+                endOfArchive = true;
+                return false;
+            } else if (entry.isDirectory()) {
+                continue;
+            } else {
+                return true;
+            }
+        }
+    }
+}

data/src/main/java/org/embulk/decoder/CommonsCompressDecoderPlugin.java

@@ -0,0 +1,52 @@
+package org.embulk.decoder;
+
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.ConfigInject;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.Task;
+import org.embulk.config.TaskSource;
+import org.embulk.spi.BufferAllocator;
+import org.embulk.spi.DecoderPlugin;
+import org.embulk.spi.FileInput;
+import org.embulk.spi.util.FileInputInputStream;
+
+public class CommonsCompressDecoderPlugin
+        implements DecoderPlugin
+{
+    public interface PluginTask
+            extends Task
+    {
+        @Config("format")
+        @ConfigDefault("\"\"")
+        public String getFormat();
+
+        @ConfigInject
+        public BufferAllocator getBufferAllocator();
+    }
+
+    @Override
+    public void transaction(ConfigSource config, DecoderPlugin.Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+        control.run(task.dump());
+    }
+
+    @Override
+    public FileInput open(TaskSource taskSource, FileInput input)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+        return new CommonsCompressFileInput(
+                task.getBufferAllocator(),
+                new CommonsCompressProvider(task, new FileInputInputStream(input) {
+                    // NOTE: This is workaround code to avoid hanging issue.
+                    // This issue will be fixed after merging #112.
+                    // https://github.com/embulk/embulk/pull/112
+                    @Override
+                    public long skip(long len) {
+                        long skipped = super.skip(len);
+                        return skipped > 0 ? skipped : 0;
+                    }
+                }));
+    }
+}

data/src/main/java/org/embulk/decoder/CommonsCompressFileInput.java

@@ -0,0 +1,73 @@
+package org.embulk.decoder;
+
+import java.io.IOException;
+import java.io.InputStream;
+
+import org.embulk.spi.Buffer;
+import org.embulk.spi.BufferAllocator;
+import org.embulk.spi.FileInput;
+import org.embulk.spi.util.InputStreamFileInput.Provider;
+
+
+class CommonsCompressFileInput implements FileInput
+{
+    private final BufferAllocator allocator;
+    private final Provider provider;
+    private InputStream current;
+
+    public CommonsCompressFileInput(BufferAllocator allocator, Provider provider)
+    {
+        this.allocator = allocator;
+        this.provider = provider;
+        this.current = null;
+    }
+
+    @Override
+    public Buffer poll()
+    {
+        if (current == null) {
+            throw new IllegalStateException("nextFile() must be called before poll()");
+        }
+        Buffer buffer = allocator.allocate();
+        try {
+            int n = current.read(buffer.array(), buffer.offset(), buffer.capacity());
+            if (n < 0) {
+                return null;
+            }
+            buffer.limit(n);
+            Buffer b = buffer;
+            buffer = null;
+            return b;
+        } catch (IOException ex) {
+            throw new RuntimeException(ex);
+        } finally {
+            if (buffer != null) {
+                buffer.release();
+                buffer = null;
+            }
+        }
+    }
+
+    @Override
+    public boolean nextFile()
+    {
+        try {
+            // NOTE: DO NOT close current because this stream may
+            // be one of a file in an archive. Provider manage it.
+            current = provider.openNext();
+            return current != null;
+        } catch (IOException ex) {
+            throw new RuntimeException(ex);
+        }
+    }
+
+    @Override
+    public void close()
+    {
+        try {
+            provider.close();
+        } catch (IOException ex) {
+            throw new RuntimeException(ex);
+        }
+    }
+}

data/src/main/java/org/embulk/decoder/CommonsCompressProvider.java

@@ -0,0 +1,182 @@
+package org.embulk.decoder;
+
+import java.io.BufferedInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+
+import org.apache.commons.compress.archivers.ArchiveException;
+import org.apache.commons.compress.archivers.ArchiveInputStream;
+import org.apache.commons.compress.archivers.ArchiveStreamFactory;
+import org.apache.commons.compress.compressors.CompressorException;
+import org.apache.commons.compress.compressors.CompressorInputStream;
+import org.apache.commons.compress.compressors.CompressorStreamFactory;
+import org.embulk.decoder.CommonsCompressDecoderPlugin.PluginTask;
+import org.embulk.spi.util.FileInputInputStream;
+import org.embulk.spi.util.InputStreamFileInput.Provider;
+
+class CommonsCompressProvider implements Provider {
+    private static final String AUTO_DETECT_FORMAT = "";
+
+    private final FileInputInputStream files;
+    private final boolean formatAutoDetection;
+    private Iterator<InputStream> inputStreamIterator;
+    private String[] formats;
+
+    CommonsCompressProvider(PluginTask task, FileInputInputStream files) {
+        this.files = files;
+        this.formatAutoDetection = task == null
+                || CommonsCompressUtil.isAutoDetect(task.getFormat());
+        if (!this.formatAutoDetection) {
+            formats = CommonsCompressUtil.toFormats(task.getFormat());
+            if (formats == null) {
+                throw new RuntimeException("Failed to get a format.");
+            }
+        }
+    }
+
+    @Override
+    public InputStream openNext() throws IOException {
+        while (true) {
+            if (inputStreamIterator == null) {
+                if (!files.nextFile()) {
+                    return null;
+                }
+                inputStreamIterator = formatAutoDetection ? createInputStreamIterator(files)
+                        : createInputStreamIterator(formats, 0, files);
+            } else {
+                if (inputStreamIterator.hasNext()) {
+                    InputStream in = inputStreamIterator.next();
+                    if (in == null) {
+                        inputStreamIterator = null;
+                    } else {
+                        return in;
+                    }
+                } else {
+                    inputStreamIterator = null;
+                }
+            }
+        }
+    }
+
+    @Override
+    public void close() throws IOException {
+        inputStreamIterator = null;
+        if (files != null) {
+            files.close();
+        }
+    }
+
+    boolean isFormatAutoDetection() {
+        return formatAutoDetection;
+    }
+    
+    String[] getFormats() {
+        return formats;
+    }
+
+    Iterator<InputStream> createInputStreamIterator(InputStream in)
+            throws IOException {
+        // It is required to support mark to detect a file format.
+        in = in.markSupported() ? in : new BufferedInputStream(in);
+        try {
+            return new ArchiveInputStreamIterator(
+                    createArchiveInputStream(AUTO_DETECT_FORMAT, in));
+        } catch (IOException | ArchiveException e) {
+            // ArchiveStreamFactory set mark and reset the stream.
+            // So, we can use the same stream to check compressor.
+            try {
+                return toIterator(createCompressorInputStream(AUTO_DETECT_FORMAT, in));
+            } catch (CompressorException e2) {
+                throw new IOException("Failed to detect a file format.", e2);
+            }
+        }
+    }
+
+    /**
+     * Create iterator to list InputStream for each archived/compressed file.
+     * 
+     * This can handle like the following formats:
+     * 1 archived format which defined in ArchiveStreamFactory(e.g. tar)
+     * 1 archived format and 1 compressor format defined in CompressorStreamFactory.(e.g. tar.bz2)
+     * 1 compressor format defined in CompressorStreamFactory.(e.g. bz2)
+     * (Actually, compressor formats can use two or more times in this code.
+     *  But it is not common case.)
+     */
+    Iterator<InputStream> createInputStreamIterator(String[] formats,
+            int pos, InputStream in) throws IOException {
+        if (pos >= formats.length) {
+            return toIterator(in);
+        }
+
+        try {
+            String format = formats[pos];
+            if (CommonsCompressUtil.isArchiveFormat(format)) {
+                return new ArchiveInputStreamIterator(
+                        createArchiveInputStream(format, in));
+            } else if (CommonsCompressUtil.isCompressorFormat(format)) {
+                return createInputStreamIterator(formats, pos + 1,
+                        createCompressorInputStream(format, in));
+            }
+            throw new IOException("Unsupported format is configured. format:"
+                    + format);
+        } catch (ArchiveException | CompressorException e) {
+            throw new IOException(e);
+        }
+    }
+
+    /**
+     * Create a new ArchiveInputStream to read an archive file based on a format
+     * parameter.
+     * 
+     * If format is not set, this method tries to detect file format
+     * automatically. In this case, BufferedInputStream is used to wrap
+     * FileInputInputStream instance. BufferedInputStream may read a data
+     * partially when calling files.nextFile(). However, it doesn't matter
+     * because the partial read data should be discarded. And then this method
+     * is called again to create a new ArchiveInputStream.
+     * 
+     * @return a new ArchiveInputStream instance.
+     */
+    ArchiveInputStream createArchiveInputStream(String format, InputStream in)
+            throws IOException, ArchiveException {
+        ArchiveStreamFactory factory = new ArchiveStreamFactory();
+        if (CommonsCompressUtil.isAutoDetect(format)) {
+            in = in.markSupported() ? in : new BufferedInputStream(in);
+            try {
+                return factory.createArchiveInputStream(in);
+            } catch (ArchiveException e) {
+                throw new IOException(
+                        "Failed to detect a file format. Please try to set a format explicitly.",
+                        e);
+            }
+        } else {
+            return factory.createArchiveInputStream(format, in);
+        }
+    }
+
+    CompressorInputStream createCompressorInputStream(String format,
+            InputStream in) throws IOException, CompressorException {
+        CompressorStreamFactory factory = new CompressorStreamFactory();
+        if (CommonsCompressUtil.isAutoDetect(format)) {
+            in = in.markSupported() ? in : new BufferedInputStream(in);
+            try {
+                return factory.createCompressorInputStream(in);
+            } catch (CompressorException e) {
+                throw new IOException(
+                        "Failed to detect a file format. Please try to set a format explicitly.",
+                        e);
+            }
+        } else {
+            return factory.createCompressorInputStream(format, in);
+        }
+    }
+
+    private Iterator<InputStream> toIterator(InputStream in) {
+        List<InputStream> list = new ArrayList<InputStream>(1);
+        list.add(in);
+        return list.iterator();
+    }
+}

data/src/main/java/org/embulk/decoder/CommonsCompressUtil.java

@@ -0,0 +1,135 @@
+package org.embulk.decoder;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import org.apache.commons.compress.archivers.ArchiveStreamFactory;
+import org.apache.commons.compress.compressors.CompressorStreamFactory;
+
+class CommonsCompressUtil {
+    // TODO: It may be better to check performance between Set and array.
+    // NOTE: Some file types may not work in an environment because some required
+    // libraries are not found.
+    static final String[] archiveFormats = {
+        ArchiveStreamFactory.AR,
+        ArchiveStreamFactory.ARJ,
+        ArchiveStreamFactory.CPIO,
+        ArchiveStreamFactory.DUMP,
+        ArchiveStreamFactory.JAR,
+        ArchiveStreamFactory.SEVEN_Z,
+        ArchiveStreamFactory.TAR,
+        ArchiveStreamFactory.ZIP,
+    };
+
+    static final String[] compressorFormats = {
+        CompressorStreamFactory.BZIP2,
+        CompressorStreamFactory.DEFLATE,
+        CompressorStreamFactory.GZIP,
+        CompressorStreamFactory.LZMA,
+        CompressorStreamFactory.PACK200,
+        CompressorStreamFactory.SNAPPY_FRAMED,
+        CompressorStreamFactory.SNAPPY_RAW,
+        CompressorStreamFactory.XZ,
+        CompressorStreamFactory.Z,
+    };
+
+    // This table is even indexes have short extensions and odd indexes has
+    // split formats for each short extensions.
+    private static final String[] solidCompressionFormats = {
+        "tgz",      ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.GZIP,
+        "tar.gz",   ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.GZIP,
+        "tbz",      ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.BZIP2,
+        "tbz2",     ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.BZIP2,
+        "tb2",      ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.BZIP2,
+        "tar.bz2",  ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.BZIP2,
+        "taz",      ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.Z,
+        "tz",       ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.Z,
+        "tar.Z",    ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.Z,
+        "tlz",      ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.LZMA,
+        "tar.lz",   ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.LZMA,
+        "tar.lzma", ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.LZMA,
+        "txz",      ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.XZ,
+        "tar.xz",   ArchiveStreamFactory.TAR + " " + CompressorStreamFactory.XZ
+    };
+
+    static boolean isArchiveFormat(String format) {
+        for (String fmt : archiveFormats) {
+            if (fmt.equalsIgnoreCase(format)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    static boolean isCompressorFormat(String format) {
+        for (String fmt : compressorFormats) {
+            if (fmt.equalsIgnoreCase(format)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    static boolean isAutoDetect(String format) {
+        return format == null || format.length() == 0;
+    }
+
+    /**
+     * Split solid compresson formats and reorder to decode the formats
+     * based on this order.
+     * 
+     * If format is a single format like "tar", then return
+     * new String[]{"tar"}.
+     * If format is a solid compresson format like "tgz", then return
+     * new String[]{"gzip", "tar"}.
+     * If format is "tar bzip2", then return
+     * new String[]{"bzip2", "tar"}.
+     * 
+     * @param format contains a file format or some file formats.
+     * @return a single format or multi format values.
+     * Otherwise, returns null.
+     */
+    static String[] toFormats(String format) {
+        if (isAutoDetect(format)) {
+            return null;
+        } else if (isArchiveFormat(format) || isCompressorFormat(format)) {
+            return splitAndReverse(format);
+        }
+
+        String[] formats = toSolidCompressionFormats(format);
+        if (formats != null) {
+            return formats;
+        }
+
+        formats = splitAndReverse(format);
+        
+        for (String s : formats) {
+            if (!(isArchiveFormat(s) || isCompressorFormat(s))) {
+                return null;
+            }
+        }
+        
+        return formats;
+    }
+
+    private static String[] toSolidCompressionFormats(String format) {
+        for (int i = 0;i < solidCompressionFormats.length; i+= 2) {
+            if (solidCompressionFormats[i].equalsIgnoreCase(format)) {
+                return splitAndReverse(solidCompressionFormats[i + 1]);
+            }
+        }
+        return null;
+    }
+
+    private static String[] splitAndReverse(String format) {
+        List<String> result = new ArrayList<>();
+        for (String s : format.split(" ")) {
+            if (s.length() > 0) {
+                result.add(s);
+            }
+        }
+        Collections.reverse(result);
+        return result.toArray(new String[result.size()]);
+    }
+}