embulk-output-elasticsearch 0.3.1 → 0.4.0

data/src/main/java/org/embulk/output/elasticsearch/ElasticsearchOutputPluginDelegate.java

@@ -0,0 +1,286 @@
+package org.embulk.output.elasticsearch;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonValue;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Optional;
+import com.google.common.base.Throwables;
+import org.eclipse.jetty.util.ssl.SslContextFactory;
+import org.embulk.base.restclient.RestClientOutputPluginDelegate;
+import org.embulk.base.restclient.RestClientOutputTaskBase;
+import org.embulk.base.restclient.jackson.JacksonServiceRequestMapper;
+import org.embulk.base.restclient.jackson.JacksonTopLevelValueLocator;
+import org.embulk.base.restclient.jackson.scope.JacksonAllInObjectScope;
+import org.embulk.base.restclient.record.RecordBuffer;
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.ConfigDiff;
+import org.embulk.config.ConfigException;
+import org.embulk.config.Task;
+import org.embulk.config.TaskReport;
+import org.embulk.spi.Exec;
+import org.embulk.spi.Schema;
+import org.embulk.spi.time.TimestampFormatter;
+import org.embulk.util.retryhelper.jetty92.Jetty92ClientCreator;
+import org.embulk.util.retryhelper.jetty92.Jetty92RetryHelper;
+import org.joda.time.DateTimeZone;
+import org.slf4j.Logger;
+
+import java.util.List;
+import java.util.Locale;
+
+public class ElasticsearchOutputPluginDelegate
+        implements RestClientOutputPluginDelegate<ElasticsearchOutputPluginDelegate.PluginTask>
+{
+    private final Logger log;
+    private final ElasticsearchHttpClient client;
+
+    public ElasticsearchOutputPluginDelegate()
+    {
+        this.log = Exec.getLogger(getClass());
+        this.client = new ElasticsearchHttpClient();
+    }
+
+    public interface NodeAddressTask
+            extends Task
+    {
+        @Config("host")
+        String getHost();
+
+        @Config("port")
+        @ConfigDefault("9200")
+        int getPort();
+    }
+
+    public interface PluginTask
+            extends RestClientOutputTaskBase, TimestampFormatter.Task
+    {
+        @Config("mode")
+        @ConfigDefault("\"insert\"")
+        Mode getMode();
+
+        @Config("nodes")
+        List<NodeAddressTask> getNodes();
+
+        @Config("cluster_name")
+        @ConfigDefault("\"elasticsearch\"")
+        String getClusterName();
+
+        @Config("index")
+        String getIndex();
+        void setIndex(String indexName);
+
+        @Config("alias")
+        @ConfigDefault("null")
+        Optional<String> getAlias();
+        void setAlias(Optional<String> aliasName);
+
+        @Config("index_type")
+        String getType();
+
+        @Config("id")
+        @ConfigDefault("null")
+        Optional<String> getId();
+
+        @Config("use_ssl")
+        @ConfigDefault("false")
+        boolean getUseSsl();
+
+        @Config("auth_method")
+        @ConfigDefault("\"none\"")
+        AuthMethod getAuthMethod();
+
+        @Config("user")
+        @ConfigDefault("null")
+        Optional<String> getUser();
+
+        @Config("password")
+        @ConfigDefault("null")
+        Optional<String> getPassword();
+
+        @Config("bulk_actions")
+        @ConfigDefault("1000")
+        int getBulkActions();
+
+        @Config("bulk_size")
+        @ConfigDefault("5242880")
+        long getBulkSize();
+
+        @Config("concurrent_requests")
+        @ConfigDefault("5")
+        int getConcurrentRequests();
+
+        @Config("maximum_retries")
+        @ConfigDefault("7")
+        int getMaximumRetries();
+
+        @Config("initial_retry_interval_millis")
+        @ConfigDefault("1000")
+        int getInitialRetryIntervalMillis();
+
+        @Config("maximum_retry_interval_millis")
+        @ConfigDefault("120000")
+        int getMaximumRetryIntervalMillis();
+
+        @Config("timeout_millis")
+        @ConfigDefault("60000")
+        int getTimeoutMills();
+
+        @Config("time_zone")
+        @ConfigDefault("\"UTC\"")
+        String getTimeZone();
+    }
+
+    public enum Mode
+    {
+        INSERT,
+        REPLACE;
+
+        @JsonValue
+        @Override
+        public String toString()
+        {
+            return name().toLowerCase(Locale.ENGLISH);
+        }
+
+        @JsonCreator
+        public static Mode fromString(String value)
+        {
+            switch (value) {
+                case "insert":
+                    return INSERT;
+                case "replace":
+                    return REPLACE;
+                default:
+                    throw new ConfigException(String.format("Unknown mode '%s'. Supported modes are insert, truncate_insert, replace", value));
+            }
+        }
+    }
+
+    public enum AuthMethod
+    {
+        NONE,
+        BASIC;
+
+        @JsonValue
+        @Override
+        public String toString()
+        {
+            return name().toLowerCase(Locale.ENGLISH);
+        }
+
+        @JsonCreator
+        public static AuthMethod fromString(String value)
+        {
+            switch (value) {
+                case "none":
+                    return NONE;
+                case "basic":
+                    return BASIC;
+                default:
+                    throw new ConfigException(String.format("Unknown auth_method '%s'. Supported auth_method are none, basic", value));
+            }
+        }
+    }
+
+    @Override  // Overridden from |OutputTaskValidatable|
+    public void validateOutputTask(PluginTask task, Schema embulkSchema, int taskCount)
+    {
+        if (task.getNodes().size() > 0) {
+            for (NodeAddressTask node : task.getNodes()) {
+                if (node.getHost().endsWith("es.amazonaws.com")) {
+                    throw new ConfigException("This plugin does't support AWS Elasticsearch Service.");
+                }
+                if (node.getPort() == 9300) {
+                    log.warn("Port:9300 is usually used by TransportClient. HTTP/Rest Client uses 9200.");
+                }
+            }
+        }
+
+        try (Jetty92RetryHelper retryHelper = createRetryHelper(task)) {
+            log.info(String.format("Connecting to Elasticsearch version:%s", client.getEsVersion(task, retryHelper)));
+            log.info("Executing plugin with '{}' mode.", task.getMode());
+            client.validateIndexOrAliasName(task.getIndex(), "index");
+            client.validateIndexOrAliasName(task.getType(), "index_type");
+
+            if (task.getMode().equals(Mode.REPLACE)) {
+                task.setAlias(Optional.of(task.getIndex()));
+                task.setIndex(client.generateNewIndexName(task.getIndex()));
+                if (client.isIndexExisting(task.getAlias().orNull(), task, retryHelper) && !client.isAliasExisting(task.getAlias().orNull(), task, retryHelper)) {
+                    throw new ConfigException(String.format("Invalid alias name [%s], an index exists with the same name as the alias", task.getAlias().orNull()));
+                }
+            }
+            log.info(String.format("Inserting data into index[%s]", task.getIndex()));
+        }
+
+        if (task.getAuthMethod() == AuthMethod.BASIC) {
+            if (!task.getUser().isPresent() || !task.getPassword().isPresent()) {
+                throw new ConfigException("'user' and 'password' are required when auth_method='basic'");
+            }
+        }
+    }
+
+    @Override  // Overridden from |ServiceRequestMapperBuildable|
+    public JacksonServiceRequestMapper buildServiceRequestMapper(PluginTask task)
+    {
+        TimestampFormatter formatter = new TimestampFormatter(task.getJRuby(), "%Y-%m-%dT%H:%M:%S.%3N%z", DateTimeZone.forID(task.getTimeZone()));
+
+        return JacksonServiceRequestMapper.builder()
+                .add(new JacksonAllInObjectScope(formatter), new JacksonTopLevelValueLocator("record"))
+                .build();
+    }
+
+    @Override  // Overridden from |RecordBufferBuildable|
+    public RecordBuffer buildRecordBuffer(PluginTask task)
+    {
+        Jetty92RetryHelper retryHelper = createRetryHelper(task);
+        return new ElasticsearchRecordBuffer("records", task, retryHelper);
+    }
+
+    @Override
+    public ConfigDiff egestEmbulkData(final PluginTask task,
+                                      Schema schema,
+                                      int taskIndex,
+                                      List<TaskReport> taskReports)
+    {
+        long totalInserted = 0;
+        for (TaskReport taskReport : taskReports) {
+            if (taskReport.has("inserted")) {
+                totalInserted += taskReport.get(Long.class, "inserted");
+            }
+        }
+
+        log.info("Insert completed. {} records", totalInserted);
+        try (Jetty92RetryHelper retryHelper = createRetryHelper(task)) {
+            // Re assign alias only when repale mode
+            if (task.getMode().equals(Mode.REPLACE)) {
+                client.reassignAlias(task.getAlias().orNull(), task.getIndex(), task, retryHelper);
+            }
+        }
+
+        return Exec.newConfigDiff();
+    }
+
+    @VisibleForTesting
+    protected Jetty92RetryHelper createRetryHelper(PluginTask task)
+    {
+        return new Jetty92RetryHelper(
+                task.getMaximumRetries(),
+                task.getInitialRetryIntervalMillis(),
+                task.getMaximumRetryIntervalMillis(),
+                new Jetty92ClientCreator() {
+                    @Override
+                    public org.eclipse.jetty.client.HttpClient createAndStart()
+                    {
+                        org.eclipse.jetty.client.HttpClient client = new org.eclipse.jetty.client.HttpClient(new SslContextFactory());
+                        try {
+                            client.start();
+                            return client;
+                        }
+                        catch (Exception e) {
+                            throw Throwables.propagate(e);
+                        }
+                    }
+                });
+    }
+}

