embulk 0.1.0

data/embulk-standards/src/main/java/org/embulk/standards/LocalFileInputPlugin.java

@@ -0,0 +1,138 @@
+package org.embulk.standards;
+
+import java.util.List;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.InputStream;
+import java.io.IOException;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.Files;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.FileVisitResult;
+import java.nio.file.attribute.BasicFileAttributes;
+import javax.validation.constraints.NotNull;
+import com.google.common.collect.ImmutableList;
+import com.fasterxml.jackson.annotation.JacksonInject;
+import org.embulk.config.Config;
+import org.embulk.config.Task;
+import org.embulk.config.TaskSource;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.NextConfig;
+import org.embulk.config.CommitReport;
+import org.embulk.spi.BufferAllocator;
+import org.embulk.spi.Exec;
+import org.embulk.spi.FileInputPlugin;
+import org.embulk.spi.TransactionalFileInput;
+import org.embulk.spi.util.InputStreamFileInput;
+import org.slf4j.Logger;
+
+import static org.embulk.spi.util.Inputs.formatPath;
+
+public class LocalFileInputPlugin
+        implements FileInputPlugin
+{
+    public interface PluginTask
+            extends Task
+    {
+        @Config("paths")
+        @NotNull
+        public List<String> getPathPrefixes();
+
+        public List<String> getFiles();
+        public void setFiles(List<String> files);
+
+        @JacksonInject
+        public BufferAllocator getBufferAllocator();
+    }
+
+    private final Logger log = Exec.getLogger(getClass());
+
+    @Override
+    public NextConfig transaction(ConfigSource config, FileInputPlugin.Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+
+        // list files recursively
+        task.setFiles(listFiles(task));
+
+        // run with threads. number of processors is same with number of files
+        control.run(task.dump(), task.getFiles().size());
+
+        return Exec.newNextConfig();
+    }
+
+    public List<String> listFiles(PluginTask task)
+    {
+        final ImmutableList.Builder<String> builder = ImmutableList.builder();
+        for (String prefix : task.getPathPrefixes()) {
+            String formatted = formatPath(prefix);
+            try {
+                log.info("Listing local files with prefix '{}'", formatted);
+                Files.walkFileTree(Paths.get(formatted), new SimpleFileVisitor<Path>() {
+                    @Override
+                    public FileVisitResult visitFile(Path file, BasicFileAttributes aAttrs)
+                    {
+                        builder.add(file.toString());
+                        return FileVisitResult.CONTINUE;
+                    }
+                });
+            } catch (IOException ex) {
+                throw new RuntimeException(String.format("Failed get a list of local files at '%s'", formatted), ex);
+            }
+        }
+        return builder.build();
+    }
+
+    @Override
+    public TransactionalFileInput open(TaskSource taskSource, int processorIndex)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+        return new LocalFileInput(task, processorIndex);
+    }
+
+    public static class LocalFileInput
+            extends InputStreamFileInput
+            implements TransactionalFileInput
+    {
+        // TODO create single-file InputStreamFileInput utility
+        private static class SingleFileProvider
+                implements InputStreamFileInput.Provider
+        {
+            private final File file;
+            private boolean opened = false;
+
+            public SingleFileProvider(File file)
+            {
+                this.file = file;
+            }
+
+            @Override
+            public InputStream openNext() throws IOException
+            {
+                if (opened) {
+                    return null;
+                }
+                opened = true;
+                return new FileInputStream(file);
+            }
+
+            @Override
+            public void close() { }
+        }
+
+        public LocalFileInput(PluginTask task, int processorIndex)
+        {
+            super(task.getBufferAllocator(), new SingleFileProvider(new File(task.getFiles().get(processorIndex))));
+        }
+
+        @Override
+        public void abort() { }
+
+        @Override
+        public CommitReport commit()
+        {
+            return Exec.newCommitReport();
+        }
+    }
+}

data/embulk-standards/src/main/java/org/embulk/standards/LocalFileOutputPlugin.java

