Enable support for compressed response within RestHighLevelClient (#63087)

Co-authored-by: Hakky54 <[email protected]>

Author: swallez

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

@@ -21,6 +21,7 @@
 
 import org.apache.http.Header;
 import org.apache.http.HttpEntity;
+import org.apache.http.client.entity.GzipDecompressingEntity;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.ActionListener;
@@ -2037,11 +2038,18 @@ protected final ElasticsearchStatusException parseResponseException(ResponseExce
         return elasticsearchException;
     }
 
-    protected final <Resp> Resp parseEntity(final HttpEntity entity,
+    protected final <Resp> Resp parseEntity(final HttpEntity httpEntity,
                                       final CheckedFunction<XContentParser, Resp, IOException> entityParser) throws IOException {
-        if (entity == null) {
+        if (httpEntity == null) {
             throw new IllegalStateException("Response body expected but not returned");
         }
+
+        final HttpEntity entity = Optional.ofNullable(httpEntity.getContentEncoding())
+            .map(Header::getValue)
+            .filter("gzip"::equalsIgnoreCase)
+            .map(gzipHeaderValue -> (HttpEntity) new GzipDecompressingEntity(httpEntity))
+            .orElse(httpEntity);
+
         if (entity.getContentType() == null) {
             throw new IllegalStateException("Elasticsearch didn't return the [Content-Type] header, unable to parse response body");
         }

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

@@ -0,0 +1,54 @@
+/*
+ * 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.HttpHeaders;
+import org.apache.http.client.methods.HttpPost;
+import org.apache.http.client.methods.HttpPut;
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.action.search.SearchResponse;
+
+import java.io.IOException;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class HighLevelRestClientCompressionIT extends ESRestHighLevelClientTestCase {
+
+    private static final String GZIP_ENCODING = "gzip";
+    private static final String SAMPLE_DOCUMENT = "{\"name\":{\"first name\":\"Steve\",\"last name\":\"Jobs\"}}";
+
+    public void testCompressesResponseIfRequested() throws IOException {
+        Request doc = new Request(HttpPut.METHOD_NAME, "/company/_doc/1");
+        doc.setJsonEntity(SAMPLE_DOCUMENT);
+        client().performRequest(doc);
+        client().performRequest(new Request(HttpPost.METHOD_NAME, "/_refresh"));
+
+        RequestOptions.Builder requestOptionsBuilder = RequestOptions.DEFAULT.toBuilder();
+        requestOptionsBuilder.addHeader(HttpHeaders.ACCEPT_ENCODING, GZIP_ENCODING);
+        RequestOptions requestOptions = requestOptionsBuilder.build();
+
+        SearchRequest searchRequest = new SearchRequest("company");
+        SearchResponse searchResponse = execute(searchRequest, highLevelClient()::search, highLevelClient()::searchAsync, requestOptions);
+
+        assertThat(searchResponse.status().getStatus(), equalTo(200));
+        assertEquals(1L, searchResponse.getHits().getTotalHits());
+        assertEquals(SAMPLE_DOCUMENT, searchResponse.getHits().getHits()[0].getSourceAsString());
+    }
+
+}

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

@@ -34,6 +34,7 @@
 import org.apache.http.message.BasicStatusLine;
 import org.apache.http.nio.entity.NByteArrayEntity;
 import org.apache.http.nio.entity.NStringEntity;
+import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.Build;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.Version;
@@ -89,10 +90,12 @@
 import org.hamcrest.Matchers;
 import org.junit.Before;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.net.SocketTimeoutException;
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -106,6 +109,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import java.util.zip.GZIPOutputStream;
 
 import static org.elasticsearch.common.xcontent.XContentHelper.toXContent;
 import static org.hamcrest.CoreMatchers.endsWith;
@@ -273,6 +277,59 @@ public void testParseEntity() throws IOException {
         }
     }
 
+    public void testParseCompressedEntity() throws IOException {
+        CheckedFunction<XContentParser, String, IOException> entityParser = parser -> {
+            assertEquals(XContentParser.Token.START_OBJECT, parser.nextToken());
+            assertEquals(XContentParser.Token.FIELD_NAME, parser.nextToken());
+            assertTrue(parser.nextToken().isValue());
+            String value = parser.text();
+            assertEquals(XContentParser.Token.END_OBJECT, parser.nextToken());
+            return value;
+        };
+
+        HttpEntity jsonEntity = createGzipEncodedEntity("{\"field\":\"value\"}", ContentType.APPLICATION_JSON);
+        assertEquals("value", restHighLevelClient.parseEntity(jsonEntity, entityParser));
+        HttpEntity yamlEntity = createGzipEncodedEntity("---\nfield: value\n", ContentType.create("application/yaml"));
+        assertEquals("value", restHighLevelClient.parseEntity(yamlEntity, entityParser));
+        HttpEntity smileEntity = createGzipEncodedEntity(SmileXContent.contentBuilder(), ContentType.create("application/smile"));
+        assertEquals("value", restHighLevelClient.parseEntity(smileEntity, entityParser));
+        HttpEntity cborEntity = createGzipEncodedEntity(CborXContent.contentBuilder(), ContentType.create("application/cbor"));
+        assertEquals("value", restHighLevelClient.parseEntity(cborEntity, entityParser));
+    }
+
+    private HttpEntity createGzipEncodedEntity(String content, ContentType contentType) throws IOException {
+        byte[] gzipEncodedContent = compressContentWithGzip(content.getBytes(StandardCharsets.UTF_8));
+        NByteArrayEntity httpEntity = new NByteArrayEntity(gzipEncodedContent, contentType);
+        httpEntity.setContentEncoding("gzip");
+
+        return httpEntity;
+    }
+
+    private HttpEntity createGzipEncodedEntity(XContentBuilder xContentBuilder, ContentType contentType) throws IOException {
+        try (XContentBuilder builder = xContentBuilder) {
+            builder.startObject();
+            builder.field("field", "value");
+            builder.endObject();
+
+            BytesRef bytesRef = BytesReference.bytes(xContentBuilder).toBytesRef();
+            byte[] gzipEncodedContent = compressContentWithGzip(bytesRef.bytes);
+            NByteArrayEntity httpEntity = new NByteArrayEntity(gzipEncodedContent, contentType);
+            httpEntity.setContentEncoding("gzip");
+
+            return httpEntity;
+        }
+    }
+
+    private static byte[] compressContentWithGzip(byte[] content) throws IOException {
+        ByteArrayOutputStream bos = new ByteArrayOutputStream(content.length);
+        GZIPOutputStream gzip = new GZIPOutputStream(bos);
+        gzip.write(content);
+        gzip.close();
+        bos.close();
+
+        return bos.toByteArray();
+    }
+
     private static HttpEntity createBinaryEntity(XContentBuilder xContentBuilder, ContentType contentType) throws IOException {
         try (XContentBuilder builder = xContentBuilder) {
             builder.startObject();
@@ -763,12 +820,12 @@ public void testApiNamingConventions() throws Exception {
                     Collectors.mapping(Tuple::v2, Collectors.toSet())));
 
         // TODO remove in 8.0 - we will undeprecate indices.get_template because the current getIndexTemplate
-        // impl will replace the existing getTemplate method. 
+        // impl will replace the existing getTemplate method.
         // The above general-purpose code ignores all deprecated methods which in this case leaves `getTemplate`
-        // looking like it doesn't have a valid implementatation when it does. 
+        // looking like it doesn't have a valid implementatation when it does.
         apiUnsupported.remove("indices.get_template");
-        
-        
+
+
 
         for (Map.Entry<String, Set<Method>> entry : methods.entrySet()) {
             String apiName = entry.getKey();
@@ -803,7 +860,7 @@ public void testApiNamingConventions() throws Exception {
                                 apiName.startsWith("index_lifecycle.") == false &&
                                 apiName.startsWith("ccr.") == false &&
                                 apiName.endsWith("freeze") == false &&
-                                // IndicesClientIT.getIndexTemplate should be renamed "getTemplate" in version 8.0 when we 
+                                // IndicesClientIT.getIndexTemplate should be renamed "getTemplate" in version 8.0 when we
                                 // can get rid of 7.0's deprecated "getTemplate"
                                 apiName.equals("indices.get_index_template") == false) {
                                 apiNotFound.add(apiName);