Add support for named xcontent parsers to high level REST client

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

@@ -41,22 +41,46 @@
 import org.elasticsearch.rest.RestStatus;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 import java.util.Objects;
 import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Stream;
 
 import static java.util.Collections.emptySet;
 import static java.util.Collections.singleton;
+import static java.util.stream.Collectors.toList;
 
 /**
  * High level REST client that wraps an instance of the low level {@link RestClient} and allows to build requests and read responses.
  * The provided {@link RestClient} is externally built and closed.
+ * Can be sub-classed to expose additional client methods that make use of endpoints added to Elasticsearch through plugins, or to
+ * add support for custom response sections, again added to Elasticsearch through plugins.
  */
 public class RestHighLevelClient {
 
     private final RestClient client;
+    private final NamedXContentRegistry registry;
 
-    public RestHighLevelClient(RestClient client) {
-        this.client = Objects.requireNonNull(client);
+    /**
+     * Creates a {@link RestHighLevelClient} given the low level {@link RestClient} that it should use to perform requests.
+     */
+    public RestHighLevelClient(RestClient restClient) {
+        this(restClient, Collections.emptyList());
+    }
+
+    /**
+     * Creates a {@link RestHighLevelClient} given the low level {@link RestClient} that it should use to perform requests and
+     * a list of entries that allow to parse custom response sections added to Elasticsearch through plugins.
+     */
+    protected RestHighLevelClient(RestClient restClient, List<NamedXContentRegistry.Entry> namedXContentEntries) {
+        this.client = Objects.requireNonNull(restClient);
+        this.registry = new NamedXContentRegistry(Stream.of(
+                getNamedXContents().stream(),
+                namedXContentEntries.stream()
+        ).flatMap(Function.identity()).collect(toList()));
     }
 
     /**
@@ -204,7 +228,7 @@ <Req extends ActionRequest, Resp> void performRequestAsync(Req request,
         client.performRequestAsync(req.method, req.endpoint, req.params, req.entity, responseListener, headers);
     }
 
-    static <Resp> ResponseListener wrapResponseListener(CheckedFunction<Response, Resp, IOException> responseConverter,
+    <Resp> ResponseListener wrapResponseListener(CheckedFunction<Response, Resp, IOException> responseConverter,
                                                         ActionListener<Resp> actionListener, Set<Integer> ignores) {
         return new ResponseListener() {
             @Override
@@ -249,7 +273,7 @@ public void onFailure(Exception exception) {
      * that wraps the original {@link ResponseException}. The potential exception obtained while parsing is added to the returned
      * exception as a suppressed exception. This method is guaranteed to not throw any exception eventually thrown while parsing.
      */
-    static ElasticsearchStatusException parseResponseException(ResponseException responseException) {
+    ElasticsearchStatusException parseResponseException(ResponseException responseException) {
         Response response = responseException.getResponse();
         HttpEntity entity = response.getEntity();
         ElasticsearchStatusException elasticsearchException;
@@ -269,7 +293,7 @@ static ElasticsearchStatusException parseResponseException(ResponseException res
         return elasticsearchException;
     }
 
-    static <Resp> Resp parseEntity(
+    <Resp> Resp parseEntity(
             HttpEntity entity, CheckedFunction<XContentParser, Resp, IOException> entityParser) throws IOException {
         if (entity == null) {
             throw new IllegalStateException("Response body expected but not returned");
@@ -281,12 +305,18 @@ static <Resp> Resp parseEntity(
         if (xContentType == null) {
             throw new IllegalStateException("Unsupported Content-Type: " + entity.getContentType().getValue());
         }
-        try (XContentParser parser = xContentType.xContent().createParser(NamedXContentRegistry.EMPTY, entity.getContent())) {
+        try (XContentParser parser = xContentType.xContent().createParser(registry, entity.getContent())) {
             return entityParser.apply(parser);
         }
     }
 
     static boolean convertExistsResponse(Response response) {
         return response.getStatusLine().getStatusCode() == 200;
     }
+
+    static List<NamedXContentRegistry.Entry> getNamedXContents() {
+        List<NamedXContentRegistry.Entry> namedXContents = new ArrayList<>();
+        //namedXContents.add(new NamedXContentRegistry.Entry(Aggregation.class, new ParseField("sterms"), StringTerms::fromXContent));
+        return namedXContents;
+    }
 }

client/rest-high-level/src/test/java/org/elasticsearch/client/RestHighLevelClientExtTests.java

@@ -0,0 +1,138 @@
+/*
+ * 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.HttpEntity;
+import org.apache.http.entity.ContentType;
+import org.apache.http.entity.StringEntity;
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.xcontent.NamedXContentRegistry;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.test.ESTestCase;
+import org.junit.Before;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.hamcrest.CoreMatchers.instanceOf;
+import static org.mockito.Mockito.mock;
+
+/**
+ * This test works against a {@link RestHighLevelClient} subclass that simulats how custom response sections returned by
+ * Elasticsearch plugins can be parsed using the high level client.
+ */
+public class RestHighLevelClientExtTests extends ESTestCase {
+
+    private RestHighLevelClient restHighLevelClient;
+
+    @Before
+    public void initClient() throws IOException {
+        RestClient restClient = mock(RestClient.class);
+        restHighLevelClient = new RestHighLevelClientExt(restClient);
+    }
+
+    public void testParseEntityCustomResponseSection() throws IOException {
+        {
+            HttpEntity jsonEntity = new StringEntity("{\"custom1\":{ \"field\":\"value\"}}", ContentType.APPLICATION_JSON);
+            BaseCustomResponseSection customSection = restHighLevelClient.parseEntity(jsonEntity, BaseCustomResponseSection::fromXContent);
+            assertThat(customSection, instanceOf(CustomResponseSection1.class));
+            CustomResponseSection1 customResponseSection1 = (CustomResponseSection1) customSection;
+            assertEquals("value", customResponseSection1.value);
+        }
+        {
+            HttpEntity jsonEntity = new StringEntity("{\"custom2\":{ \"array\": [\"item1\", \"item2\"]}}", ContentType.APPLICATION_JSON);
+            BaseCustomResponseSection customSection = restHighLevelClient.parseEntity(jsonEntity, BaseCustomResponseSection::fromXContent);
+            assertThat(customSection, instanceOf(CustomResponseSection2.class));
+            CustomResponseSection2 customResponseSection2 = (CustomResponseSection2) customSection;
+            assertArrayEquals(new String[]{"item1", "item2"}, customResponseSection2.values);
+        }
+    }
+
+    private static class RestHighLevelClientExt extends RestHighLevelClient {
+
+        private RestHighLevelClientExt(RestClient restClient) {
+            super(restClient, getNamedXContentsExt());
+        }
+
+        private static List<NamedXContentRegistry.Entry> getNamedXContentsExt() {
+            List<NamedXContentRegistry.Entry> entries = new ArrayList<>();
+            entries.add(new NamedXContentRegistry.Entry(BaseCustomResponseSection.class, new ParseField("custom1"),
+                    CustomResponseSection1::fromXContent));
+            entries.add(new NamedXContentRegistry.Entry(BaseCustomResponseSection.class, new ParseField("custom2"),
+                    CustomResponseSection2::fromXContent));
+            return entries;
+        }
+    }
+
+    private abstract static class BaseCustomResponseSection {
+
+        static BaseCustomResponseSection fromXContent(XContentParser parser) throws IOException {
+            assertEquals(XContentParser.Token.START_OBJECT, parser.nextToken());
+            assertEquals(XContentParser.Token.FIELD_NAME, parser.nextToken());
+            BaseCustomResponseSection custom = parser.namedObject(BaseCustomResponseSection.class, parser.currentName(), null);
+            assertEquals(XContentParser.Token.END_OBJECT, parser.nextToken());
+            return custom;
+        }
+    }
+
+    private static class CustomResponseSection1 extends BaseCustomResponseSection {
+
+        private final String value;
+
+        private CustomResponseSection1(String value) {
+            this.value = value;
+        }
+
+        static CustomResponseSection1 fromXContent(XContentParser parser) throws IOException {
+            assertEquals(XContentParser.Token.START_OBJECT, parser.nextToken());
+            assertEquals(XContentParser.Token.FIELD_NAME, parser.nextToken());
+            assertEquals("field", parser.currentName());
+            assertEquals(XContentParser.Token.VALUE_STRING, parser.nextToken());
+            CustomResponseSection1 responseSection1 = new CustomResponseSection1(parser.text());
+            assertEquals(XContentParser.Token.END_OBJECT, parser.nextToken());
+            return responseSection1;
+        }
+    }
+
+    private static class CustomResponseSection2 extends BaseCustomResponseSection {
+
+        private final String[] values;
+
+        private CustomResponseSection2(String[] values) {
+            this.values = values;
+        }
+
+        static CustomResponseSection2 fromXContent(XContentParser parser) throws IOException {
+            assertEquals(XContentParser.Token.START_OBJECT, parser.nextToken());
+            assertEquals(XContentParser.Token.FIELD_NAME, parser.nextToken());
+            assertEquals("array", parser.currentName());
+            assertEquals(XContentParser.Token.START_ARRAY, parser.nextToken());
+            List<String> values = new ArrayList<>();
+            while(parser.nextToken().isValue()) {
+                values.add(parser.text());
+            }
+            assertEquals(XContentParser.Token.END_ARRAY, parser.currentToken());
+            CustomResponseSection2 responseSection2 = new CustomResponseSection2(values.toArray(new String[values.size()]));
+            assertEquals(XContentParser.Token.END_OBJECT, parser.nextToken());
+            return responseSection2;
+        }
+    }
+}