@@ -0,0 +1,128 @@
+package org.embulk.standards;
+
+import java.io.File;
+import java.io.FileNotFoundException;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.List;
+import org.embulk.config.Config;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.NextConfig;
+import org.embulk.config.CommitReport;
+import org.embulk.config.Task;
+import org.embulk.config.TaskSource;
+import org.embulk.spi.Buffer;
+import org.embulk.spi.FileOutputPlugin;
+import org.embulk.spi.TransactionalFileOutput;
+import org.embulk.spi.Exec;
+import org.slf4j.Logger;
+
+public class LocalFileOutputPlugin
+        implements FileOutputPlugin
+{
+    public interface PluginTask
+            extends Task
+    {
+        @Config("directory")
+        public String getDirectory();
+
+        @Config("file_name")
+        public String getFileNameFormat();
+
+        @Config("file_ext")
+        public String getFileNameExtension();
+
+        // TODO support in FileInputPlugin and FileOutputPlugin
+        //@Config("compress_type")
+        //public String getCompressType();
+    }
+
+    private final Logger log = Exec.getLogger(getClass());
+
+    @Override
+    public NextConfig transaction(ConfigSource config, int processorCount,
+            FileOutputPlugin.Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+
+        control.run(task.dump());
+
+        return Exec.newNextConfig();
+    }
+
+    @Override
+    public TransactionalFileOutput open(TaskSource taskSource, final int processorIndex)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+
+        // TODO format path using timestamp
+        final String fileName = task.getFileNameFormat();
+
+        final String pathPrefix = task.getDirectory() + File.separator + fileName;
+        final String pathSuffix = task.getFileNameExtension();
+
+        final List<String> fileNames = new ArrayList<>();
+
+        return new TransactionalFileOutput() {
+            private int fileIndex = 0;
+            private FileOutputStream output = null;
+
+            public void nextFile()
+            {
+                closeFile();
+                String path = pathPrefix + String.format(".%03d.%02d.", processorIndex, fileIndex) + pathSuffix;
+                log.info("Writing local file '{}'", path);
+                fileNames.add(path);
+                try {
+                    output = new FileOutputStream(new File(path));
+                } catch (FileNotFoundException ex) {
+                    throw new RuntimeException(ex);  // TODO exception class
+                }
+                fileIndex++;
+            }
+
+            private void closeFile()
+            {
+                if (output != null) {
+                    try {
+                        output.close();
+                    } catch (IOException ex) {
+                        throw new RuntimeException(ex);
+                    }
+                }
+            }
+
+            public void add(Buffer buffer)
+            {
+                try {
+                    output.write(buffer.array(), buffer.offset(), buffer.limit());
+                } catch (IOException ex) {
+                    throw new RuntimeException(ex);
+                }
+            }
+
+            public void finish()
+            {
+                closeFile();
+            }
+
+            public void close()
+            {
+                closeFile();
+            }
+
+            public void abort() { }
+
+            public CommitReport commit()
+            {
+                CommitReport report = Exec.newCommitReport();
+                // TODO better setting for Report
+                // report.set("file_names", fileNames);
+                // report.set("file_sizes", fileSizes);
+                return report;
+            }
+        };
+    }
+}

data/embulk-standards/src/main/java/org/embulk/standards/NullOutputPlugin.java

@@ -0,0 +1,46 @@
+package org.embulk.standards;
+
+import org.embulk.config.ConfigSource;
+import org.embulk.config.TaskSource;
+import org.embulk.config.NextConfig;
+import org.embulk.config.CommitReport;
+import org.embulk.spi.Schema;
+import org.embulk.spi.Page;
+import org.embulk.spi.Exec;
+import org.embulk.spi.OutputPlugin;
+import org.embulk.spi.TransactionalPageOutput;
+
+public class NullOutputPlugin
+        implements OutputPlugin
+{
+    @Override
+    public NextConfig transaction(ConfigSource config,
+            Schema schema, int processorCount,
+            OutputPlugin.Control control)
+    {
+        control.run(Exec.newTaskSource());
+        return Exec.newNextConfig();
+    }
+
+    @Override
+    public TransactionalPageOutput open(TaskSource taskSource, Schema schema, int processorIndex)
+    {
+        return new TransactionalPageOutput() {
+            public void add(Page page)
+            {
+                page.release();
+            }
+
+            public void finish() { }
+
+            public void close() { }
+
+            public void abort() { }
+
+            public CommitReport commit()
+            {
+                return Exec.newCommitReport();
+            }
+        };
+    }
+}

