pl.allegro.tech.embeddedelasticsearch.ElasticRestClient Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of embedded-elasticsearch Show documentation
Tool that ease up creation of integration tests with Elasticsearch
There is a newer version: 2.4.0
Show newest version
package pl.allegro.tech.embeddedelasticsearch;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.commons.io.IOUtils;
import org.apache.http.HttpHeaders;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpDelete;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpHead;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.client.methods.HttpPut;
import org.apache.http.entity.StringEntity;
import org.apache.http.message.BasicHeader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import java.util.StringJoiner;
import java.util.function.Consumer;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.stream.Collectors.joining;
import static java.util.stream.Collectors.toList;
import static org.apache.http.entity.ContentType.APPLICATION_JSON;
import static pl.allegro.tech.embeddedelasticsearch.HttpStatusCodes.OK;

class ElasticRestClient {

    private static final Logger logger = LoggerFactory.getLogger(ElasticRestClient.class);

    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    private int elasticsearchHttpPort;
    private final HttpClient httpClient;
    private final IndicesDescription indicesDescription;
    private final TemplatesDescription templatesDescription;

    ElasticRestClient(int elasticsearchHttpPort, HttpClient httpClient, IndicesDescription indicesDescription, TemplatesDescription templatesDescription) {
        this.elasticsearchHttpPort = elasticsearchHttpPort;
        this.httpClient = httpClient;
        this.indicesDescription = indicesDescription;
        this.templatesDescription = templatesDescription;
    }

    void createIndices() {
        waitForClusterYellow();
        indicesDescription.getIndicesNames().forEach(this::createIndex);
    }

    void createIndex(String indexName) {
        if (!indexExists(indexName)) {
            HttpPut request = new HttpPut(url("/" + indexName));
            indicesDescription
                    .getIndexSettings(indexName)
                    .ifPresent(indexSettings -> setIndexSettingsAsEntity(request, indexSettings));
            httpClient.execute(request, response -> {
                if (response.getStatusLine().getStatusCode() != 200) {
                    String responseBody = readBodySafely(response);
                    throw new RuntimeException("Call to elasticsearch resulted in error:\n" + responseBody);
                }
            });
            waitForClusterYellow();
        }
    }

    private void setIndexSettingsAsEntity(HttpPut request, IndexSettings indexSettings) {
        request.setEntity(new StringEntity(indexSettings.toJson().toString(), APPLICATION_JSON));
    }

    private boolean indexExists(String indexName) {
        HttpHead request = new HttpHead(url("/" + indexName));
        return httpClient.execute(request, response -> response.getStatusLine().getStatusCode() == OK);
    }

    void createTemplates() {
        templatesDescription.getTemplatesNames().forEach(this::createTemplate);
    }

    void createTemplate(String templateName) {
        if (!templateExists(templateName)) {
            HttpPut request = new HttpPut(url("/_template/" + templateName));
            request.setEntity(new StringEntity(templatesDescription.getTemplateSettings(templateName), APPLICATION_JSON));
            httpClient.execute(request, response -> {
                if (response.getStatusLine().getStatusCode() != 200) {
                    String responseBody = readBodySafely(response);
                    throw new RuntimeException("Call to elasticsearch resulted in error:\n" + responseBody);
                }
            });
            waitForClusterYellow();
        }
    }

    private boolean templateExists(String templateName) {
        HttpHead request = new HttpHead(url("/" + templateName));
        return httpClient.execute(request, response ->
                response.getStatusLine().getStatusCode() == OK);
    }

    void deleteTemplates() {
        templatesDescription.getTemplatesNames().forEach(this::deleteTemplate);
    }

    void deleteTemplate(String templateName) {
        if (indexExists(templateName)) {
            HttpDelete request = new HttpDelete(url("/_template/" + templateName));
            httpClient.execute(request, (Consumer) response -> assertOk(response, "Delete request resulted in error"));
            waitForClusterYellow();
        } else {
            logger.warn("Template: {} does not exists so cannot be removed", templateName);
        }
    }

    private void waitForClusterYellow() {
        HttpGet request = new HttpGet(url("/_cluster/health?wait_for_status=yellow&timeout=60s"));
        httpClient.execute(request, (Consumer) response -> assertOk(response, "Cluster does not reached yellow status in specified timeout"));
    }

    void deleteIndices() {
        indicesDescription.getIndicesNames().forEach(this::deleteIndex);
    }

