embulk-output-td 0.1.0

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

@@ -0,0 +1,390 @@
+package org.embulk.output;
+
+import java.io.IOException;
+import java.util.List;
+import javax.validation.constraints.Min;
+import javax.validation.constraints.Max;
+
+import com.google.common.base.Optional;
+import com.google.common.base.Throwables;
+import com.treasuredata.api.TdApiClient;
+import com.treasuredata.api.TdApiClientConfig;
+import com.treasuredata.api.TdApiClientConfig.HttpProxyConfig;
+import com.treasuredata.api.TdApiConflictException;
+import com.treasuredata.api.TdApiNotFoundException;
+import com.treasuredata.api.TdApiException;
+import com.treasuredata.api.model.TDBulkImportSession;
+import com.treasuredata.api.model.TDBulkImportSession.ImportStatus;
+import com.treasuredata.api.model.TDDatabase;
+import com.treasuredata.api.model.TDTable;
+import org.embulk.config.CommitReport;
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.ConfigDiff;
+import org.embulk.config.ConfigInject;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.ConfigException;
+import org.embulk.config.Task;
+import org.embulk.config.TaskSource;
+import org.embulk.output.RecordWriter.FieldWriterSet;
+import org.embulk.spi.Exec;
+import org.embulk.spi.ExecSession;
+import org.embulk.spi.OutputPlugin;
+import org.embulk.spi.Schema;
+import org.embulk.spi.TransactionalPageOutput;
+import org.embulk.spi.time.Timestamp;
+import org.joda.time.format.DateTimeFormat;
+import org.jruby.embed.ScriptingContainer;
+import org.slf4j.Logger;
+
+public class TdOutputPlugin
+        implements OutputPlugin
+{
+    public interface PluginTask
+            extends Task
+    {
+        @Config("apikey")
+        public String getApiKey();
+
+        @Config("endpoint")
+        @ConfigDefault("\"api.treasuredata.com\"")
+        public String getEndpoint();
+
+        @Config("use_ssl")
+        @ConfigDefault("true")
+        public boolean getUseSsl();
+
+        @Config("http_proxy")
+        @ConfigDefault("null")
+        public Optional<HttpProxyTask> getHttpProxy();
+
+        //  TODO connect_timeout, read_timeout, send_timeout
+
+        @Config("auto_create_table")
+        @ConfigDefault("true")
+        public boolean getAutoCreateTable();
+
+        @Config("database")
+        public String getDatabase();
+
+        @Config("table")
+        public String getTable();
+
+        @Config("session")
+        @ConfigDefault("null")
+        public Optional<String> getSession();
+
+        @Config("time_column")
+        @ConfigDefault("null")
+        public Optional<String> getTimeColumn();
+
+        @Config("tmpdir")
+        @ConfigDefault("\"/tmp\"")
+        public String getTempDir();
+
+        @Config("upload_concurrency")
+        @ConfigDefault("2")
+        @Min(1)
+        @Max(8)
+        public int getUploadConcurrency();
+
+        @Config("file_split_size")
+        @ConfigDefault("16384") // default 16MB (unit: kb)
+        public long getFileSplitSize();
+
+        @ConfigInject
+        public ScriptingContainer getJRuby();
+
+        public boolean getDoUpload();
+        public void setDoUpload(boolean doUpload);
+
+        public String getSessionName();
+        public void setSessionName(String session);
+    }
+
+    public interface HttpProxyTask
+            extends Task
+    {
+        @Config("host")
+        public String getHost();
+
+        @Config("port")
+        public int getPort();
+
+        @Config("use_ssl")
+        @ConfigDefault("false")
+        public boolean getUseSsl();
+    }
+
+    private final Logger log;
+
+    public TdOutputPlugin()
+    {
+        this.log = Exec.getLogger(getClass());
+    }
+
+    public ConfigDiff transaction(final ConfigSource config, final Schema schema, int processorCount,
+                                  OutputPlugin.Control control)
+    {
+        final PluginTask task = config.loadConfig(PluginTask.class);
+
+        // generate session name
+        task.setSessionName(buildBulkImportSessionName(task, Exec.session()));
+
+        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);
+            }
+
+            // validate FieldWriterSet configuration before transaction is started
+            RecordWriter.validateSchema(log, task, schema);
+
+            return doRun(client, task, control);
+        }
+    }
+
+    public ConfigDiff resume(TaskSource taskSource,
+            Schema schema, int processorCount,
+            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)
+    {
+        boolean doUpload = startBulkImportSession(client, task.getSessionName(), task.getDatabase(), task.getTable());
+        task.setDoUpload(doUpload);
+        control.run(task.dump());
+        completeBulkImportSession(client, task.getSessionName(), 0);  // TODO perform job priority
+
+        ConfigDiff configDiff = Exec.newConfigDiff();
+        configDiff.set("last_session", task.getSessionName());
+        return configDiff;
+    }
+
+    public void cleanup(TaskSource taskSource,
+            Schema schema, int processorCount,
+            List<CommitReport> successCommitReports)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+        try (TdApiClient client = newTdApiClient(task)) {
+            String sessionName = task.getSessionName();
+            log.info("Deleting bulk import session '{}'", sessionName);
+            client.deleteBulkImportSession(sessionName);
+        }
+    }
+
+    private 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) {
+            throw Throwables.propagate(e);
+        }
+        return client;
+    }
+
+    private Optional<HttpProxyConfig> newHttpProxyConfig(Optional<HttpProxyTask> task)
+    {
+        Optional<HttpProxyConfig> httpProxyConfig;
+        if (task.isPresent()) {
+            HttpProxyTask pt = task.get();
+            httpProxyConfig = Optional.of(new HttpProxyConfig(pt.getHost(), pt.getPort(), pt.getUseSsl()));
+        } else {
+            httpProxyConfig = Optional.absent();
+        }
+        return httpProxyConfig;
+    }
+
+    private 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) {
+            try {
+                client.createDatabase(databaseName);
+                log.debug("Created database \"{}\"", databaseName);
+            } catch (TdApiConflictException ex) {
+                // ignorable error
+            }
+            try {
+                client.createTable(databaseName, tableName);
+                log.debug("Created table \"{}\".\"{}\"", databaseName, tableName);
+            } catch (TdApiConflictException exe) {
+                // ignorable error
+            }
+        } catch (TdApiConflictException e) {
+            // ignorable error
+        }
+    }
+
+    private void validateTableExists(TdApiClient client, String databaseName, String tableName)
+    {
+        try {
+            for (TDTable table : client.getTables(databaseName)) {
+                if (table.getName().equals(tableName)) {
+                    return;
+                }
+            }
+            throw new ConfigException(String.format("Table \"%s\".\"%s\" doesn't exist", databaseName, tableName));
+        } catch (TdApiNotFoundException ex) {
+            throw new ConfigException(String.format("Database \"%s\" doesn't exist", databaseName), ex);
+        }
+    }
+
+    private String buildBulkImportSessionName(PluginTask task, ExecSession exec)
+    {
+        if (task.getSession().isPresent()) {
+            return task.getSession().get();
+        } else {
+            Timestamp time = exec.getTransactionTime(); // TODO implement Exec.getTransactionUniqueName()
+            return String.format("embulk_%s_%09d",
+                    DateTimeFormat.forPattern("yyyyMMdd_HHmmss").withZoneUTC().print(time.getEpochSecond() * 1000),
+                    time.getNano());
+        }
+    }
+
+    // return false if all files are already uploaded
+    private 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) {
+            // ignorable error
+        }
+        session = client.getBulkImportSession(sessionName);
+        // TODO check associated databaseName and tableName
+
+        switch (session.getStatus()) {
+        case UPLOADING:
+            if (session.getUploadFrozen()) {
+                return false;
+            }
+            return true;
+        case PERFORMING:
+            return false;
+        case READY:
+            return false;
+        case COMMITTING:
+            return false;
+        case COMMITTED:
+            return false;
+        case UNKNOWN:
+        default:
+            throw new RuntimeException("Unknown bulk import status");
+        }
+    }
+
+    private void completeBulkImportSession(TdApiClient client, String sessionName, int priority)
+    {
+        TDBulkImportSession session = client.getBulkImportSession(sessionName);
+
+        switch (session.getStatus()) {
+        case UPLOADING:
+            if (!session.getUploadFrozen()) {
+                // freeze
+                try {
+                    client.freezeBulkImportSession(sessionName);
+                } catch (TdApiConflictException e) {
+                    // ignorable error
+                }
+            }
+            // perform
+            client.performBulkImportSession(sessionName, priority);
+
+            // pass
+        case PERFORMING:
+            log.info("Performing bulk import session '{}'", sessionName);
+            session = waitForStatusChange(client, sessionName,
+                    ImportStatus.PERFORMING, ImportStatus.READY,
+                    "perform");
+            log.info("    job id: {}", session.getJobId());
+
+            // pass
+        case READY:
+            // TODO add an option to make the transaction failed if error_records or error_parts is too large
+            // commit
+            log.info("Committing bulk import session '{}'", sessionName);
+            log.info("    valid records: {}", session.getValidRecords());
+            log.info("    error records: {}", session.getErrorRecords());
+            log.info("    valid parts: {}", session.getValidParts());
+            log.info("    error parts: {}", session.getErrorParts());
+            client.commitBulkImportSession(sessionName);
+
+            // pass
+        case COMMITTING:
+            session = waitForStatusChange(client, sessionName,
+                    ImportStatus.COMMITTING, ImportStatus.COMMITTED,
+                    "commit");
+
+            // pass
+        case COMMITTED:
+            return;
+
+        case UNKNOWN:
+            throw new RuntimeException("Unknown bulk import status");
+        }
+    }
+
+    private TDBulkImportSession waitForStatusChange(TdApiClient client, String sessionName,
+            ImportStatus current, ImportStatus expecting, String operation)
+    {
+        TDBulkImportSession importSession;
+        while (true) {
+            importSession = client.getBulkImportSession(sessionName);
+
+            if (importSession.is(expecting)) {
+                return importSession;
+
+            } else if (importSession.is(current)) {
+                // in progress
+
+            } else {
+                throw new RuntimeException(String.format("Failed to %s bulk import session '%s'",
+                            operation, sessionName));
+            }
+
+            try {
+                Thread.sleep(3000);
+            } catch (InterruptedException e) {
+            }
+        }
+    }
+
+    @Override
+    public TransactionalPageOutput open(TaskSource taskSource, Schema schema, int processorIndex)
+    {
+        final PluginTask task = taskSource.loadTask(PluginTask.class);
+
+        RecordWriter closeLater = null;
+        try {
+            FieldWriterSet fieldWriters = new FieldWriterSet(log, task, schema);
+            RecordWriter recordWriter = closeLater = new RecordWriter(task, newTdApiClient(task), fieldWriters);
+            recordWriter.open(schema);
+            closeLater = null;
+            return recordWriter;
+
+        } catch (IOException e) {
+            throw Throwables.propagate(e);
+        } finally {
+            if (closeLater != null) {
+                closeLater.close();
+            }
+        }
+    }
+}