data/src/main/java/org/embulk/output/elasticsearch/ElasticsearchRecordBuffer.java

@@ -0,0 +1,99 @@
+package org.embulk.output.elasticsearch;
+
+import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.JsonNodeFactory;
+import com.google.common.base.Throwables;
+import org.embulk.base.restclient.jackson.JacksonServiceRecord;
+import org.embulk.base.restclient.record.RecordBuffer;
+import org.embulk.base.restclient.record.ServiceRecord;
+import org.embulk.config.TaskReport;
+import org.embulk.output.elasticsearch.ElasticsearchOutputPluginDelegate.PluginTask;
+import org.embulk.spi.Exec;
+import org.embulk.util.retryhelper.jetty92.Jetty92RetryHelper;
+import org.slf4j.Logger;
+
+import java.io.IOException;
+
+/**
+ * ElasticsearchRecordBuffer is an implementation of {@code RecordBuffer} which includes JSON output directly to Elasticsearch server.
+ */
+public class ElasticsearchRecordBuffer
+        extends RecordBuffer
+{
+    private final String attributeName;
+    private final PluginTask task;
+    private final long bulkActions;
+    private final long bulkSize;
+    private final ElasticsearchHttpClient client;
+    private final Jetty92RetryHelper retryHelper;
+    private final ObjectMapper mapper;
+    private final Logger log;
+    private long totalCount;
+    private int requestCount;
+    private long requestBytes;
+    private ArrayNode records;
+
+    public ElasticsearchRecordBuffer(String attributeName, PluginTask task, Jetty92RetryHelper retryHelper)
+    {
+        this.attributeName = attributeName;
+        this.task = task;
+        this.bulkActions = task.getBulkActions();
+        this.bulkSize = task.getBulkSize();
+        this.client = new ElasticsearchHttpClient();
+        this.retryHelper = retryHelper;
+        this.mapper = new ObjectMapper()
+                .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)
+                .configure(com.fasterxml.jackson.core.JsonParser.Feature.ALLOW_UNQUOTED_CONTROL_CHARS, false);
+        this.records = JsonNodeFactory.instance.arrayNode();
+        this.totalCount = 0;
+        this.requestCount = 0;
+        this.requestBytes = 0;
+        this.log = Exec.getLogger(getClass());
+    }
+
+    @Override
+    public void bufferRecord(ServiceRecord serviceRecord)
+    {
+        JacksonServiceRecord jacksonServiceRecord;
+        try {
+            jacksonServiceRecord = (JacksonServiceRecord) serviceRecord;
+            JsonNode record = mapper.readTree(jacksonServiceRecord.toString()).get("record");
+
+            requestCount++;
+            totalCount++;
+            requestBytes += record.toString().getBytes().length;
+
+            records.add(record);
+            if (requestCount >= bulkActions || requestBytes >= bulkSize) {
+                client.push(records, task, retryHelper);
+                if (totalCount % 10000 == 0) {
+                    log.info("Inserted {} records", totalCount);
+                }
+                records = JsonNodeFactory.instance.arrayNode();
+                requestBytes = 0;
+                requestCount = 0;
+            }
+        }
+        catch (ClassCastException ex) {
+            throw new RuntimeException(ex);
+        }
+        catch (IOException ex) {
+            throw Throwables.propagate(ex);
+        }
+    }
+
+    @Override
+    public TaskReport commitWithTaskReportUpdated(TaskReport taskReport)
+    {
+        if (records.size() > 0) {
+            client.push(records, task, retryHelper);
+            log.info("Inserted {} records", records.size());
+        }
+
+        this.retryHelper.close();
+        return Exec.newTaskReport().set("inserted", totalCount);
+    }
+}

