embulk-filter-protobuf 0.1.0

data/lib/embulk/filter/protobuf.rb

@@ -0,0 +1,3 @@
+Embulk::JavaPlugin.register_filter(
+  "protobuf", "org.embulk.filter.protobuf.ProtobufFilterPlugin",
+  File.expand_path('../../../../classpath', __FILE__))

data/src/main/java/org/embulk/filter/protobuf/ColumnVisitorImpl.java

@@ -0,0 +1,264 @@
+package org.embulk.filter.protobuf;
+
+import com.google.common.io.BaseEncoding;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.Message;
+import com.google.protobuf.MessageOrBuilder;
+import com.google.protobuf.util.JsonFormat;
+
+import org.embulk.filter.protobuf.ProtobufFilterPlugin.ColumnTask;
+import org.embulk.filter.protobuf.ProtobufFilterPlugin.PluginTask;
+
+import org.embulk.plugin.PluginClassLoader;
+import org.embulk.spi.Column;
+import org.embulk.spi.ColumnVisitor;
+import org.embulk.spi.DataException;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.PageReader;
+
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.net.URLClassLoader;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class ColumnVisitorImpl implements ColumnVisitor
+{
+    private final PluginTask pluginTask;
+    private final PageReader pageReader;
+    private final PageBuilder pageBuilder;
+    private final Map<String, ColumnTask> columnTaskMap;
+
+    ColumnVisitorImpl(PluginTask task, PageReader reader, PageBuilder builder)
+    {
+        this.pluginTask = task;
+        this.pageReader = reader;
+        this.pageBuilder = builder;
+        this.columnTaskMap = getColumnMap(task.getColumns());
+        addProtobufJarToClasspath();
+    }
+
+    private Map<String, ColumnTask> getColumnMap(
+            List<ColumnTask> columnTasks)
+    {
+        Map<String, ColumnTask> m = new HashMap<>();
+        for (ColumnTask columnTask : columnTasks) {
+            m.put(columnTask.getName(), columnTask);
+        }
+        return m;
+    }
+
+    private void addProtobufJarToClasspath()
+    {
+        Path protobufJarPath = Paths.get(pluginTask.getProtobufJarPath());
+        // FIXME:
+        // getClass().getClassLoader() returns sun.misc.Launcher$AppClassLoader
+        // and it cannot be cast to PluginClassLoader in gradle test.
+        try {
+            PluginClassLoader loader = (PluginClassLoader) getClass().getClassLoader();
+            loader.addPath(protobufJarPath);
+        }
+        catch (ClassCastException e) {
+        }
+    }
+
+    private ColumnTask getColumnTask(Column column)
+    {
+        String colName = column.getName();
+        return columnTaskMap.get(colName);
+    }
+
+    private byte[] decodeMessage(String messageAsString)
+    {
+        byte[] decoded = null;
+        String encoding = pluginTask.getEncoding();
+        if (encoding.equals("Base64")) {
+            decoded = BaseEncoding.base64().decode(messageAsString);
+        }
+        return decoded;
+    }
+
+    private String encodeMessage(byte[] messageAsBytes)
+    {
+        String encoded = null;
+        String encoding = pluginTask.getEncoding();
+        if (encoding.equals("Base64")) {
+            encoded = BaseEncoding.base64().encode(messageAsBytes);
+        }
+        return encoded;
+    }
+
+    private String convertMessageBytesToJson(
+            byte[] messageAsBytes, String messageName)
+    {
+        URLClassLoader loader = (URLClassLoader) getClass().getClassLoader();
+        // Get a message object
+        Object message = null;
+        try {
+            Class<?> messageClass = loader.loadClass(messageName);
+            Method parseFrom = messageClass.getMethod(
+                "parseFrom", byte[].class);
+            message = parseFrom.invoke(
+                (Object) messageClass, (Object) messageAsBytes);
+        }
+        catch (ClassNotFoundException e) {
+            throw new DataException(e);
+        }
+        catch (NoSuchMethodException e) {
+            throw new DataException(e);
+        }
+        catch (IllegalAccessException e) {
+            throw new DataException(e);
+        }
+        catch (InvocationTargetException e) {
+            throw new DataException(e);
+        }
+        // Convert message object to json string
+        String messageAsString = null;
+        try {
+            messageAsString = JsonFormat.printer()
+                .omittingInsignificantWhitespace()
+                .print((MessageOrBuilder) message);
+        }
+        catch (InvalidProtocolBufferException e) {
+            throw new DataException(e);
+        }
+        return messageAsString;
+    }
+
+    private byte[] convertJsonToMessageBytes(
+            String messageAsJson, String messageName)
+    {
+        URLClassLoader loader = (URLClassLoader) getClass().getClassLoader();
+        // Get a message builder object
+        Message.Builder builder = null;
+        try {
+            Class<?> messageClass = loader.loadClass(messageName);
+            Method newBuilder = messageClass.getMethod("newBuilder");
+            builder = (Message.Builder) newBuilder.invoke(
+                (Object) messageClass);
+        }
+        catch (ClassNotFoundException e) {
+            throw new DataException(e);
+        }
+        catch (NoSuchMethodException e) {
+            throw new DataException(e);
+        }
+        catch (IllegalAccessException e) {
+            throw new DataException(e);
+        }
+        catch (InvocationTargetException e) {
+            throw new DataException(e);
+        }
+        // Convert message json to binary
+        byte[] messageAsBytes = null;
+        try {
+            JsonFormat.parser().merge(messageAsJson, builder);
+            messageAsBytes = builder.build().toByteArray();
+        }
+        catch (InvalidProtocolBufferException e) {
+            throw new DataException(e);
+        }
+        return messageAsBytes;
+    }
+
+    private String executeTask(ColumnTask colTask, Column column)
+    {
+        String messageName = colTask.getMessage();
+        // serialize
+        if (pluginTask.getDoSerialize().get()) {
+            String messageAsJson = pageReader.getString(column);
+            byte[] messageAsBytes = convertJsonToMessageBytes(
+                messageAsJson, messageName);
+            return encodeMessage(messageAsBytes);
+        }
+        // deserialize
+        else {
+            String messageAsString = pageReader.getString(column);
+            byte[] messageAsBytes = decodeMessage(messageAsString);
+            return convertMessageBytesToJson(messageAsBytes, messageName);
+        }
+    }
+
+    @Override
+    public void booleanColumn(Column outputColumn)
+    {
+        if (pageReader.isNull(outputColumn)) {
+            pageBuilder.setNull(outputColumn);
+        }
+        else {
+            pageBuilder.setBoolean(
+                outputColumn, pageReader.getBoolean(outputColumn));
+        }
+    }
+
+    @Override
+    public void longColumn(Column outputColumn)
+    {
+        if (pageReader.isNull(outputColumn)) {
+            pageBuilder.setNull(outputColumn);
+        }
+        else {
+            pageBuilder.setLong(
+                outputColumn, pageReader.getLong(outputColumn));
+        }
+    }
+
+    @Override
+    public void doubleColumn(Column outputColumn)
+    {
+        if (pageReader.isNull(outputColumn)) {
+            pageBuilder.setNull(outputColumn);
+        }
+        else {
+            pageBuilder.setDouble(
+                outputColumn, pageReader.getDouble(outputColumn));
+        }
+    }
+
+    @Override
+    public void stringColumn(Column outputColumn)
+    {
+        if (pageReader.isNull(outputColumn)) {
+            pageBuilder.setNull(outputColumn);
+        }
+        else {
+            ColumnTask task = getColumnTask(outputColumn);
+            if (task == null) {
+                pageBuilder.setString(
+                    outputColumn, pageReader.getString(outputColumn));
+            }
+            else {
+                pageBuilder.setString(
+                    outputColumn, executeTask(task, outputColumn));
+            }
+        }
+    }
+
+    @Override
+    public void timestampColumn(Column outputColumn)
+    {
+        if (pageReader.isNull(outputColumn)) {
+            pageBuilder.setNull(outputColumn);
+        }
+        else {
+            pageBuilder.setTimestamp(
+                outputColumn, pageReader.getTimestamp(outputColumn));
+        }
+    }
+
+    @Override
+    public void jsonColumn(Column outputColumn)
+    {
+        if (pageReader.isNull(outputColumn)) {
+            pageBuilder.setNull(outputColumn);
+        }
+        else {
+            pageBuilder.setJson(
+                outputColumn, pageReader.getJson(outputColumn));
+        }
+    }
+}