    void deleteIndex(String indexName) {
        if (indexExists(indexName)) {
            HttpDelete request = new HttpDelete(url("/" + indexName));
            httpClient.execute(request, (Consumer) response -> assertOk(response, "Delete request resulted in error"));
            waitForClusterYellow();
        } else {
            logger.warn("Index: {} does not exists so cannot be removed", indexName);
        }
    }

    void bulkIndex(Collection indexRequests) {
        String bulkRequestBody = indexRequests.stream()
                .flatMap(request ->
                    Stream.of(
                            indexMetadataJson(request.getIndexName(), request.getIndexType(), request.getId(), request.getRouting()),
                            request.getJson()
                    )
                )
                .map((jsonNodes) -> jsonNodes.replace('\n', ' ').replace('\r', ' '))
                .collect(joining("\n")) + "\n";

        performBulkRequest(url("/_bulk"), bulkRequestBody);
    }

    private String indexMetadataJson(String indexName, String indexType, String id, String routing) {
        StringJoiner joiner = new StringJoiner(",");

        if(indexName != null) {
            joiner.add("\"_index\": \"" + indexName + "\"");
        }

        if(indexType != null) {
            joiner.add("\"_type\": \"" + indexType + "\"");
        }

        if(id != null) {
            joiner.add("\"_id\": \"" + id + "\"");
        }

        if(routing != null) {
            joiner.add("\"_routing\": \"" + routing + "\"");
        }

        return "{ \"index\": {" + joiner.toString() + "} }";
    }

    void refresh() {
        HttpPost request = new HttpPost(url("/_refresh"));
        try {
            httpClient.execute(request);
        } finally {
            request.releaseConnection();
        }
    }

    private void performBulkRequest(String requestUrl, String bulkRequestBody) {
        HttpPost request = new HttpPost(requestUrl);
        request.setHeader(new BasicHeader(HttpHeaders.CONTENT_TYPE, "application/json"));
        request.setEntity(new StringEntity(bulkRequestBody, UTF_8));
        httpClient.execute(request, (Consumer) response -> assertOk(response, "Request finished with error"));
        refresh();
    }

    private String url(String path) {
        return "http://localhost:" + elasticsearchHttpPort + path;
    }

    private void assertOk(CloseableHttpResponse response, String message) {
        if (response.getStatusLine().getStatusCode() != OK) {
            throw new IllegalStateException(message + "\nResponse body:\n" + readBodySafely(response));
        }
    }

    private String readBodySafely(CloseableHttpResponse response) {
        try {
            return IOUtils.toString(response.getEntity().getContent(), UTF_8);
        } catch (IOException e) {
            logger.error("Error during reading response body", e);
            return "";
        }
    }

    List fetchAllDocuments(String... indices) {
        return fetchAllDocuments(null, indices);
    }

    List fetchAllDocuments(String routing, String... indices) {
        if (indices.length == 0) {
            return searchForDocuments(Optional.empty()).collect(toList());
        } else {
            return Stream.of(indices)
                    .flatMap((index) -> searchForDocuments(Optional.of(index), Optional.ofNullable(routing)))
                    .collect(toList());
        }
    }

    private Stream searchForDocuments(Optional indexMaybe) {
        return searchForDocuments(indexMaybe, Optional.empty());
    }

    private Stream searchForDocuments(Optional indexMaybe, Optional routing) {
        String searchCommand = prepareQuery(indexMaybe, routing);
        String body = fetchDocuments(searchCommand);
        return parseDocuments(body);
    }

    private String prepareQuery(Optional indexMaybe, Optional routing) {

        String routingQueryParam = routing
                .map(r -> "?routing=" + r)
                .orElse("");

        return indexMaybe
                .map(index -> "/" + index + "/_search" + routingQueryParam)
                .orElse("/_search");
    }

    private String fetchDocuments(String searchCommand) {
        HttpGet request = new HttpGet(url(searchCommand));
        return httpClient.execute(request, response -> {
            assertOk(response, "Error during search (" + searchCommand + ")");
            return readBodySafely(response);
        });
    }

    private Stream parseDocuments(String body) {
        try {
            JsonNode jsonNode = OBJECT_MAPPER.readTree(body);
            return StreamSupport.stream(jsonNode.get("hits").get("hits").spliterator(), false)
                    .map(hitNode -> hitNode.get("_source"))
                    .map(JsonNode::toString);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
}