data/src/test/java/org/embulk/output/TestTdOutputPlugin.java

@@ -0,0 +1,5 @@
+package org.embulk.output;
+
+public class TestTdOutputPlugin
+{
+}

metadata

@@ -0,0 +1,119 @@
+--- !ruby/object:Gem::Specification
+name: embulk-output-td
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Muga Nishizawa
+autorequire:
+bindir: bin
+cert_chain: []
+date: 2015-06-22 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: bundler
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ~>
+      - !ruby/object:Gem::Version
+        version: '1.0'
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ~>
+      - !ruby/object:Gem::Version
+        version: '1.0'
+  prerelease: false
+  type: :development
+- !ruby/object:Gem::Dependency
+  name: rake
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - '>='
+      - !ruby/object:Gem::Version
+        version: '10.0'
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - '>='
+      - !ruby/object:Gem::Version
+        version: '10.0'
+  prerelease: false
+  type: :development
+description: TreasureData output plugin is an Embulk plugin that loads records to TreasureData read by any input plugins. Search the input plugins by 'embulk-output' keyword.
+email:
+- muga.nishizawa@gmail.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- .gitignore
+- README.md
+- build.gradle
+- embulk-output-td.gemspec
+- gradle/wrapper/gradle-wrapper.jar
+- gradle/wrapper/gradle-wrapper.properties
+- gradlew
+- gradlew.bat
+- lib/embulk/output/td.rb
+- settings.gradle
+- src/main/java/com/treasuredata/api/TdApiClient.java
+- src/main/java/com/treasuredata/api/TdApiClientConfig.java
+- src/main/java/com/treasuredata/api/TdApiConflictException.java
+- src/main/java/com/treasuredata/api/TdApiConstants.java
+- src/main/java/com/treasuredata/api/TdApiException.java
+- src/main/java/com/treasuredata/api/TdApiExecutionException.java
+- src/main/java/com/treasuredata/api/TdApiExecutionInterruptedException.java
+- src/main/java/com/treasuredata/api/TdApiExecutionTimeoutException.java
+- src/main/java/com/treasuredata/api/TdApiNotFoundException.java
+- src/main/java/com/treasuredata/api/TdApiResponseException.java
+- src/main/java/com/treasuredata/api/model/TDArrayColumnType.java
+- src/main/java/com/treasuredata/api/model/TDBulkImportSession.java
+- src/main/java/com/treasuredata/api/model/TDColumn.java
+- src/main/java/com/treasuredata/api/model/TDColumnType.java
+- src/main/java/com/treasuredata/api/model/TDColumnTypeDeserializer.java
+- src/main/java/com/treasuredata/api/model/TDDatabase.java
+- src/main/java/com/treasuredata/api/model/TDDatabaseList.java
+- src/main/java/com/treasuredata/api/model/TDMapColumnType.java
+- src/main/java/com/treasuredata/api/model/TDPrimitiveColumnType.java
+- src/main/java/com/treasuredata/api/model/TDTable.java
+- src/main/java/com/treasuredata/api/model/TDTableList.java
+- src/main/java/com/treasuredata/api/model/TDTablePermission.java
+- src/main/java/com/treasuredata/api/model/TDTableSchema.java
+- src/main/java/com/treasuredata/api/model/TDTableType.java
+- src/main/java/org/embulk/output/FinalizableExecutorService.java
+- src/main/java/org/embulk/output/MsgpackGZFileBuilder.java
+- src/main/java/org/embulk/output/RecordWriter.java
+- src/main/java/org/embulk/output/TdOutputPlugin.java
+- src/test/java/org/embulk/output/TestTdOutputPlugin.java
+- classpath/embulk-output-td-0.1.0.jar
+- classpath/javassist-3.18.1-GA.jar
+- classpath/jetty-client-9.2.2.v20140723.jar
+- classpath/jetty-http-9.2.2.v20140723.jar
+- classpath/jetty-io-9.2.2.v20140723.jar
+- classpath/jetty-util-9.2.2.v20140723.jar
+- classpath/json-simple-1.1.1.jar
+- classpath/msgpack-0.6.11.jar
+homepage: https://github.com/treasure-data/embulk-output-td
+licenses:
+- Apache 2.0
+metadata: {}
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - '>='
+    - !ruby/object:Gem::Version
+      version: '0'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - '>='
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubyforge_project:
+rubygems_version: 2.1.9
+signing_key:
+specification_version: 4
+summary: TreasureData output plugin for Embulk
+test_files: []