com.github.joschi.nosqlunit.elasticsearch.http.parser.DataReader Maven / Gradle / Ivy

Go to download
package com.github.joschi.nosqlunit.elasticsearch.http.parser;

import com.google.common.base.MoreObjects;
import com.google.gson.Gson;
import io.searchbox.client.JestClient;
import io.searchbox.client.JestResult;
import io.searchbox.core.Bulk;
import io.searchbox.core.BulkResult;
import io.searchbox.core.Index;
import io.searchbox.indices.CreateIndex;
import io.searchbox.indices.Refresh;
import io.searchbox.indices.template.DeleteTemplate;
import io.searchbox.indices.template.PutTemplate;

import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.Reader;
import java.io.UncheckedIOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class DataReader {
    private static final Gson GSON = new Gson();

    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 final JestClient client;
    private final boolean createIndices;
    private final Map indexSettings;
    private final Map> templates;

    public DataReader(JestClient client,
                      boolean createIndices,
                      Map indexSettings,
                      Map> templates) {
        this.client = client;
        this.createIndices = createIndices;
        this.indexSettings = indexSettings;
        this.templates = templates;
    }

    public void read(InputStream data) {
        try {
            final List> documents = getDocuments(data);
            if (!templates.isEmpty()) {
                createTemplates(templates);
            }

            if (createIndices) {
                createIndices(documents, indexSettings);
            }

            insertDocuments(documents);

            if (!templates.isEmpty()) {
                deleteTemplates(templates.keySet());
                refreshNode();
            }

            refreshNode();
        } catch (IOException e) {
            throw new IllegalArgumentException(e);
        }
    }

    private void deleteTemplates(Set templates) throws IOException {
        for (String template : templates) {
            final DeleteTemplate deleteTemplate = new DeleteTemplate.Builder(template).build();
            final JestResult result = client.execute(deleteTemplate);
            if (!result.isSucceeded()) {
                throw new IllegalStateException("Error while deleting template \"" + template + "\": " + result.getErrorMessage());
            }
        }
    }

    private void createTemplates(Map> templates) throws IOException {
        for (Map.Entry> template : templates.entrySet()) {
            final String templateName = template.getKey();
            final PutTemplate putTemplate = new PutTemplate.Builder(templateName, template.getValue()).build();
            final JestResult result = client.execute(putTemplate);
            if (!result.isSucceeded()) {
                throw new IllegalStateException("Error while creating template \"" + templateName + "\": " + result.getErrorMessage());
            }
        }
    }

    private void refreshNode() {
        final Refresh request = new Refresh.Builder().build();
        try {
            final JestResult result = client.execute(request);
            if (!result.isSucceeded()) {
                throw new IllegalStateException("Error while refreshing indices: " + result.getErrorMessage());
            }
        } catch (IOException e) {
            throw new UncheckedIOException("Unable to refresh indices", e);
        }
    }

    @SuppressWarnings("unchecked")
    private void createIndices(List> documents, Map indexSettings) throws IOException {
        for (Map document : documents) {
            final Object object = document.get(DOCUMENT_ELEMENT);

            if (object instanceof List) {
                final List> properties = (List>) object;
                createIndex(properties, indexSettings);
            } else {
                throw new IllegalArgumentException("Array of Indexes and Data are required.");
            }
        }
    }

    private void insertDocuments(List> documents) throws IOException {
        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.");
            }
        }
    }

    @SuppressWarnings("unchecked")
    private void createIndex(List> properties, Map indexSettings) throws IOException {
        for (Map property : properties) {
            if (property.containsKey(INDEX_ELEMENT)) {
                final Map indexInformation = (Map) property.get(INDEX_ELEMENT);
                final String indexName = indexInformation.get(INDEX_NAME_ELEMENT);
                if (indexName == null) {
                    throw new IllegalArgumentException("Missing index name element in " + indexInformation);
                }

                final CreateIndex request = new CreateIndex.Builder(indexName)
                        .settings(indexSettings)
                        .build();
                final JestResult result = client.execute(request);
                if (!result.isSucceeded()) {
                    throw new IllegalStateException("Error while creating index " + indexName + ": " + result.getErrorMessage());
                }
            }
        }
    }

    private void insertDocument(List> properties) throws IOException {
        final List> indexes = new ArrayList<>();
        Map dataOfDocument = new HashMap<>();

        for (Map property : properties) {
            if (property.containsKey(INDEX_ELEMENT)) {
                @SuppressWarnings("unchecked") final Map indexInformation = (Map) property.get(INDEX_ELEMENT);
                indexes.add(indexInformation);
            } else {
                if (property.containsKey(DATA_ELEMENT)) {
                    dataOfDocument = dataOfDocument(property.get(DATA_ELEMENT));
                }
            }
        }

        insertIndexes(indexes, dataOfDocument);
    }

    private void insertIndexes(List> indexes, Map dataOfDocument) throws IOException {
        final Bulk.Builder bulkBuilder = new Bulk.Builder();
        for (Map indexInformation : indexes) {
            final Index indexDocument = indexDocument(indexInformation, dataOfDocument);
            bulkBuilder.addAction(indexDocument);
        }

        final BulkResult result = client.execute(bulkBuilder.build());
        if (!result.isSucceeded()) {
            final StringBuilder sb = new StringBuilder("Error while bulk indexing documents: ")
                    .append(result.getErrorMessage());

            if(!result.getFailedItems().isEmpty()) {
                sb.append('\n');
                for (BulkResult.BulkResultItem bulkResultItem : result.getFailedItems()) {
                    final String s = MoreObjects.toStringHelper(BulkResult.BulkResultItem.class)
                            .omitNullValues()
                            .add("operation", bulkResultItem.operation)
                            .add("index", bulkResultItem.index)
                            .add("type", bulkResultItem.type)
                            .add("id", bulkResultItem.id)
                            .add("status", bulkResultItem.status)
                            .add("error", bulkResultItem.error)
                            .add("errorType", bulkResultItem.errorType)
                            .add("errorReason", bulkResultItem.errorReason)
                            .add("version", bulkResultItem.version)
                            .toString();
                    sb.append('\n').append(s);
                }
            }

            throw new IllegalStateException(sb.toString());
        }
    }

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

    private Index indexDocument(Map indexInformation, Object source) {
        if (!indexInformation.containsKey(INDEX_NAME_ELEMENT)) {
            throw new IllegalArgumentException("Missing index name element in " + indexInformation);
        }

        final Index.Builder createIndexBuilder = new Index.Builder(source)
                .index(indexInformation.get(INDEX_NAME_ELEMENT));

        if (indexInformation.containsKey(INDEX_TYPE_ELEMENT)) {
            createIndexBuilder.type(indexInformation.get(INDEX_TYPE_ELEMENT));
        }

        if (indexInformation.containsKey(INDEX_ID_ELEMENT)) {
            createIndexBuilder.id(indexInformation.get(INDEX_ID_ELEMENT));
        }

        return createIndexBuilder.build();
    }

    @SuppressWarnings("unchecked")
    public static List> getDocuments(InputStream data) throws IOException {
        try (Reader reader = new InputStreamReader(data)) {
            final Map rootNode = (Map) GSON.fromJson(reader, Map.class);
            final Object dataElements = rootNode.get(DOCUMENTS_ELEMENT);

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