embulk 0.1.0

data/embulk-core/src/test/java/org/embulk/spi/MockParserPlugin.java

@@ -0,0 +1,73 @@
+package org.embulk.spi;
+
+import org.embulk.config.Config;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.Task;
+import org.embulk.config.TaskSource;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.type.Type;
+import org.embulk.spi.Column;
+import org.embulk.spi.Schema;
+import org.embulk.spi.SchemaConfig;
+import org.embulk.spi.PageOutput;
+
+public class MockParserPlugin implements ParserPlugin
+{
+    public static boolean raiseException = false;
+
+    public interface PluginTask extends Task
+    {
+        @Config("columns")
+        public SchemaConfig getSchemaConfig();
+    }
+
+    @Override
+    public void transaction(ConfigSource config, Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+        control.run(task.dump(), task.getSchemaConfig().toSchema());
+    }
+
+    @Override
+    public void run(TaskSource taskSource, Schema schema,
+            FileInput input, PageOutput output)
+    {
+        try (final PageBuilder pageBuilder = new PageBuilder(
+                Exec.getBufferAllocator(), schema, output)) {
+            while (input.nextFile()) {
+                Buffer buffer = input.poll();
+                if (buffer != null) {
+                    for (Column column : schema.getColumns()) {
+                        Type type = column.getType();
+                        switch (type.getName()) {
+                        case "boolean":
+                            pageBuilder.setBoolean(column, true);
+                            break;
+                        case "long":
+                            pageBuilder.setLong(column, 2L);
+                            break;
+                        case "double":
+                            pageBuilder.setDouble(column, 3.0D);
+                            break;
+                        case "string":
+                            pageBuilder.setString(column, "45");
+                            break;
+                        case "timestamp":
+                            pageBuilder.setTimestamp(column,
+                                    Timestamp.ofEpochMilli(678L));
+                            break;
+                        default:
+                            throw new IllegalStateException("Unknown type: "
+                                    + type.getName());
+                        }
+                    }
+                    pageBuilder.addRecord();
+                    if (raiseException) {
+                        throw new RuntimeException("emulated exception");
+                    }
+                }
+            }
+            pageBuilder.finish();
+        }
+    }
+}

data/embulk-core/src/test/java/org/embulk/spi/PageTestUtils.java

@@ -0,0 +1,78 @@
+package org.embulk.spi;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.embulk.spi.TestPageBuilderReader.MockPageOutput;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.Column;
+import org.embulk.spi.ColumnConfig;
+import org.embulk.spi.Schema;
+import org.embulk.spi.SchemaConfig;
+import org.embulk.spi.SchemaVisitor;
+import org.embulk.spi.type.Type;
+
+import com.google.common.collect.ImmutableList;
+
+public class PageTestUtils
+{
+    private PageTestUtils()
+    {
+    }
+
+    public static Schema newSchema()
+    {
+        return new Schema(ImmutableList.<Column> of());
+    }
+
+    public static Schema newSchema(Column... columns)
+    {
+        return new Schema(ImmutableList.copyOf(columns));
+    }
+
+    public static Schema newSchema(ColumnConfig... columns)
+    {
+        return new SchemaConfig(ImmutableList.copyOf(columns)).toSchema();
+    }
+
+    public static ColumnConfig newColumn(String name, Type type)
+    {
+        return new ColumnConfig(name, type, null);
+    }
+
+    public static List<Page> buildPage(BufferAllocator bufferAllocator,
+            Schema schema, Object... values)
+    {
+        MockPageOutput output = new MockPageOutput();
+        try (PageBuilder builder = new PageBuilder(bufferAllocator, schema,
+                output)) {
+            int idx = 0;
+            while (idx < values.length) {
+                for (int column = 0; column < builder.getSchema()
+                        .getColumnCount(); ++column) {
+                    Object value = values[idx++];
+                    if (value == null) {
+                        builder.setNull(column);
+                    } else if (value instanceof Boolean) {
+                        builder.setBoolean(column, (Boolean) value);
+                    } else if (value instanceof Double) {
+                        builder.setDouble(column, (Double) value);
+                    } else if (value instanceof Long) {
+                        builder.setLong(column, (Long) value);
+                    } else if (value instanceof String) {
+                        builder.setString(column, (String) value);
+                    } else if (value instanceof Timestamp) {
+                        builder.setTimestamp(column, (Timestamp) value);
+                    } else {
+                        throw new IllegalStateException(
+                                "Unsupported type in test utils: "
+                                        + value.toString());
+                    }
+                }
+                builder.addRecord();
+            }
+            builder.finish();
+        }
+        return output.pages;
+    }
+}

