Please wait. This can take some minutes ...
Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $
You can buy this project and download/modify it how often you want.
org.janusgraph.diskstorage.es.rest.RestElasticSearchClient Maven / Gradle / Ivy
// Copyright 2017 JanusGraph Authors
//
// Licensed 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.janusgraph.diskstorage.es.rest;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterators;
import com.google.common.collect.PeekingIterator;
import org.apache.http.HttpEntity;
import org.apache.http.HttpStatus;
import org.apache.http.entity.ByteArrayEntity;
import org.apache.http.entity.ContentType;
import org.apache.tinkerpop.shaded.jackson.annotation.JsonIgnoreProperties;
import org.apache.tinkerpop.shaded.jackson.core.JsonParseException;
import org.apache.tinkerpop.shaded.jackson.core.JsonProcessingException;
import org.apache.tinkerpop.shaded.jackson.core.type.TypeReference;
import org.apache.tinkerpop.shaded.jackson.databind.JsonMappingException;
import org.apache.tinkerpop.shaded.jackson.databind.ObjectMapper;
import org.apache.tinkerpop.shaded.jackson.databind.ObjectReader;
import org.apache.tinkerpop.shaded.jackson.databind.ObjectWriter;
import org.apache.tinkerpop.shaded.jackson.databind.SerializationFeature;
import org.apache.tinkerpop.shaded.jackson.databind.module.SimpleModule;
import org.elasticsearch.client.Request;
import org.elasticsearch.client.Response;
import org.elasticsearch.client.ResponseException;
import org.elasticsearch.client.RestClient;
import org.janusgraph.core.attribute.Geoshape;
import org.janusgraph.diskstorage.es.ElasticMajorVersion;
import org.janusgraph.diskstorage.es.ElasticSearchClient;
import org.janusgraph.diskstorage.es.ElasticSearchMutation;
import org.janusgraph.diskstorage.es.mapping.IndexMapping;
import org.janusgraph.diskstorage.es.mapping.TypedIndexMappings;
import org.janusgraph.diskstorage.es.mapping.TypelessIndexMappings;
import org.janusgraph.diskstorage.es.script.ESScriptResponse;
import org.javatuples.Pair;
import org.javatuples.Triplet;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;
import static org.janusgraph.util.encoding.StringEncoding.UTF8_CHARSET;
public class RestElasticSearchClient implements ElasticSearchClient {
private static final Logger log = LoggerFactory.getLogger(RestElasticSearchClient.class);
private static final String REQUEST_TYPE_DELETE = "DELETE";
private static final String REQUEST_TYPE_GET = "GET";
private static final String REQUEST_TYPE_POST = "POST";
private static final String REQUEST_TYPE_PUT = "PUT";
private static final String REQUEST_TYPE_HEAD = "HEAD";
private static final String REQUEST_SEPARATOR = "/";
private static final String REQUEST_PARAM_BEGINNING = "?";
private static final String REQUEST_PARAM_SEPARATOR = "&";
public static final String INCLUDE_TYPE_NAME_PARAMETER = "include_type_name";
private static final byte[] NEW_LINE_BYTES = "\n".getBytes(UTF8_CHARSET);
private static final Request INFO_REQUEST = new Request(REQUEST_TYPE_GET, REQUEST_SEPARATOR);
private static final ObjectMapper mapper;
private static final ObjectReader mapReader;
private static final ObjectWriter mapWriter;
static {
final SimpleModule module = new SimpleModule();
module.addSerializer(new Geoshape.GeoshapeGsonSerializerV2d0());
mapper = new ObjectMapper();
mapper.registerModule(module);
mapper.disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS);
mapReader = mapper.readerWithView(Map.class).forType(HashMap.class);
mapWriter = mapper.writerWithView(Map.class);
}
private static final ElasticMajorVersion DEFAULT_VERSION = ElasticMajorVersion.EIGHT;
private static final Function APPEND_OP = sb -> sb.append(sb.length() == 0 ? REQUEST_PARAM_BEGINNING : REQUEST_PARAM_SEPARATOR);
private final RestClient delegate;
private ElasticMajorVersion majorVersion;
private String bulkRefresh;
private boolean bulkRefreshEnabled = false;
private final String scrollKeepAlive;
private final boolean useMappingTypes;
private final boolean esVersion7;
private Integer retryOnConflict;
private final String retryOnConflictKey;
private final int retryAttemptLimit;
private final Set retryOnErrorCodes;
private final long retryInitialWaitMs;
private final long retryMaxWaitMs;
private final int bulkChunkSerializedLimitBytes;
public RestElasticSearchClient(RestClient delegate, int scrollKeepAlive, boolean useMappingTypesForES7,
int retryAttemptLimit, Set retryOnErrorCodes, long retryInitialWaitMs,
long retryMaxWaitMs, int bulkChunkSerializedLimitBytes) {
this.delegate = delegate;
majorVersion = getMajorVersion();
this.scrollKeepAlive = scrollKeepAlive+"s";
esVersion7 = ElasticMajorVersion.SEVEN.equals(majorVersion);
useMappingTypes = majorVersion.getValue() < 7 || (useMappingTypesForES7 && esVersion7);
retryOnConflictKey = majorVersion.getValue() >= 7 ? "retry_on_conflict" : "_retry_on_conflict";
this.retryAttemptLimit = retryAttemptLimit;
this.retryOnErrorCodes = Collections.unmodifiableSet(retryOnErrorCodes);
this.retryInitialWaitMs = retryInitialWaitMs;
this.retryMaxWaitMs = retryMaxWaitMs;
this.bulkChunkSerializedLimitBytes = bulkChunkSerializedLimitBytes;
}
@Override
public void close() throws IOException {
delegate.close();
}
@Override
public ElasticMajorVersion getMajorVersion() {
if (majorVersion != null) {
return majorVersion;
}
majorVersion = DEFAULT_VERSION;
try {
final Response response = delegate.performRequest(INFO_REQUEST);
try (final InputStream inputStream = response.getEntity().getContent()) {
final ClusterInfo info = mapper.readValue(inputStream, ClusterInfo.class);
majorVersion = ElasticMajorVersion.parse(info.getVersion() != null ? (String) info.getVersion().get("number") : null);
}
} catch (final IOException e) {
log.warn("Unable to determine Elasticsearch server version. Default to {}.", majorVersion, e);
}
return majorVersion;
}
@Override
public void clusterHealthRequest(String timeout) throws IOException {
Request clusterHealthRequest = new Request(REQUEST_TYPE_GET,
REQUEST_SEPARATOR + "_cluster" + REQUEST_SEPARATOR + "health");
clusterHealthRequest.addParameter("wait_for_status", "yellow");
clusterHealthRequest.addParameter("timeout", timeout);
final Response response = delegate.performRequest(clusterHealthRequest);
try (final InputStream inputStream = response.getEntity().getContent()) {
final Map values = mapReader.readValue(inputStream);
if (!values.containsKey("timed_out")) {
throw new IOException("Unexpected response for Elasticsearch cluster health request");
} else if (!Objects.equals(values.get("timed_out"), false)) {
throw new IOException("Elasticsearch timeout waiting for yellow status");
}
}
}
@Override
public boolean indexExists(String indexName) throws IOException {
final Response response = delegate.performRequest(new Request(REQUEST_TYPE_HEAD, REQUEST_SEPARATOR + indexName));
return response.getStatusLine().getStatusCode() == 200;
}
@Override
public boolean isIndex(String indexName) {
try {
final Response response = delegate.performRequest(new Request(REQUEST_TYPE_GET, REQUEST_SEPARATOR + indexName));
try (final InputStream inputStream = response.getEntity().getContent()) {
return mapper.readValue(inputStream, Map.class).containsKey(indexName);
}
} catch (final IOException ignored) {
}
return false;
}
@Override
public boolean isAlias(String aliasName) {
try {
delegate.performRequest(new Request(REQUEST_TYPE_GET, REQUEST_SEPARATOR + "_alias" + REQUEST_SEPARATOR + aliasName));
return true;
} catch (final IOException ignored) {
}
return false;
}
@Override
public void createStoredScript(String scriptName, Map script) throws IOException {
Request request = new Request(REQUEST_TYPE_POST, REQUEST_SEPARATOR + "_scripts" + REQUEST_SEPARATOR + scriptName);
performRequest(request, mapWriter.writeValueAsBytes(script));
}
@Override
public ESScriptResponse getStoredScript(String scriptName) throws IOException {
Request request = new Request(REQUEST_TYPE_GET,
REQUEST_SEPARATOR + "_scripts" + REQUEST_SEPARATOR + scriptName);
try{
final Response response = delegate.performRequest(request);
if(response.getStatusLine().getStatusCode() != 200){
throw new IOException("Error executing request: " + response.getStatusLine().getReasonPhrase());
}
try (final InputStream inputStream = response.getEntity().getContent()) {
return mapper.readValue(inputStream, new TypeReference() {});
} catch (final JsonParseException | JsonMappingException | ResponseException e) {
throw new IOException("Error when we try to parse ES script: "+response.getEntity().getContent());
}
} catch (ResponseException e){
final Response response = e.getResponse();
if(e.getResponse().getStatusLine().getStatusCode() == HttpStatus.SC_NOT_FOUND){
ESScriptResponse esScriptResponse = new ESScriptResponse();
esScriptResponse.setFound(false);
return esScriptResponse;
}
throw new IOException("Error executing request: " + response.getStatusLine().getReasonPhrase());
}
}
@Override
public void createIndex(String indexName, Map settings) throws IOException {
Request request = new Request(REQUEST_TYPE_PUT, REQUEST_SEPARATOR + indexName);
if(majorVersion.getValue() > 6){
if(useMappingTypes) {
request.addParameter(INCLUDE_TYPE_NAME_PARAMETER, "true");
}
if(settings != null && settings.size() > 0){
Map updatedSettings = new HashMap<>();
updatedSettings.put("settings", settings);
settings = updatedSettings;
}
}
performRequest(request, mapWriter.writeValueAsBytes(settings));
}
@Override
public void updateIndexSettings(String indexName, Map settings) throws IOException {
performRequest(REQUEST_TYPE_PUT, REQUEST_SEPARATOR + indexName + REQUEST_SEPARATOR +"_settings",
mapWriter.writeValueAsBytes(settings));
}
@Override
public void updateClusterSettings(Map settings) throws IOException {
performRequest(REQUEST_TYPE_PUT, REQUEST_SEPARATOR + "_cluster" + REQUEST_SEPARATOR + "settings",
mapWriter.writeValueAsBytes(settings));
}
@Override
public void addAlias(String alias, String index) throws IOException {
final Map actionAlias = ImmutableMap.of("actions", ImmutableList.of(ImmutableMap.of("add", ImmutableMap.of("index", index, "alias", alias))));
performRequest(REQUEST_TYPE_POST, REQUEST_SEPARATOR + "_aliases", mapWriter.writeValueAsBytes(actionAlias));
}
@Override
public Map getIndexSettings(String indexName) throws IOException {
final Response response = performRequest(REQUEST_TYPE_GET, REQUEST_SEPARATOR + indexName + REQUEST_SEPARATOR + "_settings", null);
try (final InputStream inputStream = response.getEntity().getContent()) {
final Map settings = mapper.readValue(inputStream, new TypeReference>() {});
return settings == null ? null : settings.get(indexName).getSettings().getMap();
}
}
@Override
public void createMapping(String indexName, String typeName, Map mapping) throws IOException {
Request request;
if(useMappingTypes){
request = new Request(REQUEST_TYPE_PUT,
REQUEST_SEPARATOR + indexName + REQUEST_SEPARATOR + "_mapping" + REQUEST_SEPARATOR + typeName);
if(esVersion7){
request.addParameter(INCLUDE_TYPE_NAME_PARAMETER, "true");
}
} else {
request = new Request(REQUEST_TYPE_PUT,
REQUEST_SEPARATOR + indexName + REQUEST_SEPARATOR + "_mapping");
}
performRequest(request, mapWriter.writeValueAsBytes(mapping));
}
@Override
public IndexMapping getMapping(String indexName, String typeName) throws IOException {
Request request;
if(useMappingTypes){
request = new Request(REQUEST_TYPE_GET,
REQUEST_SEPARATOR + indexName + REQUEST_SEPARATOR + "_mapping" + REQUEST_SEPARATOR + typeName);
if(esVersion7){
request.addParameter(INCLUDE_TYPE_NAME_PARAMETER, "true");
}
} else {
request = new Request(REQUEST_TYPE_GET,
REQUEST_SEPARATOR + indexName + REQUEST_SEPARATOR + "_mapping");
}
try (final InputStream inputStream = performRequest(request, null).getEntity().getContent()) {
if(useMappingTypes){
final Map settings = mapper.readValue(inputStream,
new TypeReference>() {});
return settings != null ? settings.get(indexName).getMappings().get(typeName) : null;
}
final Map settings = mapper.readValue(inputStream,
new TypeReference>() {});
return settings != null ? settings.get(indexName).getMappings() : null;
} catch (final JsonParseException | JsonMappingException | ResponseException e) {
log.info("Error when we try to get ES mapping", e);
return null;
}
}
@Override
public void deleteIndex(String indexName) throws IOException {
if (isAlias(indexName)) {
// aliased multi-index case
final String path = REQUEST_SEPARATOR + "_alias" + REQUEST_SEPARATOR + indexName;
final Response response = performRequest(REQUEST_TYPE_GET, path, null);
try (final InputStream inputStream = response.getEntity().getContent()) {
final Map records = mapper.readValue(inputStream, new TypeReference>() {});
if (records == null) return;
for (final String index : records.keySet()) {
if (indexExists(index)) {
performRequest(REQUEST_TYPE_DELETE, REQUEST_SEPARATOR + index, null);
}
}
}
}
}
@Override
public void clearStore(String indexName, String storeName) throws IOException {
String name = indexName + "_" + storeName;
if (indexExists(name)) {
performRequest(REQUEST_TYPE_DELETE, REQUEST_SEPARATOR + indexName + "_" + storeName, null);
}
}
@VisibleForTesting
class RequestBytes {
final byte [] requestBytes;
final byte [] requestSource;
@VisibleForTesting
RequestBytes(final ElasticSearchMutation request) throws JsonProcessingException {
Map requestData = new HashMap<>();
if (useMappingTypes) {
requestData.put("_index", request.getIndex());
requestData.put("_type", request.getType());
requestData.put("_id", request.getId());
} else {
requestData.put("_index", request.getIndex());
requestData.put("_id", request.getId());
}
if (retryOnConflict != null && request.getRequestType() == ElasticSearchMutation.RequestType.UPDATE) {
requestData.put(retryOnConflictKey, retryOnConflict);
}
this.requestBytes = mapWriter.writeValueAsBytes(ImmutableMap.of(request.getRequestType().name().toLowerCase(), requestData));
if (request.getSource() != null) {
this.requestSource = mapWriter.writeValueAsBytes(request.getSource());
} else {
this.requestSource = null;
}
}
@VisibleForTesting
int getSerializedSize() {
int serializedSize = this.requestBytes.length;
serializedSize+= 1; //For follow-up NEW_LINE_BYTES
if (this.requestSource != null) {
serializedSize += this.requestSource.length;
serializedSize+= 1; //For follow-up NEW_LINE_BYTES
}
return serializedSize;
}
private void writeTo(OutputStream outputStream) throws IOException {
outputStream.write(this.requestBytes);
outputStream.write(NEW_LINE_BYTES);
if (this.requestSource != null) {
outputStream.write(requestSource);
outputStream.write(NEW_LINE_BYTES);
}
}
}
private Pair buildBulkRequestInput(List requests, String ingestPipeline) throws IOException {
final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
for (final RequestBytes request : requests) {
request.writeTo(outputStream);
}
final StringBuilder bulkRequestQueryParameters = new StringBuilder();
if (ingestPipeline != null) {
APPEND_OP.apply(bulkRequestQueryParameters).append("pipeline=").append(ingestPipeline);
}
if (bulkRefreshEnabled) {
APPEND_OP.apply(bulkRequestQueryParameters).append("refresh=").append(bulkRefresh);
}
final String bulkRequestPath = REQUEST_SEPARATOR + "_bulk" + bulkRequestQueryParameters;
return Pair.with(bulkRequestPath, outputStream.toByteArray());
}
private List> pairErrorsWithSubmittedMutation(
//Bulk API is documented to return bulk item responses in the same order of submission
//https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html#bulk-api-response-body
//As such we only need to retry elements that failed
final List> bulkResponseItems,
final List submittedBulkRequestItems) {
final List> errors = new ArrayList<>(bulkResponseItems.size());
for (int itemIndex = 0; itemIndex < bulkResponseItems.size(); itemIndex++) {
Collection bulkResponseItem = bulkResponseItems.get(itemIndex).values();
if (bulkResponseItem.size() > 1) {
throw new IllegalStateException("There should only be a single item per bulk reponse item entry");
}
RestBulkResponse.RestBulkItemResponse item = bulkResponseItem.iterator().next();
if (item.getError() != null && item.getStatus() != HttpStatus.SC_NOT_FOUND) {
errors.add(Triplet.with(item.getError(), item.getStatus(), submittedBulkRequestItems.get(itemIndex)));
}
}
return errors;
}
@VisibleForTesting
class BulkRequestChunker implements Iterator> {
//By default, Elasticsearch writes are limited to 100mb, so chunk a given batch of requests so they stay under
//the specified limit
//https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html#docs-bulk-api-desc
//There is no "correct" number of actions to perform in a single bulk request. Experiment with different
// settings to find the optimal size for your particular workload. Note that Elasticsearch limits the maximum
// size of a HTTP request to 100mb by default
private final PeekingIterator requestIterator;
private final int[] exceptionallyLargeRequests;
@VisibleForTesting
BulkRequestChunker(List requests) throws JsonProcessingException {
List serializedRequests = new ArrayList<>(requests.size());
List requestSizesThatWereTooLarge = new ArrayList<>();
for (ElasticSearchMutation request : requests) {
RequestBytes requestBytes = new RequestBytes(request);
int requestSerializedSize = requestBytes.getSerializedSize();
if (requestSerializedSize <= bulkChunkSerializedLimitBytes) {
//Only keep items that we can actually send in memory
serializedRequests.add(requestBytes);
} else {
requestSizesThatWereTooLarge.add(requestSerializedSize);
}
}
this.requestIterator = Iterators.peekingIterator(serializedRequests.iterator());
//Condense request sizes that are too large into an int array to remove Boxed & List memory overhead
this.exceptionallyLargeRequests = requestSizesThatWereTooLarge.isEmpty() ? null :
requestSizesThatWereTooLarge.stream().mapToInt(Integer::intValue).toArray();
}
@Override
public boolean hasNext() {
//Make sure hasNext() still returns true if exceptionally large requests were attempted to be submitted
//This allows next() to throw after all well sized requests have been chunked for submission
return requestIterator.hasNext() || exceptionallyLargeRequests != null;
}
@Override
public List next() {
List serializedRequests = new ArrayList<>();
int chunkSerializedTotal = 0;
while (requestIterator.hasNext()) {
RequestBytes peeked = requestIterator.peek();
chunkSerializedTotal += peeked.getSerializedSize();
if (chunkSerializedTotal <= bulkChunkSerializedLimitBytes) {
serializedRequests.add(requestIterator.next());
} else {
//Adding this element would exceed the limit, so return the chunk
return serializedRequests;
}
}
//Check if we should throw an exception for items that were exceptionally large and therefore undeliverable.
//This is only done after all items that could be sent have been sent
if (serializedRequests.isEmpty() && this.exceptionallyLargeRequests != null) {
throw new IllegalArgumentException(String.format(
"Bulk request item(s) larger than permitted chunk limit. Limit is %s. Serialized item size(s) %s",
bulkChunkSerializedLimitBytes, Arrays.toString(this.exceptionallyLargeRequests)));
}
//All remaining requests fit in this chunk
return serializedRequests;
}
}
@Override
public void bulkRequest(final List requests, String ingestPipeline) throws IOException {
BulkRequestChunker bulkRequestChunker = new BulkRequestChunker(requests);
while (bulkRequestChunker.hasNext()) {
List bulkRequestChunk = bulkRequestChunker.next();
int retryCount = 0;
while (true) {
final Pair bulkRequestInput = buildBulkRequestInput(bulkRequestChunk, ingestPipeline);
final Response response = performRequest(REQUEST_TYPE_POST, bulkRequestInput.getValue0(), bulkRequestInput.getValue1());
try (final InputStream inputStream = response.getEntity().getContent()) {
final RestBulkResponse bulkResponse = mapper.readValue(inputStream, RestBulkResponse.class);
List> bulkItemsThatFailed = pairErrorsWithSubmittedMutation(bulkResponse.getItems(), bulkRequestChunk);
if (!bulkItemsThatFailed.isEmpty()) {
//Only retry the bulk request if *all* the bulk response item error codes are retry error codes
final Set errorCodes = bulkItemsThatFailed.stream().map(Triplet::getValue1).collect(Collectors.toSet());
if (retryCount < retryAttemptLimit && retryOnErrorCodes.containsAll(errorCodes)) {
//Build up the next request batch, of only the failed mutations
bulkRequestChunk = bulkItemsThatFailed.stream().map(Triplet::getValue2).collect(Collectors.toList());
performRetryWait(retryCount);
retryCount++;
} else {
final List errorItems = bulkItemsThatFailed.stream().map(Triplet::getValue0).collect(Collectors.toList());
errorItems.forEach(error -> log.error("Failed to execute ES query: {}", error));
throw new IOException("Failure(s) in Elasticsearch bulk request: " + errorItems);
}
} else {
//The entire bulk request was successful, leave the loop
break;
}
}
}
}
}
public void setRetryOnConflict(Integer retryOnConflict) {
this.retryOnConflict = retryOnConflict;
}
@Override
public long countTotal(String indexName, Map requestData) throws IOException {
final Request request = new Request(REQUEST_TYPE_GET, REQUEST_SEPARATOR + indexName + REQUEST_SEPARATOR + "_count");
final byte[] requestDataBytes = mapper.writeValueAsBytes(requestData);
if (log.isDebugEnabled()) {
log.debug("Elasticsearch request: " + mapper.writerWithDefaultPrettyPrinter().writeValueAsString(requestData));
}
final Response response = performRequest(request, requestDataBytes);
try (final InputStream inputStream = response.getEntity().getContent()) {
return mapper.readValue(inputStream, RestCountResponse.class).getCount();
}
}
/**
* Execute the aggregation request using Elasticsearch index.
* Elasticsearch uses double values to hold and represent numeric data. As a result, aggregations on long numbers
* greater than 2^53 are approximate.
* Elasticsearch, limits for long values
*
* @param indexName the name of the ElasticSearch index on which the aggregation is executed
* @param requestData the filter query
* @param agg the name of the aggregation operation (min, max, avg, sum)
* @param fieldName the name of the field on which the aggregation is computed
* @return the result of the aggregation
* @throws IOException
*/
private double executeAggs(String indexName, Map requestData, String agg, String fieldName) throws IOException {
final Request request = new Request(REQUEST_TYPE_GET, REQUEST_SEPARATOR + indexName + REQUEST_SEPARATOR + "_search");
requestData.put("aggs", ImmutableMap.of("agg_result", ImmutableMap.of(agg, ImmutableMap.of("field", fieldName))));
final byte[] requestDataBytes = mapper.writeValueAsBytes(requestData);
if (log.isDebugEnabled()) {
log.debug("Elasticsearch request: " + mapper.writerWithDefaultPrettyPrinter().writeValueAsString(requestData));
}
final Response response = performRequest(request, requestDataBytes);
try (final InputStream inputStream = response.getEntity().getContent()) {
return mapper.readValue(inputStream, RestAggResponse.class).getAggregations().getAggResult().getValue();
}
}
private Number adaptNumberType(double value, Class extends Number> expectedType) {
if (expectedType == null) return value;
else if (Byte.class.isAssignableFrom(expectedType)) return (byte)value;
else if (Short.class.isAssignableFrom(expectedType)) return (short)value;
else if (Integer.class.isAssignableFrom(expectedType)) return (int)value;
else if (Long.class.isAssignableFrom(expectedType)) return (long)value;
else if (Float.class.isAssignableFrom(expectedType)) return (float)value;
else return value;
}
@Override
public Number min(String indexName, Map requestData, String fieldName, Class extends Number> expectedType) throws IOException {
return adaptNumberType(executeAggs(indexName, requestData, "min", fieldName), expectedType);
}
@Override
public Number max(String indexName, Map requestData, String fieldName, Class extends Number> expectedType) throws IOException {
return adaptNumberType(executeAggs(indexName, requestData, "max", fieldName), expectedType);
}
@Override
public double avg(String indexName, Map requestData, String fieldName) throws IOException {
return executeAggs(indexName, requestData, "avg", fieldName);
}
@Override
public Number sum(String indexName, Map requestData, String fieldName, Class extends Number> expectedType) throws IOException {
Class extends Number> returnType;
double sum = executeAggs(indexName, requestData, "sum", fieldName);
if (Float.class.isAssignableFrom(expectedType) || Double.class.isAssignableFrom(expectedType))
return sum;
else
return (long)sum;
}
@Override
public RestSearchResponse search(String indexName, Map requestData, boolean useScroll) throws IOException {
final StringBuilder path = new StringBuilder(REQUEST_SEPARATOR).append(indexName);
path.append(REQUEST_SEPARATOR).append("_search");
if (useScroll) {
path.append(REQUEST_PARAM_BEGINNING).append("scroll=").append(scrollKeepAlive);
}
return search(requestData, path.toString());
}
@Override
public RestSearchResponse search(String scrollId) throws IOException {
final Map requestData = new HashMap<>();
requestData.put("scroll", scrollKeepAlive);
requestData.put("scroll_id", scrollId);
return search(requestData, REQUEST_SEPARATOR + "_search" + REQUEST_SEPARATOR + "scroll");
}
@Override
public void deleteScroll(String scrollId) throws IOException {
delegate.performRequest(new Request(REQUEST_TYPE_DELETE, REQUEST_SEPARATOR + "_search" + REQUEST_SEPARATOR + "scroll" + REQUEST_SEPARATOR + scrollId));
}
public void setBulkRefresh(String bulkRefresh) {
this.bulkRefresh = bulkRefresh;
bulkRefreshEnabled = bulkRefresh != null && !bulkRefresh.equalsIgnoreCase("false");
}
private RestSearchResponse search(Map requestData, String path) throws IOException {
final Request request = new Request(REQUEST_TYPE_POST, path);
final byte[] requestDataBytes = mapper.writeValueAsBytes(requestData);
if (log.isDebugEnabled()) {
log.debug("Elasticsearch request: " + mapper.writerWithDefaultPrettyPrinter().writeValueAsString(requestData));
}
final Response response = performRequest(request, requestDataBytes);
try (final InputStream inputStream = response.getEntity().getContent()) {
return mapper.readValue(inputStream, RestSearchResponse.class);
}
}
private Response performRequest(String method, String path, byte[] requestData) throws IOException {
return performRequest(new Request(method, path), requestData);
}
private Response performRequestWithRetry(Request request) throws IOException {
int retryCount = 0;
while (true) {
try {
return delegate.performRequest(request);
} catch (ResponseException e) {
if (!retryOnErrorCodes.contains(e.getResponse().getStatusLine().getStatusCode()) || retryCount >= retryAttemptLimit) {
throw e;
}
performRetryWait(retryCount);
}
retryCount++;
}
}
private void performRetryWait(int retryCount) {
long waitDurationMs = Math.min((long) (retryInitialWaitMs * Math.pow(10, retryCount)), retryMaxWaitMs);
log.warn("Retrying Elasticsearch request in {} ms. Attempt {} of {}", waitDurationMs, retryCount, retryAttemptLimit);
try {
Thread.sleep(waitDurationMs);
} catch (InterruptedException interruptedException) {
throw new RuntimeException(String.format("Thread interrupted while waiting for retry attempt %d of %d", retryCount, retryAttemptLimit), interruptedException);
}
}
private Response performRequest(Request request, byte[] requestData) throws IOException {
final HttpEntity entity = requestData != null ? new ByteArrayEntity(requestData, ContentType.APPLICATION_JSON) : null;
request.setEntity(entity);
final Response response = performRequestWithRetry(request);
if (response.getStatusLine().getStatusCode() >= 400) {
throw new IOException("Error executing request: " + response.getStatusLine().getReasonPhrase());
}
return response;
}
@JsonIgnoreProperties(ignoreUnknown=true)
private static final class ClusterInfo {
private Map version;
public Map getVersion() {
return version;
}
public void setVersion(Map version) {
this.version = version;
}
}
}