HLRest: Allow caller to set per request options

client/rest-high-level/src/main/java/org/elasticsearch/client/RestHighLevelClient.java

@@ -26,8 +26,6 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRequest;
 import org.elasticsearch.action.ActionRequestValidationException;
-import org.elasticsearch.action.admin.cluster.repositories.get.GetRepositoriesRequest;
-import org.elasticsearch.action.admin.cluster.repositories.get.GetRepositoriesResponse;
 import org.elasticsearch.action.bulk.BulkRequest;
 import org.elasticsearch.action.bulk.BulkResponse;
 import org.elasticsearch.action.delete.DeleteRequest;
@@ -168,6 +166,7 @@
 import java.util.Objects;
 import java.util.ServiceLoader;
 import java.util.Set;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -186,6 +185,7 @@
  * plugins, or to add support for custom response sections, again added to Elasticsearch through plugins.
  */
 public class RestHighLevelClient implements Closeable {
+    private static final Consumer<SafeRequest> NOT_CUSTOMIZED = r -> {};
 
     private final RestClient client;
     private final NamedXContentRegistry registry;
@@ -271,6 +271,28 @@ public final SnapshotClient snapshot() {
      *
      * See <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html">Bulk API on elastic.co</a>
      */
+    public final BulkResponse bulk(BulkRequest bulkRequest) throws IOException {
+        return performRequestAndParseEntity(bulkRequest, RequestConverters::bulk, NOT_CUSTOMIZED,
+                BulkResponse::fromXContent, emptySet());
+    }
+
+    /**
+     * Executes a bulk request using the Bulk API
+     *
+     * See <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html">Bulk API on elastic.co</a>
+     */
+    public final BulkResponse bulk(BulkRequest bulkRequest, Consumer<SafeRequest> customizer) throws IOException {
+        return performRequestAndParseEntity(bulkRequest, RequestConverters::bulk, customizer,
+                BulkResponse::fromXContent, emptySet());
+    }
+
+    /**
+     * Executes a bulk request using the Bulk API
+     *
+     * See <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html">Bulk API on elastic.co</a>
+     * @deprecated Prefer {@link #bulk(BulkRequest, Consumer)}
+     */
+    @Deprecated
     public final BulkResponse bulk(BulkRequest bulkRequest, Header... headers) throws IOException {
         return performRequestAndParseEntity(bulkRequest, RequestConverters::bulk, BulkResponse::fromXContent, emptySet(), headers);
     }
@@ -280,6 +302,27 @@ public final BulkResponse bulk(BulkRequest bulkRequest, Header... headers) throw
      *
      * See <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html">Bulk API on elastic.co</a>
      */
+    public final void bulkAsync(BulkRequest bulkRequest, ActionListener<BulkResponse> listener) {
+        performRequestAsyncAndParseEntity(bulkRequest, RequestConverters::bulk, BulkResponse::fromXContent, listener, emptySet());
+    }
+
+    /**
+     * Asynchronously executes a bulk request using the Bulk API
+     *
+     * See <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html">Bulk API on elastic.co</a>
+     */
+    public final void bulkAsync(BulkRequest bulkRequest, Consumer<SafeRequest> customizer, ActionListener<BulkResponse> listener) {
+        performRequestAsyncAndParseEntity(bulkRequest, RequestConverters::bulk, customizer,
+                BulkResponse::fromXContent, listener, emptySet());
+    }
+
+    /**
+     * Asynchronously executes a bulk request using the Bulk API
+     *
+     * See <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html">Bulk API on elastic.co</a>
+     * @deprecated Prefer {@link #bulkAsync(BulkRequest, Consumer, ActionListener)}
+     */
+    @Deprecated
     public final void bulkAsync(BulkRequest bulkRequest, ActionListener<BulkResponse> listener, Header... headers) {
         performRequestAsyncAndParseEntity(bulkRequest, RequestConverters::bulk, BulkResponse::fromXContent, listener, emptySet(), headers);
     }
@@ -576,23 +619,46 @@ public final void fieldCapsAsync(FieldCapabilitiesRequest fieldCapabilitiesReque
             FieldCapabilitiesResponse::fromXContent, listener, emptySet(), headers);
     }
 
+    @Deprecated
+    @SuppressWarnings("overloads") // The overload *should* be ok. And one of them is deprecated.
     protected final <Req extends ActionRequest, Resp> Resp performRequestAndParseEntity(Req request,
                                                                             CheckedFunction<Req, Request, IOException> requestConverter,
                                                                             CheckedFunction<XContentParser, Resp, IOException> entityParser,
                                                                             Set<Integer> ignores, Header... headers) throws IOException {
         return performRequest(request, requestConverter, (response) -> parseEntity(response.getEntity(), entityParser), ignores, headers);
     }
 
+    @SuppressWarnings("overloads") // The overload *should* be ok. And one of them is deprecated.
+    protected final <Req extends ActionRequest, Resp> Resp performRequestAndParseEntity(Req request,
+                                                                            CheckedFunction<Req, Request, IOException> requestConverter,
+                                                                            Consumer<SafeRequest> requestCustomizer,
+                                                                            CheckedFunction<XContentParser, Resp, IOException> entityParser,
+                                                                            Set<Integer> ignores) throws IOException {
+        return performRequest(request, requestConverter, requestCustomizer,
+                response -> parseEntity(response.getEntity(), entityParser), ignores);
+    }
+
+    @Deprecated
+    @SuppressWarnings("overloads") // The overload *should* be ok. And one of them is deprecated.
     protected final <Req extends ActionRequest, Resp> Resp performRequest(Req request,
                                                           CheckedFunction<Req, Request, IOException> requestConverter,
                                                           CheckedFunction<Response, Resp, IOException> responseConverter,
                                                           Set<Integer> ignores, Header... headers) throws IOException {
+        return performRequest(request, requestConverter, r -> r.setHeaders(headers), responseConverter, ignores);
+    }
+
+    @SuppressWarnings("overloads") // The overload *should* be ok. And one of them is deprecated.
+    protected final <Req extends ActionRequest, Resp> Resp performRequest(Req request,
+                                                          CheckedFunction<Req, Request, IOException> requestConverter,
+                                                          Consumer<SafeRequest> requestCustomizer,
+                                                          CheckedFunction<Response, Resp, IOException> responseConverter,
+                                                          Set<Integer> ignores) throws IOException {
         ActionRequestValidationException validationException = request.validate();
         if (validationException != null) {
             throw validationException;
         }
         Request req = requestConverter.apply(request);
-        req.setHeaders(headers);
+        requestCustomizer.accept(new SafeRequest(req));
         Response response;
         try {
             response = client.performRequest(req);
@@ -618,6 +684,8 @@ protected final <Req extends ActionRequest, Resp> Resp performRequest(Req reques
         }
     }
 
+    @Deprecated
+    @SuppressWarnings("overloads") // The overload *should* be ok. And one of them is deprecated.
     protected final <Req extends ActionRequest, Resp> void performRequestAsyncAndParseEntity(Req request,
                                                                  CheckedFunction<Req, Request, IOException> requestConverter,
                                                                  CheckedFunction<XContentParser, Resp, IOException> entityParser,
@@ -626,10 +694,31 @@ protected final <Req extends ActionRequest, Resp> void performRequestAsyncAndPar
                 listener, ignores, headers);
     }
 
+    @SuppressWarnings("overloads") // The overload *should* be ok. And one of them is deprecated.
+    protected final <Req extends ActionRequest, Resp> void performRequestAsyncAndParseEntity(Req request,
+                                                                 CheckedFunction<Req, Request, IOException> requestConverter,
+                                                                 Consumer<SafeRequest> requestCustomizer,
+                                                                 CheckedFunction<XContentParser, Resp, IOException> entityParser,
+                                                                 ActionListener<Resp> listener, Set<Integer> ignores) {
+        performRequestAsync(request, requestConverter, requestCustomizer,
+                response -> parseEntity(response.getEntity(), entityParser), listener, ignores);
+    }
+
+    @Deprecated
+    @SuppressWarnings("overloads") // The overload *should* be ok. And one of them is deprecated.
     protected final <Req extends ActionRequest, Resp> void performRequestAsync(Req request,
                                                                CheckedFunction<Req, Request, IOException> requestConverter,
                                                                CheckedFunction<Response, Resp, IOException> responseConverter,
                                                                ActionListener<Resp> listener, Set<Integer> ignores, Header... headers) {
+        performRequestAsync(request, requestConverter, r -> r.setHeaders(headers), responseConverter, listener, ignores);
+    }
+
+    @SuppressWarnings("overloads") // The overload *should* be ok. And one of them is deprecated.
+    protected final <Req extends ActionRequest, Resp> void performRequestAsync(Req request,
+                                                               CheckedFunction<Req, Request, IOException> requestConverter,
+                                                               Consumer<SafeRequest> requestCustomizer,
+                                                               CheckedFunction<Response, Resp, IOException> responseConverter,
+                                                               ActionListener<Resp> listener, Set<Integer> ignores) {
         ActionRequestValidationException validationException = request.validate();
         if (validationException != null) {
             listener.onFailure(validationException);
@@ -638,11 +727,11 @@ protected final <Req extends ActionRequest, Resp> void performRequestAsync(Req r
         Request req;
         try {
             req = requestConverter.apply(request);
+            requestCustomizer.accept(new SafeRequest(req));
         } catch (Exception e) {
             listener.onFailure(e);
             return;
         }
-        req.setHeaders(headers);
 
         ResponseListener responseListener = wrapResponseListener(responseConverter, listener, ignores);
         client.performRequestAsync(req, responseListener);

client/rest-high-level/src/main/java/org/elasticsearch/client/SafeRequest.java

@@ -0,0 +1,57 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you 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.elasticsearch.client;
+
+import org.apache.http.Header;
+import org.apache.http.nio.protocol.HttpAsyncResponseConsumer;
+
+/**
+ * Wrapper around {@link Request} that exposes things that may be customized
+ * by clients of the {@link RestHighLevelClient} without modifying the behavior
+ * of the requests.
+ */
+public class SafeRequest {
+    private final Request request;
+
+    /**
+     * Build the {@linkplain SafeRequest}. Package private because it is only
+     * called by {@link RestHighLevelClient}.
+     */
+    SafeRequest(Request request) {
+        this.request = request;
+    }
+
+    /**
+     * Set the headers to attach to the request.
+     */
+    public void setHeaders(Header... headers) {
+        request.setHeaders(headers);
+    }
+
+    /**
+     * set the {@link HttpAsyncResponseConsumerFactory} used to create one
+     * {@link HttpAsyncResponseConsumer} callback per retry. Controls how the
+     * response body gets streamed from a non-blocking HTTP connection on the
+     * client side.
+     */
+    public void setHttpAsyncResponseConsumerFactory(HttpAsyncResponseConsumerFactory httpAsyncResponseConsumerFactory) {
+        request.setHttpAsyncResponseConsumerFactory(httpAsyncResponseConsumerFactory);
+    }
+}