data/embulk-core/src/test/java/org/embulk/spi/TestFileInputInputStream.java

@@ -0,0 +1,67 @@
+package org.embulk.spi;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import java.io.ByteArrayOutputStream;
+import java.util.Random;
+import org.junit.Rule;
+import org.junit.Test;
+import org.embulk.spi.util.ListFileInput;
+import org.embulk.spi.util.FileInputInputStream;
+import org.embulk.spi.util.FileOutputOutputStream;
+import org.embulk.EmbulkTestRuntime;
+
+public class TestFileInputInputStream
+{
+    @Rule
+    public EmbulkTestRuntime runtime = new EmbulkTestRuntime();
+
+    private ListFileInput fileInput;
+    private MockFileOutput fileOutput;
+
+    private FileInputInputStream in;
+    private FileOutputOutputStream out;
+
+    private void newOutputStream()
+    {
+        fileOutput = new MockFileOutput();
+        out = new FileOutputOutputStream(fileOutput, runtime.getBufferAllocator());
+    }
+
+    private void newInputStream()
+    {
+        fileInput = new ListFileInput(fileOutput.getFiles());
+        in = new FileInputInputStream(fileInput);
+    }
+
+    @Test
+    public void testRandomReadWrite() throws Exception
+    {
+        newOutputStream();
+        out.nextFile();
+        ByteArrayOutputStream expectedOut = new ByteArrayOutputStream();
+        Random rand = runtime.getRandom();
+        byte[] buffer = new byte[rand.nextInt() % 1024 + 1024];
+        for (int i = 0; i < 256; i++) {
+            rand.nextBytes(buffer);
+            expectedOut.write(buffer);
+            out.write(buffer);
+        }
+        out.finish();
+        byte[] expected = expectedOut.toByteArray();
+        byte[] actual = new byte[expected.length];
+
+        newInputStream();
+        in.nextFile();
+        int pos = 0;
+        while (pos < actual.length) {
+            int n = in.read(actual, pos, actual.length - pos);
+            if (n < 0) {
+                break;
+            }
+            pos += n;
+        }
+        assertEquals(expected.length, pos);
+        assertArrayEquals(expected, actual);
+    }
+}

data/embulk-core/src/test/java/org/embulk/spi/TestFileInputRunner.java

