METRON-1849 Elasticsearch Index Write Functionality Should be Shared

metron-platform/metron-elasticsearch/src/main/java/org/apache/metron/elasticsearch/bulk/BulkDocumentWriter.java

@@ -0,0 +1,45 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.metron.elasticsearch.bulk;
+
+import org.apache.metron.indexing.dao.update.Document;
+
+/**
+ * Writes documents to an index in bulk.
+ *
+ * @param <D> The type of document to write.
+ */
+public interface BulkDocumentWriter<D extends Document> {
+
+    /**
+     * Add a document to the batch.
+     * @param document The document to write.
+     * @param index The name of the index to write to.
+     */
+    void addDocument(D document, String index);
+
+    /**
+     * @return The number of documents waiting to be written.
+     */
+    int size();
+
+    /**
+     * Write all documents in the batch.
+     */
+    BulkDocumentWriterResults<D> write();
+}

metron-platform/metron-elasticsearch/src/main/java/org/apache/metron/elasticsearch/bulk/BulkDocumentWriterResults.java

@@ -0,0 +1,68 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.metron.elasticsearch.bulk;
+
+import org.apache.metron.indexing.dao.update.Document;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * The result of writing documents in bulk using a {@link BulkDocumentWriter}.
+ * @param <D> The type of documents to write.
+ */
+public class BulkDocumentWriterResults<D extends Document> {
+
+  private List<WriteSuccess<D>> successes;
+  private List<WriteFailure<D>> failures;
+
+  public BulkDocumentWriterResults() {
+    this.successes = new ArrayList<>();
+    this.failures = new ArrayList<>();
+  }
+
+  public void add(WriteSuccess<D> success) {
+    this.successes.add(success);
+  }
+
+  public void addSuccess(D success) {
+    add(new WriteSuccess<D>(success));
+  }
+
+  public void addSuccesses(List<D> successes) {
+    for(D success: successes) {
+      addSuccess(success);
+    }
+  }
+
+  public List<WriteSuccess<D>> getSuccesses() {
+    return successes;
+  }
+
+  public void add(WriteFailure<D> failure) {
+    this.failures.add(failure);
+  }
+
+  public void addFailure(D document, Throwable cause, String message) {
+    add(new WriteFailure(document, cause, message));
+  }
+
+  public List<WriteFailure<D>> getFailures() {
+    return failures;
+  }
+}

metron-platform/metron-elasticsearch/src/main/java/org/apache/metron/elasticsearch/bulk/ElasticsearchBulkDocumentWriter.java

