add support for headers: default ones and per request

This commit is contained in:
javanna 2016-05-13 15:20:21 +02:00 committed by Luca Cavanna
parent 85a7721185
commit c0a72c1686
3 changed files with 69 additions and 10 deletions

View File

@ -20,6 +20,7 @@ package org.elasticsearch.client;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.http.Header;
import org.apache.http.HttpEntity;
import org.apache.http.HttpHost;
import org.apache.http.client.config.RequestConfig;
@ -40,6 +41,7 @@ import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.Iterator;
@ -75,10 +77,15 @@ public final class RestClient implements Closeable {
this.connections = Collections.unmodifiableList(connections);
}
public ElasticsearchResponse performRequest(String method, String endpoint, Map<String, Object> params, HttpEntity entity)
throws IOException {
public ElasticsearchResponse performRequest(String method, String endpoint, Map<String, Object> params,
HttpEntity entity, Header... headers) throws IOException {
URI uri = buildUri(endpoint, params);
HttpRequestBase request = createHttpRequest(method, uri, entity);
if (headers.length > 0) {
for (Header header : headers) {
request.addHeader(header);
}
}
//we apply a soft margin so that e.g. if a request took 59 seconds and timeout is set to 60 we don't do another attempt
long retryTimeout = Math.round(this.maxRetryTimeout / (float)100 * 98);
IOException lastSeenException = null;
@ -276,13 +283,15 @@ public final class RestClient implements Closeable {
private CloseableHttpClient httpClient;
private int maxRetryTimeout = DEFAULT_MAX_RETRY_TIMEOUT;
private HttpHost[] hosts;
private Collection<? extends Header> defaultHeaders;
private Builder() {
}
/**
* Sets the http client. A new default one will be created if not specified, by calling {@link #createDefaultHttpClient()}.
* Sets the http client. A new default one will be created if not
* specified, by calling {@link #createDefaultHttpClient(Collection)}.
*
* @see CloseableHttpClient
*/
@ -316,12 +325,29 @@ public final class RestClient implements Closeable {
return this;
}
/**
* Sets the default request headers, to be used when creating the default http client instance.
* In case the http client is set through {@link #setHttpClient(CloseableHttpClient)}, the default headers need to be
* set to it externally during http client construction.
*/
public Builder setDefaultHeaders(Collection<? extends Header> defaultHeaders) {
this.defaultHeaders = defaultHeaders;
return this;
}
/**
* Creates a new {@link RestClient} based on the provided configuration.
*/
public RestClient build() {
if (httpClient == null) {
httpClient = createDefaultHttpClient();
httpClient = createDefaultHttpClient(defaultHeaders);
} else {
if (defaultHeaders != null) {
throw new IllegalArgumentException("defaultHeaders need to be set to the HttpClient directly when manually provided");
}
}
if (hosts == null || hosts.length == 0) {
throw new IllegalArgumentException("no hosts provided");
}
return new RestClient(httpClient, maxRetryTimeout, hosts);
}
@ -331,7 +357,7 @@ public final class RestClient implements Closeable {
*
* @see CloseableHttpClient
*/
public static CloseableHttpClient createDefaultHttpClient() {
public static CloseableHttpClient createDefaultHttpClient(Collection<? extends Header> defaultHeaders) {
PoolingHttpClientConnectionManager connectionManager = new PoolingHttpClientConnectionManager();
//default settings may be too constraining
connectionManager.setDefaultMaxPerRoute(10);
@ -342,7 +368,11 @@ public final class RestClient implements Closeable {
.setSocketTimeout(DEFAULT_SOCKET_TIMEOUT)
.setConnectionRequestTimeout(DEFAULT_CONNECTION_REQUEST_TIMEOUT).build();
return HttpClientBuilder.create().setConnectionManager(connectionManager).setDefaultRequestConfig(requestConfig).build();
HttpClientBuilder httpClientBuilder = HttpClientBuilder.create();
if (defaultHeaders != null) {
httpClientBuilder.setDefaultHeaders(defaultHeaders);
}
return httpClientBuilder.setConnectionManager(connectionManager).setDefaultRequestConfig(requestConfig).build();
}
}

View File

@ -28,6 +28,7 @@ import org.elasticsearch.client.RestClient;
import java.io.Closeable;
import java.io.IOException;
import java.util.Collection;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.Executors;
@ -201,7 +202,7 @@ public final class Sniffer extends RestClient.FailureListener implements Closeab
/**
* Sets the http client. Mandatory argument. Best practice is to use the same client used
* within {@link org.elasticsearch.client.RestClient} which can be created manually or
* through {@link RestClient.Builder#createDefaultHttpClient()}.
* through {@link RestClient.Builder#createDefaultHttpClient(Collection)}.
* @see CloseableHttpClient
*/
public Builder setRestClient(RestClient restClient) {

View File

@ -21,10 +21,15 @@ package org.elasticsearch.client;
import com.carrotsearch.randomizedtesting.generators.RandomInts;
import org.apache.http.HttpHost;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.message.BasicHeader;
import org.apache.lucene.util.LuceneTestCase;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.logging.LogManager;
public class RestClientBuilderTests extends LuceneTestCase {
@ -56,12 +61,27 @@ public class RestClientBuilderTests extends LuceneTestCase {
}
try {
RestClient.builder();
RestClient.builder().build();
fail("should have failed");
} catch(IllegalArgumentException e) {
assertEquals(e.getMessage(), "no hosts provided");
}
try {
RestClient.builder().setHosts(new HttpHost[]{new HttpHost("localhost", 9200), null}).build();
fail("should have failed");
} catch(NullPointerException e) {
assertEquals(e.getMessage(), "host cannot be null");
}
try (CloseableHttpClient httpClient = HttpClientBuilder.create().build()) {
RestClient.builder().setHttpClient(httpClient)
.setDefaultHeaders(Collections.singleton(new BasicHeader("header", "value"))).build();
fail("should have failed");
} catch(IllegalArgumentException e) {
assertEquals(e.getMessage(), "defaultHeaders need to be set to the HttpClient directly when manually provided");
}
RestClient.Builder builder = RestClient.builder();
int numNodes = RandomInts.randomIntBetween(random(), 1, 5);
HttpHost[] hosts = new HttpHost[numNodes];
@ -70,11 +90,19 @@ public class RestClientBuilderTests extends LuceneTestCase {
}
builder.setHosts(hosts);
//TODO test one host is null among others
if (random().nextBoolean()) {
builder.setHttpClient(HttpClientBuilder.create().build());
} else {
if (random().nextBoolean()) {
int numHeaders = RandomInts.randomIntBetween(random(), 1, 5);
Collection<BasicHeader> headers = new ArrayList<>(numHeaders);
for (int i = 0; i < numHeaders; i++) {
headers.add(new BasicHeader("header" + i, "value"));
}
builder.setDefaultHeaders(headers);
}
}
if (random().nextBoolean()) {
builder.setMaxRetryTimeout(RandomInts.randomIntBetween(random(), 1, Integer.MAX_VALUE));
}