diff --git a/src/main/java/ai/tecton/client/transport/TectonHttpClient.java b/src/main/java/ai/tecton/client/transport/TectonHttpClient.java index d0330c90..ed6dbfc0 100644 --- a/src/main/java/ai/tecton/client/transport/TectonHttpClient.java +++ b/src/main/java/ai/tecton/client/transport/TectonHttpClient.java @@ -17,7 +17,6 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; -import java.util.stream.Collectors; import okhttp3.Call; import okhttp3.Callback; import okhttp3.ConnectionPool; @@ -32,7 +31,6 @@ import org.apache.commons.lang3.Validate; public class TectonHttpClient { - private HttpUrl url; private final String apiKey; private final OkHttpClient client; @@ -96,7 +94,7 @@ public boolean isClosed() { public HttpResponse performRequest(String endpoint, HttpMethod method, String requestBody) { HttpRequest httpRequest = new HttpRequest(url.url().toString(), endpoint, method, apiKey, requestBody); - Request request = buildRequestWithDefaultHeaders(httpRequest); + Request request = buildRequestWithDefaultHeaders(httpRequest, 0); Call call = client.newCall(request); try (Response response = call.execute()) { return new HttpResponse(response); @@ -113,13 +111,12 @@ public List performParallelRequests( // Map request body to OkHttp Request // ordering of requests is maintained - List requestList = - requestBodyList.stream() - .map( - requestBody -> - new HttpRequest(url.url().toString(), endpoint, method, apiKey, requestBody)) - .map(this::buildRequestWithDefaultHeaders) - .collect(Collectors.toList()); + List requestList = new ArrayList<>(); + for (int i = 0; i < requestBodyList.size(); i++) { + HttpRequest httpRequest = + new HttpRequest(url.url().toString(), endpoint, method, apiKey, requestBodyList.get(i)); + requestList.add(buildRequestWithDefaultHeaders(httpRequest, i)); + } // Initialize a countdown latch for numberOfCalls. CountDownLatch countDownLatch = new CountDownLatch(requestBodyList.size()); @@ -141,7 +138,7 @@ public void onResponse(Call call, Response response) { try (ResponseBody responseBody = response.body()) { // Add response to corresponding index parallelCallHandler.set( - requestList.indexOf(call.request()), new HttpResponse(response, responseBody)); + (Integer) call.request().tag(), new HttpResponse(response, responseBody)); } catch (Exception e) { throw new TectonServiceException(e.getMessage()); } finally { @@ -171,7 +168,7 @@ public void onResponse(Call call, Response response) { } } - public Request buildRequestWithDefaultHeaders(HttpRequest httpRequest) { + public Request buildRequestWithDefaultHeaders(HttpRequest httpRequest, int index) { // Construct url Request.Builder requestBuilder = new Request.Builder().url(httpRequest.getUrl()); @@ -183,16 +180,10 @@ public Request buildRequestWithDefaultHeaders(HttpRequest httpRequest) { requestBuilder.header(HttpHeader.AUTHORIZATION.getName(), apiKeyHeader); // Add request body - switch (httpRequest.getMethod()) { - case POST: - default: - { - okhttp3.MediaType mediaType = - okhttp3.MediaType.parse(MediaType.APPLICATION_JSON.getName()); - RequestBody requestBody = RequestBody.create(httpRequest.getJsonBody(), mediaType); - requestBuilder.post(requestBody); - } - } + okhttp3.MediaType mediaType = okhttp3.MediaType.parse(MediaType.APPLICATION_JSON.getName()); + RequestBody requestBody = RequestBody.create(httpRequest.getJsonBody(), mediaType); + requestBuilder.post(requestBody); + requestBuilder.tag(index); return requestBuilder.build(); } diff --git a/src/test/java/ai/tecton/client/transport/TectonHttpClientTest.java b/src/test/java/ai/tecton/client/transport/TectonHttpClientTest.java index 034c8fe7..1d95b2f6 100644 --- a/src/test/java/ai/tecton/client/transport/TectonHttpClientTest.java +++ b/src/test/java/ai/tecton/client/transport/TectonHttpClientTest.java @@ -75,7 +75,7 @@ public void testDefaultTectonRequest() { HttpRequest httpRequest = new HttpRequest(url, endpoint, TectonHttpClient.HttpMethod.POST, apiKey, body); - Request request = httpClient.buildRequestWithDefaultHeaders(httpRequest); + Request request = httpClient.buildRequestWithDefaultHeaders(httpRequest, 0); Assert.assertNotNull(request); Assert.assertEquals(POST, request.method());