embulk-output-td 0.1.4 → 0.1.5

data/src/main/java/org/embulk/output/td/writer/FieldWriter.java

@@ -0,0 +1,38 @@
+package org.embulk.output.td.writer;
+
+import org.embulk.output.td.MsgpackGZFileBuilder;
+import org.embulk.spi.Column;
+import org.embulk.spi.PageReader;
+
+import java.io.IOException;
+
+public abstract class FieldWriter
+{
+    private final String keyName;
+
+    protected FieldWriter(String keyName)
+    {
+        this.keyName = keyName;
+    }
+
+    public void writeKeyValue(MsgpackGZFileBuilder builder, PageReader reader, Column column)
+            throws IOException
+    {
+        writeKey(builder);
+        if (reader.isNull(column)) {
+            builder.writeNil();
+        }
+        else {
+            writeValue(builder, reader, column);
+        }
+    }
+
+    private void writeKey(MsgpackGZFileBuilder builder)
+            throws IOException
+    {
+        builder.writeString(keyName);
+    }
+
+    protected abstract void writeValue(MsgpackGZFileBuilder builder, PageReader reader, Column column)
+            throws IOException;
+}

data/src/main/java/org/embulk/output/td/writer/FieldWriterSet.java

@@ -0,0 +1,206 @@
+package org.embulk.output.td.writer;
+
+import com.google.common.base.Optional;
+import org.embulk.config.ConfigException;
+import org.embulk.output.td.TdOutputPlugin;
+import org.embulk.spi.Column;
+import org.embulk.spi.Schema;
+import org.embulk.spi.time.TimestampFormatter;
+import org.embulk.spi.type.BooleanType;
+import org.embulk.spi.type.DoubleType;
+import org.embulk.spi.type.LongType;
+import org.embulk.spi.type.StringType;
+import org.embulk.spi.type.TimestampType;
+import org.embulk.spi.type.Type;
+import org.embulk.spi.util.Timestamps;
+import org.slf4j.Logger;
+
+public class FieldWriterSet
+{
+    private enum ColumnWriterMode
+    {
+        PRIMARY_KEY,
+        SIMPLE_VALUE,
+        DUPLICATE_PRIMARY_KEY;
+    }
+
+    private final int fieldCount;
+    private final FieldWriter[] fieldWriters;
+
+    public FieldWriterSet(Logger log, TdOutputPlugin.PluginTask task, Schema schema)
+    {
+        Optional<String> userDefinedPrimaryKeySourceColumnName = task.getTimeColumn();
+        boolean hasPkWriter = false;
+        int duplicatePrimaryKeySourceIndex = -1;
+        int firstTimestampColumnIndex = -1;
+
+        int fc = 0;
+        fieldWriters = new FieldWriter[schema.size()];
+        TimestampFormatter[] timestampFormatters = Timestamps.newTimestampColumnFormatters(task, schema, task.getColumnOptions());
+
+        for (int i = 0; i < schema.size(); i++) {
+            String columnName = schema.getColumnName(i);
+            Type columnType = schema.getColumnType(i);
+
+            // choose the mode
+            final ColumnWriterMode mode;
+
+            if (userDefinedPrimaryKeySourceColumnName.isPresent() &&
+                    columnName.equals(userDefinedPrimaryKeySourceColumnName.get())) {
+                // found time_column
+                if ("time".equals(userDefinedPrimaryKeySourceColumnName.get())) {
+                    mode = ColumnWriterMode.PRIMARY_KEY;
+                }
+                else {
+                    mode = ColumnWriterMode.DUPLICATE_PRIMARY_KEY;
+                }
+            }
+            else if ("time".equals(columnName)) {
+                // the column name is same with the primary key name.
+                if (userDefinedPrimaryKeySourceColumnName.isPresent()) {
+                    columnName = newColumnUniqueName(columnName, schema);
+                    mode = ColumnWriterMode.SIMPLE_VALUE;
+                    log.warn("time_column '{}' is set but 'time' column also exists. The existent 'time' column is renamed to {}",
+                            userDefinedPrimaryKeySourceColumnName.get(), "time", "time", columnName);
+                }
+                else {
+                    mode = ColumnWriterMode.PRIMARY_KEY;
+                }
+            }
+            else {
+                mode = ColumnWriterMode.SIMPLE_VALUE;
+            }
+
+            // create the fieldWriters writer depending on the mode
+            final FieldWriter writer;
+
+            switch (mode) {
+                case PRIMARY_KEY:
+                    log.info("Using {}:{} column as the data partitioning key", columnName, columnType);
+                    if (columnType instanceof LongType) {
+                        if (task.getUnixTimestampUnit() != TdOutputPlugin.UnixTimestampUnit.SEC) {
+                            log.warn("time column is converted from {} to seconds", task.getUnixTimestampUnit());
+                        }
+                        writer = new UnixTimestampLongFieldWriter(columnName, task.getUnixTimestampUnit().getFractionUnit());
+                        hasPkWriter = true;
+                    }
+                    else if (columnType instanceof TimestampType) {
+                        writer = new TimestampLongFieldWriter(columnName);
+
+                        hasPkWriter = true;
+                    }
+                    else {
+                        throw new ConfigException(String.format("Type of '%s' column must be long or timestamp but got %s",
+                                columnName, columnType));
+                    }
+                    break;
+
+                case SIMPLE_VALUE:
+                    if (columnType instanceof BooleanType) {
+                        writer = new BooleanFieldWriter(columnName);
+                    }
+                    else if (columnType instanceof LongType) {
+                        writer = new LongFieldWriter(columnName);
+                    }
+                    else if (columnType instanceof DoubleType) {
+                        writer = new DoubleFieldWriter(columnName);
+                    }
+                    else if (columnType instanceof StringType) {
+                        writer = new StringFieldWriter(columnName);
+                    }
+                    else if (columnType instanceof TimestampType) {
+                        writer = new TimestampStringFieldWriter(timestampFormatters[i], columnName);
+                        if (firstTimestampColumnIndex < 0) {
+                            firstTimestampColumnIndex = i;
+                        }
+                    }
+                    else {
+                        throw new ConfigException("Unsupported type: " + columnType);
+                    }
+                    break;
+
+                case DUPLICATE_PRIMARY_KEY:
+                    duplicatePrimaryKeySourceIndex = i;
+                    writer = null;  // handle later
+                    break;
+
+                default:
+                    throw new AssertionError();
+            }
+
+            fieldWriters[i] = writer;
+            fc += 1;
+        }
+
+        if (!hasPkWriter) {
+            // PRIMARY_KEY was not found.
+            if (duplicatePrimaryKeySourceIndex < 0) {
+                if (userDefinedPrimaryKeySourceColumnName.isPresent()) {
+                    throw new ConfigException(String.format("time_column '%s' does not exist", userDefinedPrimaryKeySourceColumnName.get()));
+                }
+                else if (firstTimestampColumnIndex >= 0) {
+                    // if time is not found, use the first timestamp column
+                    duplicatePrimaryKeySourceIndex = firstTimestampColumnIndex;
+                }
+                else {
+                    throw new ConfigException(String.format("TD output plugin requires at least one timestamp column, or a long column named 'time'"));
+                }
+            }
+
+            String columnName = schema.getColumnName(duplicatePrimaryKeySourceIndex);
+            Type columnType = schema.getColumnType(duplicatePrimaryKeySourceIndex);
+
+            FieldWriter writer;
+            if (columnType instanceof LongType) {
+                log.info("Duplicating {}:{} column (unix timestamp {}) to 'time' column as seconds for the data partitioning",
+                        columnName, columnType, task.getUnixTimestampUnit());
+                writer = new UnixTimestampFieldDuplicator(columnName, "time", task.getUnixTimestampUnit().getFractionUnit());
+            }
+            else if (columnType instanceof TimestampType) {
+                log.info("Duplicating {}:{} column to 'time' column as seconds for the data partitioning",
+                        columnName, columnType);
+                writer = new TimestampFieldLongDuplicator(timestampFormatters[duplicatePrimaryKeySourceIndex], columnName, "time");
+            }
+            else {
+                throw new ConfigException(String.format("Type of '%s' column must be long or timestamp but got %s",
+                        columnName, columnType));
+            }
+
+            // replace existint writer
+            fieldWriters[duplicatePrimaryKeySourceIndex] = writer;
+            fc += 1;
+        }
+
+        fieldCount = fc;
+    }
+
+    private static String newColumnUniqueName(String originalName, Schema schema)
+    {
+        String name = originalName;
+        do {
+            name += "_";
+        }
+        while (containsColumnName(schema, name));
+        return name;
+    }
+
+    private static boolean containsColumnName(Schema schema, String name)
+    {
+        for (Column c : schema.getColumns()) {
+            if (c.getName().equals(name)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    public FieldWriter getFieldWriter(int index)
+    {
+        return fieldWriters[index];
+    }
+
+    public int getFieldCount()
+    {
+        return fieldCount;
+    }
+}

data/src/main/java/org/embulk/output/td/writer/LongFieldWriter.java

@@ -0,0 +1,23 @@
+package org.embulk.output.td.writer;
+
+import org.embulk.output.td.MsgpackGZFileBuilder;
+import org.embulk.spi.Column;
+import org.embulk.spi.PageReader;
+
+import java.io.IOException;
+
+public class LongFieldWriter
+        extends FieldWriter
+{
+    LongFieldWriter(String keyName)
+    {
+        super(keyName);
+    }
+
+    @Override
+    public void writeValue(MsgpackGZFileBuilder builder, PageReader reader, Column column)
+            throws IOException
+    {
+        builder.writeLong(reader.getLong(column));
+    }
+}

data/src/main/java/org/embulk/output/td/writer/StringFieldWriter.java

@@ -0,0 +1,23 @@
+package org.embulk.output.td.writer;
+
+import org.embulk.output.td.MsgpackGZFileBuilder;
+import org.embulk.spi.Column;
+import org.embulk.spi.PageReader;
+
+import java.io.IOException;
+
+public class StringFieldWriter
+        extends FieldWriter
+{
+    public StringFieldWriter(String keyName)
+    {
+        super(keyName);
+    }
+
+    @Override
+    public void writeValue(MsgpackGZFileBuilder builder, PageReader reader, Column column)
+            throws IOException
+    {
+        builder.writeString(reader.getString(column));
+    }
+}

data/src/main/java/org/embulk/output/td/writer/TimestampFieldLongDuplicator.java

@@ -0,0 +1,28 @@
+package org.embulk.output.td.writer;
+
+import org.embulk.output.td.MsgpackGZFileBuilder;
+import org.embulk.spi.Column;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.time.TimestampFormatter;
+
+import java.io.IOException;
+
+public class TimestampFieldLongDuplicator
+        extends TimestampStringFieldWriter
+{
+    private final TimestampLongFieldWriter timeFieldWriter;
+
+    public TimestampFieldLongDuplicator(TimestampFormatter formatter, String keyName, String longDuplicateKeyName)
+    {
+        super(formatter, keyName);
+        timeFieldWriter = new TimestampLongFieldWriter(longDuplicateKeyName);
+    }
+
+    @Override
+    public void writeValue(MsgpackGZFileBuilder builder, PageReader reader, Column column)
+            throws IOException
+    {
+        super.writeValue(builder, reader, column);
+        timeFieldWriter.writeKeyValue(builder, reader, column);
+    }
+}

data/src/main/java/org/embulk/output/td/writer/TimestampLongFieldWriter.java

@@ -0,0 +1,23 @@
+package org.embulk.output.td.writer;
+
+import org.embulk.output.td.MsgpackGZFileBuilder;
+import org.embulk.spi.Column;
+import org.embulk.spi.PageReader;
+
+import java.io.IOException;
+
+public class TimestampLongFieldWriter
+        extends FieldWriter
+{
+    public TimestampLongFieldWriter(String keyName)
+    {
+        super(keyName);
+    }
+
+    @Override
+    public void writeValue(MsgpackGZFileBuilder builder, PageReader reader, Column column)
+            throws IOException
+    {
+        builder.writeLong(reader.getTimestamp(column).getEpochSecond());
+    }
+}

data/src/main/java/org/embulk/output/td/writer/TimestampStringFieldWriter.java

@@ -0,0 +1,27 @@
+package org.embulk.output.td.writer;
+
+import org.embulk.output.td.MsgpackGZFileBuilder;
+import org.embulk.spi.Column;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.time.TimestampFormatter;
+
+import java.io.IOException;
+
+public class TimestampStringFieldWriter
+        extends FieldWriter
+{
+    private final TimestampFormatter formatter;
+
+    public TimestampStringFieldWriter(TimestampFormatter formatter, String keyName)
+    {
+        super(keyName);
+        this.formatter = formatter;
+    }
+
+    @Override
+    public void writeValue(MsgpackGZFileBuilder builder, PageReader reader, Column column)
+            throws IOException
+    {
+        builder.writeString(formatter.format(reader.getTimestamp(column)));
+    }
+}

data/src/main/java/org/embulk/output/td/writer/UnixTimestampFieldDuplicator.java

@@ -0,0 +1,27 @@
+package org.embulk.output.td.writer;
+
+import org.embulk.output.td.MsgpackGZFileBuilder;
+import org.embulk.spi.Column;
+import org.embulk.spi.PageReader;
+
+import java.io.IOException;
+
+public class UnixTimestampFieldDuplicator
+        extends LongFieldWriter
+{
+    private final UnixTimestampLongFieldWriter timeFieldWriter;
+
+    public UnixTimestampFieldDuplicator(String keyName, String duplicateKeyName, int fractionUnit)
+    {
+        super(keyName);
+        timeFieldWriter = new UnixTimestampLongFieldWriter(duplicateKeyName, fractionUnit);
+    }
+
+    @Override
+    public void writeValue(MsgpackGZFileBuilder builder, PageReader reader, Column column)
+            throws IOException
+    {
+        super.writeValue(builder, reader, column);
+        timeFieldWriter.writeKeyValue(builder, reader, column);
+    }
+}

data/src/main/java/org/embulk/output/td/writer/UnixTimestampLongFieldWriter.java

@@ -0,0 +1,26 @@
+package org.embulk.output.td.writer;
+
+import org.embulk.output.td.MsgpackGZFileBuilder;
+import org.embulk.spi.Column;
+import org.embulk.spi.PageReader;
+
+import java.io.IOException;
+
+public class UnixTimestampLongFieldWriter
+        extends FieldWriter
+{
+    private final int fractionUnit;
+
+    UnixTimestampLongFieldWriter(String keyName, int fractionUnit)
+    {
+        super(keyName);
+        this.fractionUnit = fractionUnit;
+    }
+
+    @Override
+    public void writeValue(MsgpackGZFileBuilder builder, PageReader reader, Column column)
+            throws IOException
+    {
+        builder.writeLong(reader.getLong(column) / fractionUnit);
+    }
+}

data/src/test/java/com/treasuredata/api/TestTdApiClient.java

@@ -47,7 +47,7 @@ public class TestTdApiClient
     public void stopTdApiClient()
             throws Exception
     {
-        client.close();;
+        client.close();
     }
 
     private static final String DATABASE_LIST_JSON =

data/src/test/java/org/embulk/output/td/TestRecordWriter.java

@@ -0,0 +1,198 @@
+package org.embulk.output.td;
+
+import com.treasuredata.api.TdApiClient;
+import org.embulk.EmbulkTestRuntime;
+import org.embulk.output.td.TdOutputPlugin.PluginTask;
+import org.embulk.output.td.writer.FieldWriterSet;
+import org.embulk.spi.Page;
+import org.embulk.spi.PageTestUtils;
+import org.embulk.spi.Schema;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.type.Types;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.msgpack.MessagePack;
+import org.msgpack.type.MapValue;
+import org.msgpack.unpacker.Unpacker;
+import org.slf4j.Logger;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.util.zip.GZIPInputStream;
+
+import static org.embulk.output.td.TestTdOutputPlugin.config;
+import static org.embulk.output.td.TestTdOutputPlugin.fieldWriters;
+import static org.embulk.output.td.TestTdOutputPlugin.plugin;
+import static org.embulk.output.td.TestTdOutputPlugin.pluginTask;
+import static org.embulk.output.td.TestTdOutputPlugin.schema;
+import static org.embulk.output.td.TestTdOutputPlugin.recordWriter;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.spy;
+import static org.msgpack.type.ValueFactory.createRawValue;
+
+public class TestRecordWriter
+{
+    @Rule
+    public EmbulkTestRuntime runtime = new EmbulkTestRuntime();
+
+    private MessagePack msgpack;
+    private Logger log;
+    private Schema schema;
+    private TdOutputPlugin plugin; // mock
+    private PluginTask task;
+    private RecordWriter recordWriter; // mock
+
+    @Before
+    public void createResources()
+    {
+        msgpack = new MessagePack();
+        log = runtime.getExec().getLogger(TestRecordWriter.class);
+
+        schema = schema("time", Types.LONG, "_c0", Types.LONG, "_c1", Types.STRING,
+                "_c2", Types.BOOLEAN, "_c3", Types.DOUBLE, "_c4", Types.TIMESTAMP);
+
+        plugin = plugin();
+        task = pluginTask(config().set("session_name", "my_session"));
+
+        TdApiClient client = plugin.newTdApiClient(task);
+        FieldWriterSet fieldWriters = fieldWriters(log, task, schema);
+        recordWriter = recordWriter(task, client, fieldWriters);
+    }
+
+    @Test
+    public void checkOpenAndClose()
+            throws Exception
+    {
+        // confirm that no error happens
+        try {
+            recordWriter.open(schema);
+        }
+        finally {
+            recordWriter.close();
+        }
+    }
+
+    @Test
+    public void checkFlushAndFinish()
+            throws Exception
+    {
+        TdApiClient client = spy(plugin.newTdApiClient(task));
+
+        { // add no record
+            RecordWriter recordWriter = recordWriter(task, client, fieldWriters(log, task, schema));
+            try {
+                recordWriter.open(schema);
+            }
+            finally {
+                recordWriter.finish();
+            }
+        }
+
+        { // add 1 record
+            doNothing().when(client).uploadBulkImportPart(anyString(), anyString(), any(File.class));
+
+            RecordWriter recordWriter = recordWriter(task, client, fieldWriters(log, task, schema));
+            try {
+                recordWriter.open(schema);
+
+                // values are not null
+                for (Page page : PageTestUtils.buildPage(runtime.getBufferAllocator(), schema,
+                        1442595600L, 0L, "v", true, 0.0, Timestamp.ofEpochSecond(1442595600L))) {
+                    recordWriter.add(page);
+                }
+            }
+            finally {
+                recordWriter.finish();
+            }
+        }
+    }
+
+    @Test
+    public void addNonNullValues()
+            throws Exception
+    {
+        try {
+            recordWriter.open(schema);
+
+            // values are not null
+            for (Page page : PageTestUtils.buildPage(runtime.getBufferAllocator(), schema,
+                    1442595600L, 0L, "v", true, 0.0, Timestamp.ofEpochSecond(1442595600L))) {
+                recordWriter.add(page);
+            }
+
+            MsgpackGZFileBuilder builder = recordWriter.getBuilder();
+            builder.finish();
+
+            // record count 1
+            assertEquals(1, builder.getRecordCount());
+
+            Unpacker u = msgpack.createUnpacker(new GZIPInputStream(new FileInputStream(builder.getFile())));
+            MapValue v = u.readValue().asMapValue();
+
+            // compare actual values
+            assertEquals(1442595600L, v.get(createRawValue("time")).asIntegerValue().getLong());
+            assertEquals(0L, v.get(createRawValue("_c0")).asIntegerValue().getLong());
+            assertEquals("v", v.get(createRawValue("_c1")).asRawValue().getString());
+            assertEquals(true, v.get(createRawValue("_c2")).asBooleanValue().getBoolean());
+            assertEquals(0.0, v.get(createRawValue("_c3")).asFloatValue().getDouble(), 0.000001);
+            assertEquals("2015-09-18 17:00:00.000", v.get(createRawValue("_c4")).asRawValue().getString());
+
+        }
+        finally {
+            recordWriter.close();
+        }
+    }
+
+    @Test
+    public void addNullValues()
+            throws Exception
+    {
+        try {
+            recordWriter.open(schema);
+
+            // values are not null
+            for (Page page : PageTestUtils.buildPage(runtime.getBufferAllocator(), schema,
+                    1442595600L, null, null, null, null, null)) {
+                recordWriter.add(page);
+            }
+
+            MsgpackGZFileBuilder builder = recordWriter.getBuilder();
+            builder.finish();
+
+            // record count 1
+            assertEquals(1, builder.getRecordCount());
+
+            Unpacker u = msgpack.createUnpacker(new GZIPInputStream(new FileInputStream(builder.getFile())));
+            MapValue v = u.readValue().asMapValue();
+
+            // compare actual values
+            assertTrue(v.get(createRawValue("_c0")).isNilValue());
+            assertTrue(v.get(createRawValue("_c1")).isNilValue());
+            assertTrue(v.get(createRawValue("_c2")).isNilValue());
+            assertTrue(v.get(createRawValue("_c3")).isNilValue());
+            assertTrue(v.get(createRawValue("_c4")).isNilValue());
+
+        }
+        finally {
+            recordWriter.close();
+        }
+    }
+
+    @Test
+    public void doAbortNorthing()
+    {
+        recordWriter.abort();
+        // no error happen
+    }
+
+    @Test
+    public void checkCommitReport()
+    {
+        assertTrue(recordWriter.commit().isEmpty());
+    }
+}