Merge pull request #286 from wicknicks/merge-275-into-5.0.x

Merge 275 into 5.0.x

Author: wicknicks

src/main/java/io/confluent/connect/elasticsearch/ElasticsearchSinkConnectorConfig.java

@@ -149,6 +149,11 @@ public class ElasticsearchSinkConnectorConfig extends AbstractConfig {
       + " mapping conflict or a field name containing illegal characters. Valid options are "
       + "'ignore', 'warn', and 'fail'.";
 
+  public static final String AUTO_CREATE_INDICES_AT_START_CONFIG = "auto.create.indices.at.start";
+  private static final String AUTO_CREATE_INDICES_AT_START_DOC = "Auto create the Elasticsearch"
+      + " indices at startup. This is useful when the indices are a direct mapping "
+      + " of the Kafka topics.";
+
   protected static ConfigDef baseConfigDef() {
     final ConfigDef configDef = new ConfigDef();
     addConnectorConfigs(configDef);
@@ -274,10 +279,21 @@ private static void addConnectorConfigs(ConfigDef configDef) {
         3000, 
         Importance.LOW, 
         READ_TIMEOUT_MS_CONFIG_DOC,
-        group, 
-        ++order, 
-        Width.SHORT, 
-        "Read Timeout");
+        group,
+        ++order,
+        Width.SHORT,
+        "Read Timeout"
+    ).define(
+        AUTO_CREATE_INDICES_AT_START_CONFIG,
+        Type.BOOLEAN,
+        true,
+        Importance.LOW,
+        AUTO_CREATE_INDICES_AT_START_DOC,
+        group,
+        ++order,
+        Width.SHORT,
+        "Create indices at startup"
+      );
   }
 
   private static void addConversionConfigs(ConfigDef configDef) {

src/main/java/io/confluent/connect/elasticsearch/ElasticsearchSinkTask.java

@@ -40,6 +40,7 @@ public class ElasticsearchSinkTask extends SinkTask {
   private static final Logger log = LoggerFactory.getLogger(ElasticsearchSinkTask.class);
   private ElasticsearchWriter writer;
   private ElasticsearchClient client;
+  private Boolean createIndicesAtStartTime;
 
   @Override
   public String version() {
@@ -91,6 +92,8 @@ public void start(Map<String, String> props, ElasticsearchClient client) {
           config.getInt(ElasticsearchSinkConnectorConfig.MAX_RETRIES_CONFIG);
       boolean dropInvalidMessage =
           config.getBoolean(ElasticsearchSinkConnectorConfig.DROP_INVALID_MESSAGE_CONFIG);
+      boolean createIndicesAtStartTime =
+          config.getBoolean(ElasticsearchSinkConnectorConfig.AUTO_CREATE_INDICES_AT_START_CONFIG);
 
       DataConverter.BehaviorOnNullValues behaviorOnNullValues =
           DataConverter.BehaviorOnNullValues.forValue(
@@ -137,6 +140,8 @@ public void start(Map<String, String> props, ElasticsearchClient client) {
           .setBehaviorOnNullValues(behaviorOnNullValues)
           .setBehaviorOnMalformedDoc(behaviorOnMalformedDoc);
 
+      this.createIndicesAtStartTime = createIndicesAtStartTime;
+
       writer = builder.build();
       writer.start();
     } catch (ConfigException e) {
@@ -150,11 +155,13 @@ public void start(Map<String, String> props, ElasticsearchClient client) {
   @Override
   public void open(Collection<TopicPartition> partitions) {
     log.debug("Opening the task for topic partitions: {}", partitions);
-    Set<String> topics = new HashSet<>();
-    for (TopicPartition tp : partitions) {
-      topics.add(tp.topic());
+    if (createIndicesAtStartTime) {
+      Set<String> topics = new HashSet<>();
+      for (TopicPartition tp : partitions) {
+        topics.add(tp.topic());
+      }
+      writer.createIndicesForTopics(topics);
     }
-    writer.createIndicesForTopics(topics);
   }
 
   @Override

src/main/java/io/confluent/connect/elasticsearch/ElasticsearchWriter.java

@@ -253,6 +253,8 @@ public void write(Collection<SinkRecord> records) {
       final boolean ignoreSchema =
           ignoreSchemaTopics.contains(sinkRecord.topic()) || this.ignoreSchema;
 
+      client.createIndices(Collections.singleton(index));
+
       if (!ignoreSchema && !existingMappings.contains(index)) {
         try {
           if (Mapping.getMapping(client, index, type) == null) {

src/main/java/io/confluent/connect/elasticsearch/jest/JestElasticsearchClient.java

@@ -20,12 +20,12 @@
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.google.gson.JsonObject;
-import io.confluent.connect.elasticsearch.bulk.BulkRequest;
 import io.confluent.connect.elasticsearch.ElasticsearchClient;
 import io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig;
 import io.confluent.connect.elasticsearch.IndexableRecord;
 import io.confluent.connect.elasticsearch.Key;
 import io.confluent.connect.elasticsearch.Mapping;
+import io.confluent.connect.elasticsearch.bulk.BulkRequest;
 import io.confluent.connect.elasticsearch.bulk.BulkResponse;
 import io.searchbox.action.Action;
 import io.searchbox.action.BulkableAction;
@@ -54,6 +54,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -74,6 +75,8 @@ public class JestElasticsearchClient implements ElasticsearchClient {
   private final JestClient client;
   private final Version version;
 
+  private final Set<String> indexCache = new HashSet<>();
+
   // visible for testing
   public JestElasticsearchClient(JestClient client) {
     try {
@@ -216,7 +219,10 @@ public Version getVersion() {
   }
 
   private boolean indexExists(String index) {
-    Action<JestResult> action = new IndicesExists.Builder(index).build();
+    if (indexCache.contains(index)) {
+      return true;
+    }
+    Action<?> action = new IndicesExists.Builder(index).build();
     try {
       JestResult result = client.execute(action);
       return result.isSucceeded();
@@ -238,6 +244,7 @@ public void createIndices(Set<String> indices) {
               throw new ConnectException("Could not create index '" + index + "'" + msg);
             }
           }
+          indexCache.add(index);
         } catch (IOException e) {
           throw new ConnectException(e);
         }

src/test/java/io/confluent/connect/elasticsearch/ElasticsearchSinkTaskTest.java

@@ -22,6 +22,12 @@
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.sink.SinkRecord;
+import org.elasticsearch.action.ActionFuture;
+import org.elasticsearch.action.ActionRequest;
+import org.elasticsearch.action.ActionRequestValidationException;
+import org.elasticsearch.action.admin.indices.exists.indices.IndicesExistsRequest;
+import org.elasticsearch.action.admin.indices.exists.indices.IndicesExistsResponse;
+import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.test.InternalTestCluster;
 import org.junit.Test;
@@ -42,6 +48,11 @@ public class ElasticsearchSinkTaskTest extends ElasticsearchSinkTestBase {
   private static final int PARTITION_113 = 113;
   private static final TopicPartition TOPIC_IN_CAPS_PARTITION = new TopicPartition(TOPIC_IN_CAPS, PARTITION_113);
 
+  private static final String UNSEEN_TOPIC = "UnseenTopic";
+  private static final int PARTITION_114 = 114;
+  private static final TopicPartition UNSEEN_TOPIC_PARTITION = new TopicPartition(UNSEEN_TOPIC, PARTITION_114);
+
+
   private Map<String, String> createProps() {
     Map<String, String> props = new HashMap<>();
     props.put(ElasticsearchSinkConnectorConfig.TYPE_NAME_CONFIG, TYPE);
@@ -94,12 +105,12 @@ public void testCreateAndWriteToIndexForTopicWithUppercaseCharacters() {
     Struct record = createRecord(schema);
 
     SinkRecord sinkRecord = new SinkRecord(TOPIC_IN_CAPS,
-            PARTITION_113,
-            Schema.STRING_SCHEMA,
-            key,
-            schema,
-            record,
-            0 );
+        PARTITION_113,
+        Schema.STRING_SCHEMA,
+        key,
+        schema,
+        record,
+        0 );
 
     try {
       task.start(props, client);
@@ -111,4 +122,46 @@ public void testCreateAndWriteToIndexForTopicWithUppercaseCharacters() {
       task.stop();
     }
   }
+
+  @Test
+  public void testCreateAndWriteToIndexNotCreatedAtStartTime() {
+    InternalTestCluster cluster = ESIntegTestCase.internalCluster();
+    cluster.ensureAtLeastNumDataNodes(3);
+    Map<String, String> props = createProps();
+
+    props.put(ElasticsearchSinkConnectorConfig.AUTO_CREATE_INDICES_AT_START_CONFIG, "false");
+
+    ElasticsearchSinkTask task = new ElasticsearchSinkTask();
+
+    String key = "key";
+    Schema schema = createSchema();
+    Struct record = createRecord(schema);
+
+    SinkRecord sinkRecord = new SinkRecord(UNSEEN_TOPIC,
+        PARTITION_114,
+        Schema.STRING_SCHEMA,
+        key,
+        schema,
+        record,
+        0 );
+
+    task.start(props, client);
+    task.open(new HashSet<>(Collections.singletonList(TOPIC_IN_CAPS_PARTITION)));
+    task.put(Collections.singleton(sinkRecord));
+    task.stop();
+
+    assertTrue(UNSEEN_TOPIC + " index created without errors ",
+        verifyIndexExist(cluster, UNSEEN_TOPIC.toLowerCase()));
+
+  }
+
+  private boolean verifyIndexExist(InternalTestCluster cluster, String ... indices) {
+    ActionFuture<IndicesExistsResponse> action = cluster
+        .client()
+        .admin()
+        .indices()
+        .exists(new IndicesExistsRequest(indices));
+
+    return action.actionGet().isExists();
+  }
 }