custom generic response, checking in transport base

Author: l-trotta

java-client/src/main/java/co/elastic/clients/transport/DefaultTransportOptions.java

@@ -37,13 +37,24 @@ public class DefaultTransportOptions implements TransportOptions {
     private final HeaderMap headers;
     private final Map<String, String> parameters;
     private final Function<List<String>, Boolean> onWarnings;
+    private boolean keepResponseBodyOnException;
 
     public static final DefaultTransportOptions EMPTY = new DefaultTransportOptions();
 
     public DefaultTransportOptions() {
         this(new HeaderMap(), Collections.emptyMap(), null);
     }
 
+    public DefaultTransportOptions(
+        @Nullable HeaderMap headers,
+        @Nullable Map<String, String> parameters,
+        @Nullable Function<List<String>, Boolean> onWarnings,
+        boolean keepResponseBodyOnException
+    ) {
+        this(headers,parameters,onWarnings);
+        this.keepResponseBodyOnException = keepResponseBodyOnException;
+    }
+
     public DefaultTransportOptions(
         @Nullable HeaderMap headers,
         @Nullable Map<String, String> parameters,
@@ -53,10 +64,11 @@ public DefaultTransportOptions(
         this.parameters = (parameters == null || parameters.isEmpty()) ?
             Collections.emptyMap() : Collections.unmodifiableMap(parameters);
         this.onWarnings = onWarnings;
+        this.keepResponseBodyOnException = false;
     }
 
     protected DefaultTransportOptions(AbstractBuilder<?> builder) {
-        this(builder.headers, builder.parameters, builder.onWarnings);
+        this(builder.headers, builder.parameters, builder.onWarnings, builder.keepResponseBodyOnException);
     }
 
     public static DefaultTransportOptions of(@Nullable TransportOptions options) {
@@ -90,7 +102,7 @@ public Function<List<String>, Boolean> onWarnings() {
 
     @Override
     public boolean keepResponseBodyOnException() {
-        return false;
+        return keepResponseBodyOnException;
     }
 
     @Override
@@ -125,13 +137,13 @@ public AbstractBuilder(DefaultTransportOptions options) {
             this.headers = new HeaderMap(options.headers);
             this.parameters = copyOrNull(options.parameters);
             this.onWarnings = options.onWarnings;
-            this.keepResponseBodyOnException = options.keepResponseBodyOnException();
+            this.keepResponseBodyOnException = options.keepResponseBodyOnException;
         }
 
         protected abstract BuilderT self();
 
         @Override
-        public BuilderT keepResponseBodyOnException(boolean value){
+        public BuilderT keepResponseBodyOnException(boolean value) {
             this.keepResponseBodyOnException = value;
             return self();
         }

java-client/src/main/java/co/elastic/clients/transport/ElasticsearchTransportBase.java

@@ -29,6 +29,7 @@
 import co.elastic.clients.transport.endpoints.BooleanEndpoint;
 import co.elastic.clients.transport.endpoints.BooleanResponse;
 import co.elastic.clients.transport.http.HeaderMap;
+import co.elastic.clients.transport.http.RepeatableBodyResponse;
 import co.elastic.clients.transport.http.TransportHttpClient;
 import co.elastic.clients.transport.instrumentation.Instrumentation;
 import co.elastic.clients.transport.instrumentation.NoopInstrumentation;
@@ -306,6 +307,9 @@ private <ResponseT, ErrorT> ResponseT getApiResponse(
 
         int statusCode = clientResp.statusCode();
 
+        if(options().keepResponseBodyOnException()){
+            clientResp = RepeatableBodyResponse.of(clientResp);
+        }
         try {
             if (statusCode == 200) {
                 checkProductHeader(clientResp, endpoint);

java-client/src/main/java/co/elastic/clients/transport/http/RepeatableBodyResponse.java

@@ -0,0 +1,65 @@
+package co.elastic.clients.transport.http;
+
+import co.elastic.clients.util.BinaryData;
+import co.elastic.clients.util.ByteArrayBinaryData;
+
+import javax.annotation.Nullable;
+import java.io.IOException;
+import java.util.List;
+
+public class RepeatableBodyResponse implements TransportHttpClient.Response {
+
+    private final TransportHttpClient.Response response;
+    private final BinaryData body;
+
+    public static TransportHttpClient.Response of(TransportHttpClient.Response response) throws IOException {
+        BinaryData body = response.body();
+        if (body == null || body.isRepeatable()) {
+            return response;
+        }
+        return new RepeatableBodyResponse(response);
+    }
+
+    public RepeatableBodyResponse(TransportHttpClient.Response response) throws IOException {
+        this.response = response;
+        this.body = new ByteArrayBinaryData(response.body());
+    }
+
+    @Override
+    public TransportHttpClient.Node node() {
+        return response.node();
+    }
+
+    @Override
+    public int statusCode() {
+        return response.statusCode();
+    }
+
+    @Nullable
+    @Override
+    public String header(String name) {
+        return response.header(name);
+    }
+
+    @Override
+    public List<String> headers(String name) {
+        return response.headers(name);
+    }
+
+    @Nullable
+    @Override
+    public BinaryData body() throws IOException {
+        return this.body;
+    }
+
+    @Nullable
+    @Override
+    public Object originalResponse() {
+        return response.originalResponse();
+    }
+
+    @Override
+    public void close() throws IOException {
+        response.close();
+    }
+}

java-client/src/main/java/co/elastic/clients/transport/rest_client/RestClientHttpClient.java

@@ -23,7 +23,6 @@
 import co.elastic.clients.transport.http.HeaderMap;
 import co.elastic.clients.transport.http.TransportHttpClient;
 import co.elastic.clients.util.BinaryData;
-import co.elastic.clients.util.ByteArrayBinaryData;
 import co.elastic.clients.util.NoCopyByteArrayOutputStream;
 import org.apache.http.Header;
 import org.apache.http.HeaderElement;
@@ -35,10 +34,8 @@
 import org.elasticsearch.client.RestClient;
 
 import javax.annotation.Nullable;
-import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.InputStream;
-import java.io.InputStreamReader;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
 import java.util.AbstractList;
@@ -93,9 +90,6 @@ public Response performRequest(String endpointId, @Nullable Node node, Request r
         RestClientOptions rcOptions = RestClientOptions.of(options);
         org.elasticsearch.client.Request restRequest = createRestRequest(request, rcOptions);
         org.elasticsearch.client.Response restResponse = restClient.performRequest(restRequest);
-        if (options.keepResponseBodyOnException()) {
-            return new RepeatableBodyResponse(restResponse);
-        }
         return new RestResponse(restResponse);
     }
 
@@ -119,9 +113,6 @@ public CompletableFuture<Response> performRequestAsync(
         future.cancellable = restClient.performRequestAsync(restRequest, new ResponseListener() {
             @Override
             public void onSuccess(org.elasticsearch.client.Response response) {
-                if (options.keepResponseBodyOnException()) {
-                    future.complete(new RepeatableBodyResponse(response));
-                }
                 future.complete(new RestResponse(response));
             }
 
@@ -251,51 +242,6 @@ public void close() throws IOException {
         }
     }
 
-    public class RepeatableBodyResponse extends RestResponse {
-
-        BinaryData repeatableBody;
-
-        RepeatableBodyResponse(org.elasticsearch.client.Response restResponse) {
-            super(restResponse);
-        }
-
-        @Nullable
-        @Override
-        public BinaryData body() throws IOException {
-            if(repeatableBody != null) {
-                return repeatableBody;
-            }
-            BinaryData body = super.body();
-            if (body != null) {
-                if(body.isRepeatable()){
-                    repeatableBody = body;
-                }
-                else{
-                    repeatableBody = new ByteArrayBinaryData(body);
-                }
-            }
-            return repeatableBody;
-        }
-
-        public String getOriginalBodyAsString() throws IOException {
-            BinaryData body = body();
-
-            if (body != null) {
-                StringBuilder sb = new StringBuilder();
-                BufferedReader br = new BufferedReader(new InputStreamReader(body.asInputStream()));
-                String read;
-
-                while ((read = br.readLine()) != null) {
-                    sb.append(read);
-                }
-                br.close();
-                return sb.toString();
-            }
-            return null;
-        }
-
-    }
-
     private static class HttpEntityBinaryData implements BinaryData {
         private final HttpEntity entity;
 

java-client/src/test/java/co/elastic/clients/transport/TransportTest.java

@@ -21,9 +21,10 @@
 
 import co.elastic.clients.elasticsearch.ElasticsearchClient;
 import co.elastic.clients.json.jackson.JacksonJsonpMapper;
-import co.elastic.clients.transport.rest_client.RestClientHttpClient;
+import co.elastic.clients.transport.http.RepeatableBodyResponse;
 import co.elastic.clients.transport.rest_client.RestClientOptions;
 import co.elastic.clients.transport.rest_client.RestClientTransport;
+import co.elastic.clients.util.BinaryData;
 import com.sun.net.httpserver.HttpServer;
 import org.apache.http.HttpHost;
 import org.elasticsearch.client.RequestOptions;
@@ -32,6 +33,8 @@
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
+import java.io.BufferedReader;
+import java.io.InputStreamReader;
 import java.io.OutputStream;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
@@ -106,7 +109,7 @@ public void testOriginalJsonBodyRetrievalException() throws Exception {
             .builder(new HttpHost(address.getHostString(), address.getPort(), "http"))
             .build();
 
-        // no transport options, should throw TransportException, but original body cannot be retrieved
+        // no transport options, response is not RepeatableBodyResponse, original body cannot be retrieved
         ElasticsearchClient esClient = new ElasticsearchClient(new RestClientTransport(restClient,
             new JacksonJsonpMapper()));
 
@@ -116,7 +119,7 @@ public void testOriginalJsonBodyRetrievalException() throws Exception {
         );
 
         assertEquals(200, ex.statusCode());
-        assertNotEquals(RestClientHttpClient.RepeatableBodyResponse.class, ex.response().getClass());
+        assertNotEquals(RepeatableBodyResponse.class, ex.response().getClass());
 
         // setting transport option
         RestClientOptions options = new RestClientOptions(RequestOptions.DEFAULT, true);
@@ -134,10 +137,19 @@ public void testOriginalJsonBodyRetrievalException() throws Exception {
         httpServer.stop(0);
 
         assertEquals(200, ex.statusCode());
-        assertEquals(RestClientHttpClient.RepeatableBodyResponse.class, ex.response().getClass());
-
-        try (RestClientHttpClient.RepeatableBodyResponse repeatableResponse = (RestClientHttpClient.RepeatableBodyResponse) ex.response()){
-            assertEquals("definitely not json",repeatableResponse.getOriginalBodyAsString());
+        assertEquals(RepeatableBodyResponse.class, ex.response().getClass());
+
+        try (RepeatableBodyResponse repeatableResponse = (RepeatableBodyResponse) ex.response()){
+            BinaryData body = repeatableResponse.body();
+                StringBuilder sb = new StringBuilder();
+                BufferedReader br = new BufferedReader(new InputStreamReader(body.asInputStream()));
+                String read;
+
+                while ((read = br.readLine()) != null) {
+                    sb.append(read);
+                }
+                br.close();
+            assertEquals("definitely not json",sb.toString());
         }
     }
 }