data/embulk-standards/src/main/java/org/embulk/standards/S3FileInputPlugin.java

@@ -0,0 +1,238 @@
+package org.embulk.standards;
+
+import java.util.List;
+import java.io.IOException;
+import java.io.InputStream;
+
+import com.amazonaws.AmazonClientException;
+import com.amazonaws.AmazonServiceException;
+import com.google.common.collect.ImmutableList;
+import com.google.common.base.Optional;
+import com.fasterxml.jackson.annotation.JacksonInject;
+import com.amazonaws.auth.AWSCredentials;
+import com.amazonaws.auth.AWSCredentialsProvider;
+import com.amazonaws.auth.BasicAWSCredentials;
+import com.amazonaws.services.s3.AmazonS3Client;
+import com.amazonaws.services.s3.model.ListObjectsRequest;
+import com.amazonaws.services.s3.model.S3ObjectSummary;
+import com.amazonaws.services.s3.model.ObjectListing;
+import com.amazonaws.services.s3.model.GetObjectRequest;
+import com.amazonaws.services.s3.model.S3Object;
+import com.amazonaws.ClientConfiguration;
+import org.embulk.config.Config;
+import org.embulk.config.Task;
+import org.embulk.config.TaskSource;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.NextConfig;
+import org.embulk.config.CommitReport;
+import org.embulk.spi.BufferAllocator;
+import org.embulk.spi.Exec;
+import org.embulk.spi.FileInputPlugin;
+import org.embulk.spi.TransactionalFileInput;
+import org.embulk.spi.util.InputStreamFileInput;
+import org.slf4j.Logger;
+
+import static org.embulk.spi.util.Inputs.formatPath;
+
+public class S3FileInputPlugin
+        implements FileInputPlugin
+{
+    public interface PluginTask
+            extends Task
+    {
+        @Config("bucket")
+        public String getBucket();
+
+        @Config("paths")
+        public List<String> getPathPrefixes();
+
+        @Config("endpoint")
+        public Optional<String> getEndpoint();
+
+        // TODO timeout, ssl, etc
+
+        @Config("access_key_id")
+        public String getAccessKeyId();
+
+        @Config("secret_access_key")
+        public String getSecretAccessKey();
+
+        // TODO support more options such as STS
+
+        public List<String> getFiles();
+        public void setFiles(List<String> files);
+
+        @JacksonInject
+        public BufferAllocator getBufferAllocator();
+    }
+
+    private final Logger log = Exec.getLogger(getClass());
+
+    @Override
+    public NextConfig transaction(ConfigSource config, FileInputPlugin.Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+
+        // list files recursively
+        task.setFiles(listFiles(task));
+
+        // number of processors is same with number of files
+
+        // run
+        control.run(task.dump(), task.getFiles().size());
+
+        return Exec.newNextConfig();
+    }
+
+    public static AWSCredentialsProvider getCredentialsProvider(PluginTask task)
+    {
+        final AWSCredentials cred = new BasicAWSCredentials(
+                task.getAccessKeyId(), task.getSecretAccessKey());
+        return new AWSCredentialsProvider() {
+            public AWSCredentials getCredentials()
+            {
+                return cred;
+            }
+
+            public void refresh()
+            {
+            }
+        };
+    }
+
+    private static AmazonS3Client newS3Client(PluginTask task)
+    {
+        AWSCredentialsProvider credentials = getCredentialsProvider(task);
+        AmazonS3Client client = newS3Client(credentials, task.getEndpoint());
+        return client;
+    }
+
+    private static AmazonS3Client newS3Client(AWSCredentialsProvider credentials,
+            Optional<String> endpoint)
+    {
+        // TODO get config from AmazonS3Task
+        ClientConfiguration clientConfig = new ClientConfiguration();
+        //clientConfig.setProtocol(Protocol.HTTP);
+        clientConfig.setMaxConnections(50); // SDK default: 50
+        clientConfig.setMaxErrorRetry(3); // SDK default: 3
+        clientConfig.setSocketTimeout(8*60*1000); // SDK default: 50*1000
+
+        AmazonS3Client client = new AmazonS3Client(credentials, clientConfig);
+
+        if (endpoint.isPresent()) {
+            client.setEndpoint(endpoint.get());
+        }
+
+        return client;
+    }
+
+    public List<String> listFiles(PluginTask task)
+    {
+        AmazonS3Client client = newS3Client(task);
+        String bucketName = task.getBucket();
+
+        ImmutableList.Builder<String> builder = ImmutableList.builder();
+        for (String prefix : task.getPathPrefixes()) {
+            String formatted = formatPath(prefix);
+            try {
+                log.info("Listing S3 files with prefix '{}'", formatted);
+                builder.addAll(listS3FilesByPrefix(client, bucketName, formatted));
+            } catch (RuntimeException e) {
+                throw new RuntimeException(String.format("Failed get a list of S3 files at '%s'", formatted), e);
+            }
+        }
+
+        return builder.build();
+    }
+
+    /**
+     * Lists S3 filenames filtered by prefix.
+     *
+     * The resulting list does not include the file that's size == 0.
+     */
+    public static List<String> listS3FilesByPrefix(AmazonS3Client client, String bucketName, String prefix)
+    {
+        // TODO implement retrying
+        ImmutableList.Builder<String> builder = ImmutableList.builder();
+
+        String lastKey = null;
+        do {
+            ListObjectsRequest req = new ListObjectsRequest(bucketName, prefix, lastKey, null, 1024);
+            ObjectListing ol = client.listObjects(req);
+            for(S3ObjectSummary s : ol.getObjectSummaries()) {
+                if (s.getSize() > 0) {
+                    builder.add(s.getKey());
+                }
+            }
+            lastKey = ol.getNextMarker();
+        } while(lastKey != null);
+
+        return builder.build();
+    }
+
+    @Override
+    public TransactionalFileInput open(TaskSource taskSource, int processorIndex)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+        return new S3FileInput(task, processorIndex);
+    }
+
+    public static class S3FileInput
+            extends InputStreamFileInput
+            implements TransactionalFileInput
+    {
+        // TODO create single-file InputStreamFileInput utility
+        private static class SingleFileProvider
+                implements InputStreamFileInput.Provider
+        {
+            private AmazonS3Client client;
+            private final String bucket;
+            private final String key;
+            private boolean opened = false;
+
+            public SingleFileProvider(PluginTask task, int processorIndex)
+            {
+                this.client = newS3Client(task);
+                this.bucket = task.getBucket();
+                this.key = task.getFiles().get(processorIndex);
+            }
+
+            @Override
+            public InputStream openNext() throws IOException
+            {
+                if (opened) {
+                    return null;
+                }
+                opened = true;
+                GetObjectRequest request = new GetObjectRequest(bucket, key);
+                //if (pos > 0) {
+                //    request.setRange(pos, contentLength);
+                //}
+                S3Object obj = client.getObject(request);
+                //if (pos <= 0) {
+                //    // first call
+                //    contentLength = obj.getObjectMetadata().getContentLength();
+                //}
+                return obj.getObjectContent();
+            }
+
+            @Override
+            public void close() { }
+        }
+
+        public S3FileInput(PluginTask task, int processorIndex)
+        {
+            super(task.getBufferAllocator(), new SingleFileProvider(task, processorIndex));
+        }
+
+        public void abort() { }
+
+        public CommitReport commit()
+        {
+            return Exec.newCommitReport();
+        }
+
+        @Override
+        public void close() { }
+    }
+}