data/src/test/java/org/embulk/output/elasticsearch/ElasticsearchTestUtils.java

@@ -0,0 +1,159 @@
+package org.embulk.output.elasticsearch;
+
+import com.google.common.base.Throwables;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import org.eclipse.jetty.util.ssl.SslContextFactory;
+import org.embulk.config.ConfigSource;
+import org.embulk.output.elasticsearch.ElasticsearchOutputPluginDelegate.PluginTask;
+import org.embulk.spi.Exec;
+import org.embulk.util.retryhelper.jetty92.Jetty92ClientCreator;
+import org.embulk.util.retryhelper.jetty92.Jetty92RetryHelper;
+
+import java.lang.reflect.Method;
+import java.util.Arrays;
+import java.util.List;
+
+import static org.junit.Assume.assumeNotNull;
+
+public class ElasticsearchTestUtils
+{
+    public static String ES_HOST;
+    public static int ES_PORT;
+    public static List ES_NODES;
+    public static String ES_INDEX;
+    public static String ES_INDEX_TYPE;
+    public static String ES_ID;
+    public static int ES_BULK_ACTIONS;
+    public static int ES_BULK_SIZE;
+    public static int ES_CONCURRENT_REQUESTS;
+    public static String PATH_PREFIX;
+    public static String ES_INDEX2;
+    public static String ES_ALIAS;
+
+    /*
+     * This test case requires environment variables
+     *   ES_HOST
+     *   ES_INDEX
+     *   ES_INDEX_TYPE
+     */
+    public void initializeConstant()
+    {
+        ES_HOST = System.getenv("ES_HOST") != null ? System.getenv("ES_HOST") : "";
+        ES_PORT = System.getenv("ES_PORT") != null ? Integer.valueOf(System.getenv("ES_PORT")) : 9200;
+
+        ES_INDEX = System.getenv("ES_INDEX");
+        ES_INDEX2 = ES_INDEX + "_02";
+        ES_ALIAS = ES_INDEX + "_alias";
+        ES_INDEX_TYPE = System.getenv("ES_INDEX_TYPE");
+        ES_ID = "id";
+        ES_BULK_ACTIONS = System.getenv("ES_BULK_ACTIONS") != null ? Integer.valueOf(System.getenv("ES_BULK_ACTIONS")) : 1000;
+        ES_BULK_SIZE = System.getenv("ES_BULK_SIZE") != null ? Integer.valueOf(System.getenv("ES_BULK_SIZE")) : 5242880;
+        ES_CONCURRENT_REQUESTS = System.getenv("ES_CONCURRENT_REQUESTS") != null ? Integer.valueOf(System.getenv("ES_CONCURRENT_REQUESTS")) : 5;
+
+        assumeNotNull(ES_HOST, ES_INDEX, ES_INDEX_TYPE);
+
+        ES_NODES = Arrays.asList(ImmutableMap.of("host", ES_HOST, "port", ES_PORT));
+
+        PATH_PREFIX = ElasticsearchTestUtils.class.getClassLoader().getResource("sample_01.csv").getPath();
+    }
+
+    public void prepareBeforeTest(PluginTask task) throws Exception
+    {
+        ElasticsearchHttpClient client = new ElasticsearchHttpClient();
+        try (Jetty92RetryHelper retryHelper = createRetryHelper()) {
+            Method deleteIndex = ElasticsearchHttpClient.class.getDeclaredMethod("deleteIndex", String.class, PluginTask.class, Jetty92RetryHelper.class);
+            deleteIndex.setAccessible(true);
+
+            // Delete alias
+            if (client.isAliasExisting(ES_ALIAS, task, retryHelper)) {
+                deleteIndex.invoke(client, ES_ALIAS, task, retryHelper);
+            }
+
+            // Delete index
+            if (client.isIndexExisting(ES_INDEX, task, retryHelper)) {
+                deleteIndex.invoke(client, ES_INDEX, task, retryHelper);
+            }
+
+            if (client.isIndexExisting(ES_INDEX2, task, retryHelper)) {
+                deleteIndex.invoke(client, ES_INDEX2, task, retryHelper);
+            }
+        }
+    }
+
+    public ConfigSource config()
+    {
+        return Exec.newConfigSource()
+                .set("in", inputConfig())
+                .set("parser", parserConfig(schemaConfig()))
+                .set("type", "elasticsearch")
+                .set("mode", "insert")
+                .set("nodes", ES_NODES)
+                .set("index", ES_INDEX)
+                .set("index_type", ES_INDEX_TYPE)
+                .set("id", ES_ID)
+                .set("bulk_actions", ES_BULK_ACTIONS)
+                .set("bulk_size", ES_BULK_SIZE)
+                .set("concurrent_requests", ES_CONCURRENT_REQUESTS);
+    }
+
+    public ImmutableMap<String, Object> inputConfig()
+    {
+        ImmutableMap.Builder<String, Object> builder = new ImmutableMap.Builder<>();
+        builder.put("type", "file");
+        builder.put("path_prefix", PATH_PREFIX);
+        builder.put("last_path", "");
+        return builder.build();
+    }
+
+    public ImmutableMap<String, Object> parserConfig(ImmutableList<Object> schemaConfig)
+    {
+        ImmutableMap.Builder<String, Object> builder = new ImmutableMap.Builder<>();
+        builder.put("type", "csv");
+        builder.put("newline", "CRLF");
+        builder.put("delimiter", ",");
+        builder.put("quote", "\"");
+        builder.put("escape", "\"");
+        builder.put("trim_if_not_quoted", false);
+        builder.put("skip_header_lines", 1);
+        builder.put("allow_extra_columns", false);
+        builder.put("allow_optional_columns", false);
+        builder.put("columns", schemaConfig);
+        return builder.build();
+    }
+
+    public ImmutableList<Object> schemaConfig()
+    {
+        ImmutableList.Builder<Object> builder = new ImmutableList.Builder<>();
+        builder.add(ImmutableMap.of("name", "id", "type", "long"));
+        builder.add(ImmutableMap.of("name", "account", "type", "long"));
+        builder.add(ImmutableMap.of("name", "time", "type", "timestamp", "format", "%Y-%m-%d %H:%M:%S"));
+        builder.add(ImmutableMap.of("name", "purchase", "type", "timestamp", "format", "%Y%m%d"));
+        builder.add(ImmutableMap.of("name", "flg", "type", "boolean"));
+        builder.add(ImmutableMap.of("name", "score", "type", "double"));
+        builder.add(ImmutableMap.of("name", "comment", "type", "string"));
+        return builder.build();
+    }
+
+    public Jetty92RetryHelper createRetryHelper()
+    {
+        return new Jetty92RetryHelper(
+                2,
+                1000,
+                32000,
+                new Jetty92ClientCreator() {
+                    @Override
+                    public org.eclipse.jetty.client.HttpClient createAndStart()
+                    {
+                        org.eclipse.jetty.client.HttpClient client = new org.eclipse.jetty.client.HttpClient(new SslContextFactory());
+                        try {
+                            client.start();
+                            return client;
+                        }
+                        catch (Exception e) {
+                            throw Throwables.propagate(e);
+                        }
+                    }
+                });
+    }
+}