embulk-output-td 0.1.4 → 0.1.5

data/src/main/java/org/embulk/output/td/TdOutputPlugin.java

@@ -6,6 +6,7 @@ import java.util.Map;
 import javax.validation.constraints.Min;
 import javax.validation.constraints.Max;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Optional;
 import com.google.common.base.Throwables;
 import com.fasterxml.jackson.annotation.JsonCreator;
@@ -26,7 +27,7 @@ import org.embulk.config.ConfigSource;
 import org.embulk.config.ConfigException;
 import org.embulk.config.Task;
 import org.embulk.config.TaskSource;
-import org.embulk.output.td.RecordWriter.FieldWriterSet;
+import org.embulk.output.td.writer.FieldWriterSet;
 import org.embulk.spi.Exec;
 import org.embulk.spi.ExecSession;
 import org.embulk.spi.OutputPlugin;
@@ -34,7 +35,6 @@ import org.embulk.spi.Schema;
 import org.embulk.spi.TransactionalPageOutput;
 import org.embulk.spi.time.Timestamp;
 import org.embulk.spi.time.TimestampFormatter;
-import org.joda.time.DateTimeZone;
 import org.joda.time.format.DateTimeFormat;
 import org.slf4j.Logger;
 
@@ -61,7 +61,9 @@ public class TdOutputPlugin
 
         //  TODO connect_timeout, read_timeout, send_timeout
 
-        //  TODO mode[append, replace]
+        @Config("mode")
+        @ConfigDefault("\"append\"")
+        public Mode getMode();
 
         @Config("auto_create_table")
         @ConfigDefault("true")
@@ -73,6 +75,9 @@ public class TdOutputPlugin
         @Config("table")
         public String getTable();
 
+        public void setLoadTargetTableName(String name);
+        public String getLoadTargetTableName();
+
         @Config("session")
         @ConfigDefault("null")
         public Optional<String> getSession();
@@ -130,6 +135,37 @@ public class TdOutputPlugin
             extends Task, TimestampFormatter.TimestampColumnOption
     {}
 
+    public enum Mode
+    {
+        APPEND, REPLACE;
+
+        @JsonCreator
+        public static Mode fromConfig(String value)
+        {
+            switch(value) {
+            case "append":
+                return APPEND;
+            case "replace":
+                return REPLACE;
+            default:
+                throw new ConfigException(String.format("Unknown mode '%s'. Supported modes are [append, replace]", value));
+            }
+        }
+
+        @JsonValue
+        public String toString()
+        {
+            switch(this) {
+            case APPEND:
+                return "append";
+            case REPLACE:
+                return "replace";
+            default:
+                throw new IllegalStateException();
+            }
+        }
+    }
+
     public interface HttpProxyTask
             extends Task
     {
@@ -173,7 +209,7 @@ public class TdOutputPlugin
             case "nano": return NANO;
             default:
                 throw new ConfigException(
-                        String.format("Unknown unix_timestamp_unit '%s'. Supported units are sec, milli, micro, and nano"));
+                        String.format("Unknown unix_timestamp_unit '%s'. Supported units are sec, milli, micro, and nano", s));
             }
         }
 