@@ -0,0 +1,166 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.metron.elasticsearch.bulk;
+
+import org.apache.commons.lang3.exception.ExceptionUtils;
+import org.apache.metron.elasticsearch.client.ElasticsearchClient;
+import org.apache.metron.indexing.dao.update.Document;
+import org.elasticsearch.action.DocWriteRequest;
+import org.elasticsearch.action.bulk.BulkItemResponse;
+import org.elasticsearch.action.bulk.BulkRequest;
+import org.elasticsearch.action.bulk.BulkResponse;
+import org.elasticsearch.action.index.IndexRequest;
+import org.elasticsearch.action.support.WriteRequest;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.lang.invoke.MethodHandles;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+
+/**
+ * Writes documents to an Elasticsearch index in bulk.
+ *
+ * @param <D> The type of document to write.
+ */
+public class ElasticsearchBulkDocumentWriter<D extends Document> implements BulkDocumentWriter<D> {
+
+    /**
+     * A {@link Document} along with the index it will be written to.
+     */
+    private class Indexable {
+        D document;
+        String index;
+
+        public Indexable(D document, String index) {
+            this.document = document;
+            this.index = index;
+        }
+    }
+
+    private static final Logger LOG = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
+    private ElasticsearchClient client;
+    private List<Indexable> documents;
+    private WriteRequest.RefreshPolicy refreshPolicy;
+
+    public ElasticsearchBulkDocumentWriter(ElasticsearchClient client) {
+        this.client = client;
+        this.documents = new ArrayList<>();
+        this.refreshPolicy = WriteRequest.RefreshPolicy.NONE;
+    }
+
+    @Override
+    public void addDocument(D document, String indexName) {
+        documents.add(new Indexable(document, indexName));
+        LOG.debug("Adding document to batch; document={}, index={}", document, indexName);
+    }
+
+    @Override
+    public BulkDocumentWriterResults<D> write() {
+        BulkDocumentWriterResults<D> results = new BulkDocumentWriterResults<>();
+        try {
+            // create an index request for each document
+            BulkRequest bulkRequest = new BulkRequest();
+            bulkRequest.setRefreshPolicy(refreshPolicy);
+            for(Indexable doc: documents) {
+                DocWriteRequest request = createRequest(doc.document, doc.index);
+                bulkRequest.add(request);
+            }
+
+            // submit the request and handle the response
+            BulkResponse bulkResponse = client.getHighLevelClient().bulk(bulkRequest);
+            handleBulkResponse(bulkResponse, documents, results);
+
+        } catch(IOException e) {
+            // assume all documents have failed
+            for(Indexable indexable: documents) {
+                D failed = indexable.document;
+                results.addFailure(failed, e, ExceptionUtils.getRootCauseMessage(e));
+            }
+            LOG.error("Failed to submit bulk request; all documents failed", e);
+
+        } finally {
+            // flush all documents no matter which ones succeeded or failed
+            documents.clear();
+        }
+
+        LOG.debug("Wrote document(s) to Elasticsearch; batchSize={}, success={}, failed={}",
+                documents.size(), results.getSuccesses().size(), results.getFailures().size());
+        return results;
+    }
+
+    @Override
+    public int size() {
+        return documents.size();
+    }
+
+    public ElasticsearchBulkDocumentWriter<D> withRefreshPolicy(WriteRequest.RefreshPolicy refreshPolicy) {
+        this.refreshPolicy = refreshPolicy;
+        return this;
+    }
+
+    private IndexRequest createRequest(D document, String index) {
+        if(document.getTimestamp() == null) {
+            throw new IllegalArgumentException("Document must contain the timestamp");
+        }
+        return new IndexRequest()
+                .source(document.getDocument())
+                .type(document.getSensorType() + "_doc")
+                .id(document.getGuid())
+                .index(index)
+                .timestamp(document.getTimestamp().toString());
+    }
+
+    /**
+     * Handles the {@link BulkResponse} received from Elasticsearch.
+     * @param bulkResponse The response received from Elasticsearch.
+     * @param documents The documents included in the bulk request.
+     * @param results The writer results.
+     */
+    private void handleBulkResponse(BulkResponse bulkResponse, List<Indexable> documents, BulkDocumentWriterResults<D> results) {
+        if (bulkResponse.hasFailures()) {
+
+            // interrogate the response to distinguish between those that succeeded and those that failed
+            for(BulkItemResponse response: bulkResponse) {
+                if(response.isFailed()) {
+                    // request failed
+                    D failed = getDocument(response.getItemId());
+                    Exception cause = response.getFailure().getCause();
+                    String message = response.getFailureMessage();
+                    results.addFailure(failed, cause, message);
+
+                } else {
+                    // request succeeded
+                    D success = getDocument(response.getItemId());
+                    results.addSuccess(success);
+                }
+            }
+        } else {
+            // all requests succeeded
+            for(Indexable success: documents) {
+                results.addSuccess(success.document);
+            }
+        }
+    }
+
+    private D getDocument(int index) {
+        return documents.get(index).document;
+    }
+}

metron-platform/metron-elasticsearch/src/main/java/org/apache/metron/elasticsearch/bulk/WriteFailure.java

@@ -0,0 +1,48 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.metron.elasticsearch.bulk;
+
+import org.apache.metron.indexing.dao.update.Document;
+
+/**
+ * Indicates that a document failed to be written by a {@link BulkDocumentWriter}.
+ * @param <D> The type of document that failed to write.
+ */
+public class WriteFailure <D extends Document> {
+  private D document;
+  private Throwable cause;
+  private String message;
+
+  public WriteFailure(D document, Throwable cause, String message) {
+    this.document = document;
+    this.cause = cause;
+    this.message = message;
+  }
+
+  public D getDocument() {
+    return document;
+  }
+
+  public Throwable getCause() {
+    return cause;
+  }
+
+  public String getMessage() {
+    return message;
+  }
+}

metron-platform/metron-elasticsearch/src/main/java/org/apache/metron/elasticsearch/bulk/WriteSuccess.java

@@ -0,0 +1,36 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.metron.elasticsearch.bulk;
+
+import org.apache.metron.indexing.dao.update.Document;
+
+/**
+ * Indicates that a document was successfully written by a {@link BulkDocumentWriter}.
+ * @param <D> The type of document written.
+ */
+public class WriteSuccess <D extends Document> {
+  private D document;
+
+  public WriteSuccess(D document) {
+    this.document = document;
+  }
+
+  public D getDocument() {
+    return document;
+  }
+}

metron-platform/metron-elasticsearch/src/main/java/org/apache/metron/elasticsearch/dao/ElasticsearchDao.java

@@ -191,7 +191,7 @@ public ElasticsearchDao withRefreshPolicy(WriteRequest.RefreshPolicy refreshPoli
   }
 
   protected Optional<String> getIndexName(String guid, String sensorType) throws IOException {
-    return updateDao.getIndexName(guid, sensorType);
+    return updateDao.findIndexNameByGUID(guid, sensorType);
   }
 
   protected SearchResponse search(SearchRequest request, QueryBuilder queryBuilder)