@@ -0,0 +1,180 @@
+package org.embulk.spi;
+
+import static org.junit.Assert.assertEquals;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Queue;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import org.embulk.EmbulkTestRuntime;
+import org.embulk.config.CommitReport;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.NextConfig;
+import org.embulk.config.TaskSource;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.TestPageBuilderReader.MockPageOutput;
+import org.embulk.spi.Schema;
+import org.embulk.spi.util.Pages;
+
+public class TestFileInputRunner
+{
+    @Rule
+    public EmbulkTestRuntime runtime = new EmbulkTestRuntime();
+
+    @Before
+    public void tearDown()
+    {
+        MockParserPlugin.raiseException = false;
+    }
+
+    private static class MockFileInputPlugin implements FileInputPlugin
+    {
+        Boolean transactionCompleted = null;
+        Queue<Buffer> buffers;
+
+        public MockFileInputPlugin(Queue<Buffer> buffers)
+        {
+            this.buffers = buffers;
+        }
+
+        public NextConfig transaction(ConfigSource config,
+                FileInputPlugin.Control control)
+        {
+            control.run(Exec.newTaskSource(), 1);
+            return null;
+        }
+
+        public TransactionalFileInput open(TaskSource taskSource,
+                int processorIndex)
+        {
+            return new TransactionalFileInput()
+            {
+                @Override
+                public Buffer poll()
+                {
+                    return buffers.poll();
+                }
+
+                @Override
+                public boolean nextFile()
+                {
+                    return !buffers.isEmpty();
+                }
+
+                @Override
+                public void close()
+                {
+                }
+
+                @Override
+                public void abort()
+                {
+                    transactionCompleted = false;
+                }
+
+                @Override
+                public CommitReport commit()
+                {
+                    transactionCompleted = true;
+                    return null;
+                }
+            };
+        }
+    }
+
+    @Test
+    public void testMockParserIteration()
+    {
+        Buffer[] buffers = new Buffer[] {
+                runtime.getBufferAllocator().allocate(),
+                runtime.getBufferAllocator().allocate() };
+        MockFileInputPlugin fileInputPlugin = new MockFileInputPlugin(
+                new LinkedList<Buffer>(Arrays.asList(buffers)));
+        final FileInputRunner runner = new FileInputRunner(fileInputPlugin);
+
+        ConfigSource config = Exec.newConfigSource().set(
+                "parser",
+                ImmutableMap.of("type", "mock", "columns", ImmutableList.of(
+                        ImmutableMap.of("name", "col1", "type", "boolean"),
+                        ImmutableMap.of("name", "col2", "type", "long"),
+                        ImmutableMap.of("name", "col3", "type", "double"),
+                        ImmutableMap.of("name", "col4", "type", "string"),
+                        ImmutableMap.of("name", "col5", "type", "timestamp"))));
+
+        final MockPageOutput output = new MockPageOutput();
+        runner.transaction(config, new InputPlugin.Control()
+        {
+            public List<CommitReport> run(TaskSource inputTaskSource,
+                    Schema schema, int processorCount)
+            {
+                List<CommitReport> reports = new ArrayList<>();
+                reports.add(runner.run(inputTaskSource, schema, 0, output));
+                return reports;
+            }
+        });
+
+        assertEquals(true, fileInputPlugin.transactionCompleted);
+        assertEquals(1, output.pages.size());
+
+        Schema schema = config.getNested("parser")
+                .loadConfig(MockParserPlugin.PluginTask.class)
+                .getSchemaConfig().toSchema();
+
+        List<Object[]> records = Pages.toObjects(schema, output.pages);
+        assertEquals(2, records.size());
+        for (Object[] record : records) {
+            assertEquals(5, record.length);
+            assertEquals(true, record[0]);
+            assertEquals(2L, record[1]);
+            assertEquals(3.0D, (Double) record[2], 0.01D);
+            assertEquals("45", record[3]);
+            assertEquals(678L, ((Timestamp) record[4]).toEpochMilli());
+        }
+    }
+
+    @Test
+    public void testTransactionAborted()
+    {
+        Buffer[] buffers = new Buffer[] {
+                runtime.getBufferAllocator().allocate(),
+                runtime.getBufferAllocator().allocate() };
+        MockFileInputPlugin fileInputPlugin = new MockFileInputPlugin(
+                new LinkedList<Buffer>(Arrays.asList(buffers)));
+        final FileInputRunner runner = new FileInputRunner(fileInputPlugin);
+
+        ConfigSource config = Exec.newConfigSource().set(
+                "parser",
+                ImmutableMap.of("type", "mock", "columns", ImmutableList.of(
+                        ImmutableMap.of("name", "col1", "type", "boolean"),
+                        ImmutableMap.of("name", "col2", "type", "long"),
+                        ImmutableMap.of("name", "col3", "type", "double"),
+                        ImmutableMap.of("name", "col4", "type", "string"),
+                        ImmutableMap.of("name", "col5", "type", "timestamp"))));
+
+        final MockPageOutput output = new MockPageOutput();
+
+        MockParserPlugin.raiseException = true;
+
+        try {
+            runner.transaction(config, new InputPlugin.Control()
+            {
+                public List<CommitReport> run(TaskSource inputTaskSource,
+                        Schema schema, int processorCount)
+                {
+                    List<CommitReport> reports = new ArrayList<>();
+                    reports.add(runner.run(inputTaskSource, schema, 0, output));
+                    return reports;
+                }
+            });
+        } catch (RuntimeException re) {
+        }
+        assertEquals(false, fileInputPlugin.transactionCompleted);
+        assertEquals(0, output.pages.size());
+    }
+}

data/embulk-core/src/test/java/org/embulk/spi/TestFileOutputRunner.java

