com.lordofthejars.nosqlunit.elasticsearch2.parser.DataReader Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of nosqlunit-elasticsearch2 Show documentation
The newest version!
package com.lordofthejars.nosqlunit.elasticsearch2.parser;

import org.codehaus.jackson.map.ObjectMapper;
import org.codehaus.jackson.type.TypeReference;
import org.elasticsearch.action.index.IndexRequestBuilder;
import org.elasticsearch.client.Client;

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

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 {
			final List> documents = getDocuments(data);
			insertDocuments(documents);
			refreshNode();
		} 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) {
			final Object object = document.get(DOCUMENT_ELEMENT);

			if (object instanceof List) {
				@SuppressWarnings("unchecked")
				final List> properties = (List>) object;
				insertDocument(properties);
			} else {
				throw new IllegalArgumentException("Array of Indexes and Data are required.");
			}
		}
	}

	private void insertDocument(List> properties) {
		final 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();
		}
	}

	@SuppressWarnings("unchecked")
	private Map dataOfDocument(Object object) {
		return (Map) object;
	}

	private IndexRequestBuilder createIndex(Object object) {
		@SuppressWarnings("unchecked")
		final Map indexInformation = (Map) object;

		final 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;
	}

	@SuppressWarnings("unchecked")
	public static List> getDocuments(InputStream data) throws IOException {
		final Map rootNode = MAPPER.readValue(data, new TypeReference>() {
		});
		final Object dataElements = rootNode.get(DOCUMENTS_ELEMENT);

		if (dataElements instanceof List) {
			return (List>) dataElements;
		} else {
			throw new IllegalArgumentException("Array of documents are required.");
		}
	}
}