Skip to content

Commit

Permalink
custom generic response, checking in transport base
Browse files Browse the repository at this point in the history
  • Loading branch information
l-trotta committed Oct 16, 2024
1 parent 435d2ef commit 0beeebb
Show file tree
Hide file tree
Showing 5 changed files with 104 additions and 65 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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) {
Expand Down Expand Up @@ -90,7 +102,7 @@ public Function<List<String>, Boolean> onWarnings() {

@Override
public boolean keepResponseBodyOnException() {
return false;
return keepResponseBodyOnException;
}

@Override
Expand Down Expand Up @@ -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();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
@@ -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();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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);
}

Expand All @@ -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));
}

Expand Down Expand Up @@ -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;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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()));

Expand All @@ -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);
Expand All @@ -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());
}
}
}

0 comments on commit 0beeebb

Please sign in to comment.