@@ -0,0 +1,192 @@
+package org.embulk.spi;
+
+import static org.junit.Assert.assertEquals;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.embulk.EmbulkTestRuntime;
+import org.embulk.config.CommitReport;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.NextConfig;
+import org.embulk.config.Task;
+import org.embulk.config.TaskSource;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.Schema;
+import org.junit.Rule;
+import org.junit.Test;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+
+public class TestFileOutputRunner
+{
+    @Rule
+    public EmbulkTestRuntime runtime = new EmbulkTestRuntime();
+
+    public interface PluginTask extends Task
+    {
+    }
+
+    private static class MockFileOutputPlugin implements FileOutputPlugin
+    {
+        Boolean transactionCompleted = null;
+
+        @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)
+        {
+            return new TransactionalFileOutput()
+            {
+
+                @Override
+                public void nextFile()
+                {
+                }
+
+                @Override
+                public void add(Buffer buffer)
+                {
+                }
+
+                @Override
+                public void finish()
+                {
+                }
+
+                @Override
+                public void close()
+                {
+                }
+
+                @Override
+                public void abort()
+                {
+                    transactionCompleted = false;
+                }
+
+                @Override
+                public CommitReport commit()
+                {
+                    transactionCompleted = true;
+                    return Exec.newCommitReport();
+                }
+            };
+        }
+    }
+
+    @Test
+    public void testMockFormatterIteration()
+    {
+        MockFileOutputPlugin fileOutputPlugin = new MockFileOutputPlugin();
+        final FileOutputRunner runner = new FileOutputRunner(fileOutputPlugin);
+
+        ImmutableList<ImmutableMap<String, String>> columns = ImmutableList.of(
+                ImmutableMap.of("name", "col1", "type", "boolean"),
+                ImmutableMap.of("name", "col2", "type", "long"),
+                ImmutableMap.of("name", "col3", "type", "double"),
+                ImmutableMap.of("name", "col4", "type", "string"),
+                ImmutableMap.of("name", "col5", "type", "timestamp"));
+        ConfigSource config = Exec
+                .newConfigSource()
+                .set("type", "unused?")
+                .set("formatter",
+                        ImmutableMap.of("type", "mock", "columns", columns));
+        final Schema schema = config.getNested("formatter")
+                .loadConfig(MockParserPlugin.PluginTask.class)
+                .getSchemaConfig().toSchema();
+
+        runner.transaction(config, schema, 1, new OutputPlugin.Control()
+        {
+            public List<CommitReport> run(final TaskSource outputTask)
+            {
+                TransactionalPageOutput tran = runner.open(outputTask, schema,
+                        1);
+                boolean committed = false;
+                try {
+                    for (Page page : PageTestUtils.buildPage(
+                            runtime.getBufferAllocator(), schema, true, 2L,
+                            3.0D, "45", Timestamp.ofEpochMilli(678L), true, 2L,
+                            3.0D, "45", Timestamp.ofEpochMilli(678L))) {
+                        tran.add(page);
+                    }
+                    tran.commit();
+                    committed = true;
+                } finally {
+                    if (!committed) {
+                        tran.abort();
+                    }
+                    tran.close();
+                }
+                return new ArrayList<CommitReport>();
+            }
+        });
+
+        assertEquals(true, fileOutputPlugin.transactionCompleted);
+        assertEquals(2, MockFormatterPlugin.records.size());
+        for (List<Object> record : MockFormatterPlugin.records) {
+            assertEquals(Boolean.TRUE, record.get(0));
+            assertEquals(2L, record.get(1));
+            assertEquals(3.0D, (Double) record.get(2), 0.1D);
+            assertEquals("45", record.get(3));
+            assertEquals(678L, ((Timestamp) record.get(4)).toEpochMilli());
+        }
+    }
+
+    @Test
+    public void testTransactionAborted()
+    {
+        MockFileOutputPlugin fileOutputPlugin = new MockFileOutputPlugin();
+        final FileOutputRunner runner = new FileOutputRunner(fileOutputPlugin);
+
+        ImmutableList<ImmutableMap<String, String>> columns = ImmutableList.of(
+                ImmutableMap.of("name", "col1", "type", "boolean"),
+                ImmutableMap.of("name", "col2", "type", "long"),
+                ImmutableMap.of("name", "col3", "type", "double"),
+                ImmutableMap.of("name", "col4", "type", "string"),
+                ImmutableMap.of("name", "col5", "type", "timestamp"));
+        ConfigSource config = Exec
+                .newConfigSource()
+                .set("type", "unused?")
+                .set("formatter",
+                        ImmutableMap.of("type", "mock", "columns", columns));
+        final Schema schema = config.getNested("formatter")
+                .loadConfig(MockParserPlugin.PluginTask.class)
+                .getSchemaConfig().toSchema();
+
+        try {
+            runner.transaction(config, schema, 1, new OutputPlugin.Control()
+            {
+                public List<CommitReport> run(final TaskSource outputTask)
+                {
+                    TransactionalPageOutput tran = runner.open(outputTask,
+                            schema, 1);
+                    boolean committed = false;
+                    try {
+                        tran.add(null);
+                        tran.commit();
+                        committed = true;
+                    } finally {
+                        if (!committed) {
+                            tran.abort();
+                        }
+                        tran.close();
+                    }
+                    return new ArrayList<CommitReport>();
+                }
+            });
+        } catch (NullPointerException npe) {
+        }
+
+        assertEquals(false, fileOutputPlugin.transactionCompleted);
+    }
+}