package org.embulk.spi.util; import java.util.Arrays; import java.util.Iterator; import java.io.InputStream; import java.io.Closeable; import java.io.IOException; import org.embulk.spi.Buffer; import org.embulk.spi.FileInput; import org.embulk.spi.BufferAllocator; public class InputStreamFileInput implements FileInput { public interface Provider extends Closeable { public InputStream openNext() throws IOException; public void close() throws IOException; } public interface Opener { public InputStream open() throws IOException; } public static class IteratorProvider implements Provider { private Iterator iterator; public IteratorProvider(Iterable iterable) { this.iterator = iterable.iterator(); } public IteratorProvider(Iterator iterator) { this.iterator = iterator; } @Override public InputStream openNext() throws IOException { if (!iterator.hasNext()) { return null; } return iterator.next(); } @Override public void close() throws IOException { while (iterator.hasNext()) { iterator.next().close(); } } } private static class OpenerProvider implements Provider { private Opener opener; public OpenerProvider(Opener opener) { this.opener = opener; } @Override public InputStream openNext() throws IOException { if (opener == null) { return null; } InputStream stream = opener.open(); opener = null; return stream; } @Override public void close() throws IOException { } } private static class InputStreamProvider implements Provider { private InputStream input; public InputStreamProvider(InputStream input) { this.input = input; } @Override public InputStream openNext() throws IOException { if (input == null) { return null; } InputStream ret = input; input = null; return ret; } @Override public void close() throws IOException { if (input != null) { input.close(); input = null; } } } private final BufferAllocator allocator; private final Provider provider; private InputStream current; public InputStreamFileInput(BufferAllocator allocator, Provider provider) { this.allocator = allocator; this.provider = provider; this.current = null; } public InputStreamFileInput(BufferAllocator allocator, Opener opener) { this(allocator, new OpenerProvider(opener)); } public InputStreamFileInput(BufferAllocator allocator, InputStream openedStream) { this(allocator, new InputStreamProvider(openedStream)); } 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; } } } public boolean nextFile() { try { if (current != null) { current.close(); current = null; } current = provider.openNext(); return current != null; } catch (IOException ex) { throw new RuntimeException(ex); } } public void close() { try { try { if (current != null) { current.close(); current = null; } } finally { provider.close(); } } catch (IOException ex) { throw new RuntimeException(ex); } } }