data/src/main/java/org/embulk/filter/protobuf/ProtobufFilterPlugin.java

@@ -0,0 +1,155 @@
+package org.embulk.filter.protobuf;
+
+import com.google.common.base.Optional;
+
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.ConfigException;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.Task;
+import org.embulk.config.TaskSource;
+import org.embulk.spi.Column;
+import org.embulk.spi.Exec;
+import org.embulk.spi.FilterPlugin;
+import org.embulk.spi.Page;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.PageOutput;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+import org.embulk.spi.type.Types;
+
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Arrays;
+import java.util.List;
+
+public class ProtobufFilterPlugin implements FilterPlugin
+{
+    public interface PluginTask extends Task
+    {
+        @Config("serialize")
+        @ConfigDefault("false")
+        public Optional<Boolean> getDoSerialize();
+
+        @Config("deserialize")
+        @ConfigDefault("false")
+        public Optional<Boolean> getDoDeserialize();
+
+        @Config("encoding")
+        public String getEncoding();
+
+        @Config("protobuf_jar_path")
+        public String getProtobufJarPath();
+
+        @Config("columns")
+        public List<ColumnTask> getColumns();
+    }
+
+    public interface ColumnTask extends Task
+    {
+        @Config("name")
+        public String getName();
+
+        @Config("message")
+        public String getMessage();
+    }
+
+    public void validate(PluginTask pluginTask, Schema inputSchema)
+    {
+        // validate 'serialize' and 'deserialize' in PluginTask
+        boolean doSerialize = pluginTask.getDoSerialize().get();
+        boolean doDeserialize = pluginTask.getDoDeserialize().get();
+        boolean bothTrue = doSerialize && doDeserialize;
+        boolean bothFalse = !doSerialize && !doDeserialize;
+        if (bothTrue || bothFalse) {
+            String errMsg = "Specify either 'serialize: true' or 'deserialize: true'.";
+            throw new ConfigException(errMsg);
+        }
+        // validate 'encoding' in PluginTask
+        String[] allowedEncordings = {"Base64"};
+        String encoding = pluginTask.getEncoding();
+        if (!Arrays.asList(allowedEncordings).contains(encoding)) {
+            String errMsg = "Specify 'encoding: Base64'.";
+            throw new ConfigException(errMsg);
+        }
+        // validate 'protobuf_jar_path' in PluginTask
+        Path protobufJarPath = Paths.get(pluginTask.getProtobufJarPath());
+        if (!protobufJarPath.toFile().exists()) {
+            String errMsg = "The jar file does not exist.";
+            throw new ConfigException(errMsg);
+        }
+        // validate 'name' in ColumnTask
+        for (ColumnTask colTask : pluginTask.getColumns()) {
+            // throws exception when the column does not exist
+            Column column = inputSchema.lookupColumn(colTask.getName());
+            // TODO: accept both STRING and JSON type when 'serilialize': true
+            if (!Types.STRING.equals(column.getType())) {
+                String errMsg = "Type of input columns must be string.";
+                throw new ConfigException(errMsg);
+            }
+        }
+    }
+
+    @Override
+    public void transaction(ConfigSource config, Schema inputSchema,
+            FilterPlugin.Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+        validate(task, inputSchema);
+        Schema outputSchema = inputSchema;
+        control.run(task.dump(), outputSchema);
+    }
+
+    @Override
+    public PageOutput open(TaskSource taskSource, Schema inputSchema,
+            Schema outputSchema, PageOutput output)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+        PageBuilder pageBuilder = new PageBuilder(
+            Exec.getBufferAllocator(), outputSchema, output);
+        PageReader pageReader = new PageReader(inputSchema);
+        ColumnVisitorImpl visitor = new ColumnVisitorImpl(
+            task, pageReader, pageBuilder);
+
+        return new PageOutputImpl(
+            pageReader, pageBuilder, outputSchema, visitor);
+    }
+
+    public static class PageOutputImpl implements PageOutput
+    {
+        private PageReader pageReader;
+        private PageBuilder pageBuilder;
+        private Schema outputSchema;
+        private ColumnVisitorImpl visitor;
+
+        PageOutputImpl(PageReader pageReader, PageBuilder pageBuilder, Schema outputSchema, ColumnVisitorImpl visitor)
+        {
+            this.pageReader = pageReader;
+            this.pageBuilder = pageBuilder;
+            this.outputSchema = outputSchema;
+            this.visitor = visitor;
+        }
+
+        @Override
+        public void add(Page page)
+        {
+            pageReader.setPage(page);
+            while (pageReader.nextRecord()) {
+                outputSchema.visitColumns(visitor);
+                pageBuilder.addRecord();
+            }
+        }
+
+        @Override
+        public void finish()
+        {
+            pageBuilder.finish();
+        }
+
+        @Override
+        public void close()
+        {
+            pageBuilder.close();
+        }
+    };
+}