@@ -197,12 +233,8 @@ public class TdOutputPlugin
     {
         final PluginTask task = config.loadConfig(PluginTask.class);
 
-        // TODO mode check
-
         // check column_options is valid or not
-        for (String columnName : task.getColumnOptions().keySet()) {
-            schema.lookupColumn(columnName); // throws SchemaConfigException
-        }
+        checkColumnOptions(schema, task.getColumnOptions());
 
         // generate session name
         task.setSessionName(buildBulkImportSessionName(task, Exec.session()));
@@ -210,11 +242,24 @@ public class TdOutputPlugin
         try (TdApiClient client = newTdApiClient(task)) {
             String databaseName = task.getDatabase();
             String tableName = task.getTable();
-            if (task.getAutoCreateTable()) {
-                createTableIfNotExists(client, databaseName, tableName);
-            } else {
-                // check if the database and/or table exist or not
-                validateTableExists(client, databaseName, tableName);
+
+            switch (task.getMode()) {
+            case APPEND:
+                if (task.getAutoCreateTable()) {
+                    // auto_create_table is valid only with append mode (replace mode always creates a new table)
+                    createTableIfNotExists(client, databaseName, tableName);
+                }
+                else {
+                    // check if the database and/or table exist or not
+                    validateTableExists(client, databaseName, tableName);
+                }
+                task.setLoadTargetTableName(tableName);
+                break;
+
+            case REPLACE:
+                task.setLoadTargetTableName(
+                        createTemporaryTableWithPrefix(client, databaseName, makeTablePrefix(task)));
+                break;
             }
 
             // validate FieldWriterSet configuration before transaction is started
@@ -226,20 +271,32 @@ public class TdOutputPlugin
 
     public ConfigDiff resume(TaskSource taskSource,
             Schema schema, int processorCount,
-            OutputPlugin.Control control) {
+            OutputPlugin.Control control)
+    {
         PluginTask task = taskSource.loadTask(PluginTask.class);
         try (TdApiClient client = newTdApiClient(task)) {
             return doRun(client, task, control);
         }
     }
 
-    private ConfigDiff doRun(TdApiClient client, PluginTask task, OutputPlugin.Control control)
+    @VisibleForTesting
+    ConfigDiff doRun(TdApiClient client, PluginTask task, OutputPlugin.Control control)
     {
-        boolean doUpload = startBulkImportSession(client, task.getSessionName(), task.getDatabase(), task.getTable());
+        boolean doUpload = startBulkImportSession(client, task.getSessionName(), task.getDatabase(), task.getLoadTargetTableName());
         task.setDoUpload(doUpload);
         control.run(task.dump());
         completeBulkImportSession(client, task.getSessionName(), 0);  // TODO perform job priority
 
+        // commit
+        switch (task.getMode()) {
+        case APPEND:
+            // already done
+            break;
+        case REPLACE:
+            // rename table
+            renameTable(client, task.getDatabase(), task.getLoadTargetTableName(), task.getTable());
+        }
+
         ConfigDiff configDiff = Exec.newConfigDiff();
         configDiff.set("last_session", task.getSessionName());
         return configDiff;
@@ -257,14 +314,29 @@ public class TdOutputPlugin
         }
     }
 
-    private TdApiClient newTdApiClient(final PluginTask task)
+    private String makeTablePrefix(PluginTask task)
+    {
+        return task.getTable() + "_" + task.getSessionName();
+    }
+
+    @VisibleForTesting
+    void checkColumnOptions(Schema schema, Map<String, TimestampColumnOption> columnOptions)
+    {
+        for (String columnName : columnOptions.keySet()) {
+            schema.lookupColumn(columnName); // throws SchemaConfigException
+        }
+    }
+
+    @VisibleForTesting
+    TdApiClient newTdApiClient(final PluginTask task)
     {
         Optional<HttpProxyConfig> httpProxyConfig = newHttpProxyConfig(task.getHttpProxy());
         TdApiClientConfig config = new TdApiClientConfig(task.getEndpoint(), task.getUseSsl(), httpProxyConfig);
         TdApiClient client = new TdApiClient(task.getApiKey(), config);
         try {
             client.start();
-        } catch (IOException e) {
+        }
+        catch (IOException e) {
             throw Throwables.propagate(e);
         }
         return client;
@@ -276,37 +348,63 @@ public class TdOutputPlugin
         if (task.isPresent()) {
             HttpProxyTask pt = task.get();
             httpProxyConfig = Optional.of(new HttpProxyConfig(pt.getHost(), pt.getPort(), pt.getUseSsl()));
-        } else {
+        }
+        else {
             httpProxyConfig = Optional.absent();
         }
         return httpProxyConfig;
     }
 
-    private void createTableIfNotExists(TdApiClient client, String databaseName, String tableName)
+    @VisibleForTesting
+    void createTableIfNotExists(TdApiClient client, String databaseName, String tableName)
     {
         log.debug("Creating table \"{}\".\"{}\" if not exists", databaseName, tableName);
         try {
             client.createTable(databaseName, tableName);
             log.debug("Created table \"{}\".\"{}\"", databaseName, tableName);
-        } catch (TdApiNotFoundException e) {
+        }
+        catch (TdApiNotFoundException e) {
             try {
                 client.createDatabase(databaseName);
                 log.debug("Created database \"{}\"", databaseName);
-            } catch (TdApiConflictException ex) {
+            }
+            catch (TdApiConflictException ex) {
                 // ignorable error
             }
             try {
                 client.createTable(databaseName, tableName);
                 log.debug("Created table \"{}\".\"{}\"", databaseName, tableName);
-            } catch (TdApiConflictException exe) {
+            }
+            catch (TdApiConflictException exe) {
                 // ignorable error
             }
-        } catch (TdApiConflictException e) {
+        }
+        catch (TdApiConflictException e) {
             // ignorable error
         }
     }
 
-    private void validateTableExists(TdApiClient client, String databaseName, String tableName)
+    @VisibleForTesting
+    String createTemporaryTableWithPrefix(TdApiClient client, String databaseName, String tablePrefix)
+            throws TdApiConflictException
+    {
+        String tableName = tablePrefix;
+        while (true) {
+            log.debug("Creating temporal table \"{}\".\"{}\"", databaseName, tableName);
+            try {
+                client.createTable(databaseName, tableName);
+                log.debug("Created temporal table \"{}\".\"{}\"", databaseName, tableName);
+                return tableName;
+            }
+            catch (TdApiConflictException e) {
+                log.debug("\"{}\".\"{}\" table already exists. Renaming temporal table.", databaseName, tableName);
+                tableName += "_";
+            }
+        }
+    }
+
+    @VisibleForTesting
+    void validateTableExists(TdApiClient client, String databaseName, String tableName)
     {
         try {
             for (TDTable table : client.getTables(databaseName)) {
@@ -315,16 +413,19 @@ public class TdOutputPlugin
                 }
             }
             throw new ConfigException(String.format("Table \"%s\".\"%s\" doesn't exist", databaseName, tableName));
-        } catch (TdApiNotFoundException ex) {
+        }
+        catch (TdApiNotFoundException ex) {
             throw new ConfigException(String.format("Database \"%s\" doesn't exist", databaseName), ex);
         }
     }
 
-    private String buildBulkImportSessionName(PluginTask task, ExecSession exec)
+    @VisibleForTesting
+    String buildBulkImportSessionName(PluginTask task, ExecSession exec)
     {
         if (task.getSession().isPresent()) {
             return task.getSession().get();
-        } else {
+        }
+        else {
             Timestamp time = exec.getTransactionTime(); // TODO implement Exec.getTransactionUniqueName()
             return String.format("embulk_%s_%09d",
                     DateTimeFormat.forPattern("yyyyMMdd_HHmmss").withZoneUTC().print(time.getEpochSecond() * 1000),
@@ -333,14 +434,16 @@ public class TdOutputPlugin
     }
 
     // return false if all files are already uploaded
-    private boolean startBulkImportSession(TdApiClient client,
+    @VisibleForTesting
+    boolean startBulkImportSession(TdApiClient client,
             String sessionName, String databaseName, String tableName)
     {
         log.info("Create bulk_import session {}", sessionName);
         TDBulkImportSession session;
         try {
             client.createBulkImportSession(sessionName, databaseName, tableName);
-        } catch (TdApiConflictException ex) {
+        }
+        catch (TdApiConflictException ex) {
             // ignorable error
         }
         session = client.getBulkImportSession(sessionName);
@@ -366,7 +469,8 @@ public class TdOutputPlugin
         }
     }
 
-    private void completeBulkImportSession(TdApiClient client, String sessionName, int priority)
+    @VisibleForTesting
+    void completeBulkImportSession(TdApiClient client, String sessionName, int priority)
     {
         TDBulkImportSession session = client.getBulkImportSession(sessionName);
 
@@ -376,7 +480,8 @@ public class TdOutputPlugin
                 // freeze
                 try {
                     client.freezeBulkImportSession(sessionName);
-                } catch (TdApiConflictException e) {
+                }
+                catch (TdApiConflictException e) {
                     // ignorable error
                 }
             }
@@ -417,7 +522,8 @@ public class TdOutputPlugin
         }
     }
 
-    private TDBulkImportSession waitForStatusChange(TdApiClient client, String sessionName,
+    @VisibleForTesting
+    TDBulkImportSession waitForStatusChange(TdApiClient client, String sessionName,
             ImportStatus current, ImportStatus expecting, String operation)
     {
         TDBulkImportSession importSession;
@@ -427,21 +533,47 @@ public class TdOutputPlugin
             if (importSession.is(expecting)) {
                 return importSession;
 
-            } else if (importSession.is(current)) {
+            }
+            else if (importSession.is(current)) {
                 // in progress
 
-            } else {
+            }
+            else {
                 throw new RuntimeException(String.format("Failed to %s bulk import session '%s'",
                             operation, sessionName));
             }
 
             try {
                 Thread.sleep(3000);
-            } catch (InterruptedException e) {
+            }
+            catch (InterruptedException e) {
             }
         }
     }
 
+    @VisibleForTesting
+    void renameTable(TdApiClient client, String databaseName, String oldName, String newName)
+    {
+        log.debug("Renaming table \"{}\".\"{}\" to \"{}\"", databaseName, oldName, newName);
+        try {
+            client.renameTable(databaseName, oldName, newName);
+        }
+        catch (TdApiConflictException e) {
+            try {
+                client.deleteTable(databaseName, newName);
+                log.debug("Deleted original table \"{}\".\"{}\"", databaseName, newName);
+            }
+            catch (TdApiNotFoundException ex) {
+                // ignoreable error
+            }
+            catch (IOException ex) {
+                throw Throwables.propagate(ex);
+            }
+
+            client.renameTable(databaseName, oldName, newName);
+        }
+    }
+
     @Override
     public TransactionalPageOutput open(TaskSource taskSource, Schema schema, int taskIndex)
     {
@@ -450,14 +582,17 @@ public class TdOutputPlugin
         RecordWriter closeLater = null;
         try {
             FieldWriterSet fieldWriters = new FieldWriterSet(log, task, schema);
-            RecordWriter recordWriter = closeLater = new RecordWriter(task, taskIndex, newTdApiClient(task), fieldWriters);
+            closeLater = new RecordWriter(task, taskIndex, newTdApiClient(task), fieldWriters);
+            RecordWriter recordWriter = closeLater;
             recordWriter.open(schema);
             closeLater = null;
             return recordWriter;
 
-        } catch (IOException e) {
+        }
+        catch (IOException e) {
             throw Throwables.propagate(e);
-        } finally {
+        }
+        finally {
             if (closeLater != null) {
                 closeLater.close();
             }

data/src/main/java/org/embulk/output/td/writer/BooleanFieldWriter.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 BooleanFieldWriter
+        extends FieldWriter
+{
+    public BooleanFieldWriter(String keyName)
+    {
+        super(keyName);
+    }
+
+    @Override
+    public void writeValue(MsgpackGZFileBuilder builder, PageReader reader, Column column)
+            throws IOException
+    {
+        builder.writeBoolean(reader.getBoolean(column));
+    }
+}

data/src/main/java/org/embulk/output/td/writer/DoubleFieldWriter.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 DoubleFieldWriter
+        extends FieldWriter
+{
+    public DoubleFieldWriter(String keyName)
+    {
+        super(keyName);
+    }
+
+    @Override
+    public void writeValue(MsgpackGZFileBuilder builder, PageReader reader, Column column)
+            throws IOException
+    {
+        builder.writeDouble(reader.getDouble(column));
+    }
+}