com.lordofthejars.nosqlunit.elasticsearch.parser.DataReader Maven / Gradle / Ivy
package com.lordofthejars.nosqlunit.elasticsearch.parser;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.codehaus.jackson.JsonParseException;
import org.codehaus.jackson.map.JsonMappingException;
import org.codehaus.jackson.map.ObjectMapper;
import org.elasticsearch.action.index.IndexRequestBuilder;
import org.elasticsearch.client.Client;
public class DataReader {
private static final ObjectMapper MAPPER = new ObjectMapper();
public static final String DOCUMENTS_ELEMENT = "documents";
public static final String DOCUMENT_ELEMENT = "document";
public static final String DATA_ELEMENT = "data";
public static final String INDEX_ELEMENT = "index";
public static final String INDEX_NAME_ELEMENT = "indexName";
public static final String INDEX_TYPE_ELEMENT = "indexType";
public static final String INDEX_ID_ELEMENT = "indexId";
private Client client;
public DataReader(Client client) {
this.client = client;
}
public void read(InputStream data) {
try {
List> documents = getDocuments(data);
insertDocuments(documents);
refreshNode();
} catch (JsonParseException e) {
throw new IllegalArgumentException(e);
} catch (JsonMappingException e) {
throw new IllegalArgumentException(e);
} catch (IOException e) {
throw new IllegalArgumentException(e);
}
}
private void refreshNode() {
client.admin().indices().prepareRefresh().execute().actionGet();
}
private void insertDocuments(List> documents) {
for (Map document : documents) {
Object object = document.get(DOCUMENT_ELEMENT);
if (object instanceof List) {
List> properties = (List>) object;
insertDocument(properties);
} else {
throw new IllegalArgumentException("Array of Indexes and Data are required.");
}
}
}
private void insertDocument(List> properties) {
List indexes = new ArrayList();
Map dataOfDocument = new HashMap();
for (Map property : properties) {
if (property.containsKey(INDEX_ELEMENT)) {
indexes.add(createIndex(property.get(INDEX_ELEMENT)));
} else {
if (property.containsKey(DATA_ELEMENT)) {
dataOfDocument = dataOfDocument(property.get(DATA_ELEMENT));
}
}
}
insertIndexes(indexes, dataOfDocument);
}
private void insertIndexes(List indexes, Map dataOfDocument) {
for (IndexRequestBuilder indexRequestBuilder : indexes) {
indexRequestBuilder.setSource(dataOfDocument).execute().actionGet();
}
}
private Map dataOfDocument(Object object) {
Map data = (Map) object;
return data;
}
private IndexRequestBuilder createIndex(Object object) {
Map indexInformation = (Map) object;
IndexRequestBuilder prepareIndex = client.prepareIndex();
if (indexInformation.containsKey(INDEX_NAME_ELEMENT)) {
prepareIndex.setIndex(indexInformation.get(INDEX_NAME_ELEMENT));
}
if (indexInformation.containsKey(INDEX_TYPE_ELEMENT)) {
prepareIndex.setType(indexInformation.get(INDEX_TYPE_ELEMENT));
}
if (indexInformation.containsKey(INDEX_ID_ELEMENT)) {
prepareIndex.setId(indexInformation.get(INDEX_ID_ELEMENT));
}
return prepareIndex;
}
public static List> getDocuments(InputStream data) throws JsonParseException,
JsonMappingException, IOException {
Map rootNode = MAPPER.readValue(data, Map.class);
Object dataElements = rootNode.get(DOCUMENTS_ELEMENT);
if (dataElements instanceof List) {
return (List>) dataElements;
} else {
throw new IllegalArgumentException("Array of documents are required.");
}
}
}