data/src/test/java/org/embulk/filter/protobuf/TestColumnVisitorImpl.java

@@ -0,0 +1,117 @@
+package org.embulk.filter.protobuf;
+
+import org.embulk.EmbulkTestRuntime;
+import org.embulk.filter.protobuf.ProtobufFilterPlugin.PageOutputImpl;
+import org.embulk.filter.protobuf.ProtobufFilterPlugin.PluginTask;
+import org.embulk.spi.Page;
+import org.embulk.spi.PageBuilder;
+import org.embulk.spi.PageOutput;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.PageTestUtils;
+import org.embulk.spi.Schema;
+import org.embulk.spi.TestPageBuilderReader.MockPageOutput;
+import org.embulk.spi.type.Types;
+import org.embulk.spi.util.Pages;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import static org.embulk.filter.protobuf.TestProtobufFilterPlugin.taskFromYamlString;
+import static org.junit.Assert.assertEquals;
+
+import java.io.File;
+import java.util.List;
+
+public class TestColumnVisitorImpl
+{
+    @Rule
+    public EmbulkTestRuntime runtime;
+
+    private String protobufJarPath;
+
+    public TestColumnVisitorImpl()
+    {
+        this.runtime = new EmbulkTestRuntime();
+
+        String pluginBasePath = new File(".").getAbsoluteFile().getParent();
+        this.protobufJarPath = String.format(
+            "%s/example/AddressBookProtosProto3Syntax.jar",
+            pluginBasePath);
+    }
+
+    private List<Object[]> filter(
+            PluginTask task, Schema inputSchema, Object... objects)
+    {
+        MockPageOutput output = new MockPageOutput();
+        Schema outputSchema = inputSchema;
+        PageBuilder pageBuilder = new PageBuilder(
+            runtime.getBufferAllocator(), outputSchema, output);
+        PageReader pageReader = new PageReader(inputSchema);
+        ColumnVisitorImpl visitor = new ColumnVisitorImpl(
+            task, pageReader, pageBuilder);
+
+        List<Page> pages = PageTestUtils.buildPage(
+            runtime.getBufferAllocator(), inputSchema, objects);
+        PageOutput mockPageOutput = new PageOutputImpl(
+            pageReader, pageBuilder, outputSchema, visitor);
+        for (Page page : pages) {
+            mockPageOutput.add(page);
+        }
+        mockPageOutput.finish();
+        mockPageOutput.close();
+        return Pages.toObjects(outputSchema, output.pages);
+    }
+
+    @Test
+    public void testExecuteTask_serialize()
+    {
+        PluginTask task = taskFromYamlString(
+            "type: protobuf",
+            "serialize: true",
+            "encoding: Base64",
+            "protobuf_jar_path: " + protobufJarPath,
+            "columns:",
+            "  - {name: to serialize, message: com.example.tutorial.AddressBookProtos$Person}"
+        );
+        Schema inputSchema = Schema.builder()
+            .add("to serialize", Types.STRING)
+            .build();
+        List<Object[]> records = filter(task, inputSchema,
+            // generated from proto2-syntax .proto
+            "{\"name\":\"John Doe\",\"id\":1234,\"email\":\"jdoe@example.com\",\"phone\":[{\"number\":\"111-0000\",\"type\":\"MOBILE\"},{\"number\":\"555-4321\",\"type\":\"HOME\"}]}",
+            // generated from proto3-syntax .proto
+            "{\"name\":\"John Doe\",\"id\":1234,\"email\":\"jdoe@example.com\",\"phone\":[{\"number\":\"111-0000\"},{\"number\":\"555-4321\",\"type\":\"HOME\"}]}"
+        );
+        assertEquals(2, records.size());
+        String expected = "CghKb2huIERvZRDSCRoQamRvZUBleGFtcGxlLmNvbSIKCggxMTEtMDAwMCIMCgg1NTUtNDMyMRAB";
+        assertEquals(expected, records.get(0)[0]);
+        assertEquals(expected, records.get(1)[0]);
+    }
+
+    @Test
+    public void testExecuteTask_deserialize()
+    {
+        PluginTask task = taskFromYamlString(
+            "type: protobuf",
+            "deserialize: true",
+            "encoding: Base64",
+            "protobuf_jar_path: " + protobufJarPath,
+            "columns:",
+            "  - {name: to deserialize, message: com.example.tutorial.AddressBookProtos$Person}"
+        );
+        Schema inputSchema = Schema.builder()
+            .add("to deserialize", Types.STRING)
+            .build();
+        List<Object[]> records = filter(
+            task, inputSchema,
+            // generated from proto2-syntax .proto
+            "CghKb2huIERvZRDSCRoQamRvZUBleGFtcGxlLmNvbSIMCggxMTEtMDAwMBAAIgwKCDU1NS00MzIxEAE=",
+            // generated from proto3-syntax .proto
+            "CghKb2huIERvZRDSCRoQamRvZUBleGFtcGxlLmNvbSIKCggxMTEtMDAwMCIMCgg1NTUtNDMyMRAB"
+        );
+        assertEquals(2, records.size());
+        String expected = "{\"name\":\"John Doe\",\"id\":1234,\"email\":\"jdoe@example.com\",\"phone\":[{\"number\":\"111-0000\"},{\"number\":\"555-4321\",\"type\":\"HOME\"}]}";
+        assertEquals(expected, records.get(0)[0]);
+        assertEquals(expected, records.get(1)[0]);
+    }
+}