diff --git a/core/src/main/java/org/apache/iceberg/rest/HTTPClient.java b/core/src/main/java/org/apache/iceberg/rest/HTTPClient.java index b2c4981db924..f9e42486cab7 100644 --- a/core/src/main/java/org/apache/iceberg/rest/HTTPClient.java +++ b/core/src/main/java/org/apache/iceberg/rest/HTTPClient.java @@ -32,6 +32,7 @@ import org.apache.hc.client5.http.impl.classic.CloseableHttpResponse; import org.apache.hc.client5.http.impl.classic.HttpClients; import org.apache.hc.core5.http.ContentType; +import org.apache.hc.core5.http.Header; import org.apache.hc.core5.http.HttpHeaders; import org.apache.hc.core5.http.HttpStatus; import org.apache.hc.core5.http.Method; @@ -190,6 +191,35 @@ private T execute( Class responseType, Map headers, Consumer errorHandler) { + return execute( + method, path, queryParams, requestBody, responseType, headers, errorHandler, h -> {}); + } + + /** + * Method to execute an HTTP request and process the corresponding response. + * + * @param method - HTTP method, such as GET, POST, HEAD, etc. + * @param queryParams - A map of query parameters + * @param path - URL path to send the request to + * @param requestBody - Content to place in the request body + * @param responseType - Class of the Response type. Needs to have serializer registered with + * ObjectMapper + * @param errorHandler - Error handler delegated for HTTP responses which handles server error + * responses + * @param responseHeaders The consumer of the response headers + * @param - Class type of the response for deserialization. Must be registered with the + * ObjectMapper. + * @return The response entity, parsed and converted to its type T + */ + private T execute( + Method method, + String path, + Map queryParams, + Object requestBody, + Class responseType, + Map headers, + Consumer errorHandler, + Consumer> responseHeaders) { if (path.startsWith("/")) { throw new RESTException( "Received a malformed path for a REST request: %s. Paths should not start with /", path); @@ -210,6 +240,12 @@ private T execute( } try (CloseableHttpResponse response = httpClient.execute(request)) { + Map respHeaders = Maps.newHashMap(); + for (Header header : response.getHeaders()) { + respHeaders.put(header.getName(), header.getValue()); + } + + responseHeaders.accept(respHeaders); // Skip parsing the response stream for any successful request not expecting a response body if (response.getCode() == HttpStatus.SC_NO_CONTENT @@ -269,6 +305,18 @@ public T post( return execute(Method.POST, path, null, body, responseType, headers, errorHandler); } + @Override + public T post( + String path, + RESTRequest body, + Class responseType, + Map headers, + Consumer errorHandler, + Consumer> responseHeaders) { + return execute( + Method.POST, path, null, body, responseType, headers, errorHandler, responseHeaders); + } + @Override public T delete( String path, diff --git a/core/src/main/java/org/apache/iceberg/rest/RESTClient.java b/core/src/main/java/org/apache/iceberg/rest/RESTClient.java index 08bc0c5fa58d..0f17d9a127e2 100644 --- a/core/src/main/java/org/apache/iceberg/rest/RESTClient.java +++ b/core/src/main/java/org/apache/iceberg/rest/RESTClient.java @@ -112,6 +112,30 @@ default T post( return post(path, body, responseType, headers.get(), errorHandler); } + default T post( + String path, + RESTRequest body, + Class responseType, + Supplier> headers, + Consumer errorHandler, + Consumer> responseHeaders) { + return post(path, body, responseType, headers.get(), errorHandler, responseHeaders); + } + + default T post( + String path, + RESTRequest body, + Class responseType, + Map headers, + Consumer errorHandler, + Consumer> responseHeaders) { + if (null != responseHeaders) { + throw new UnsupportedOperationException("Returning response headers is not supported"); + } + + return post(path, body, responseType, headers, errorHandler); + } + T post( String path, RESTRequest body, diff --git a/core/src/test/java/org/apache/iceberg/rest/TestHTTPClient.java b/core/src/test/java/org/apache/iceberg/rest/TestHTTPClient.java index 21a9ba5b9f1f..6445fd237e06 100644 --- a/core/src/test/java/org/apache/iceberg/rest/TestHTTPClient.java +++ b/core/src/test/java/org/apache/iceberg/rest/TestHTTPClient.java @@ -18,6 +18,7 @@ */ package org.apache.iceberg.rest; +import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; @@ -33,6 +34,7 @@ import java.util.Locale; import java.util.Map; import java.util.Objects; +import java.util.function.Consumer; import org.apache.iceberg.AssertHelpers; import org.apache.iceberg.IcebergBuild; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; @@ -125,7 +127,8 @@ public static void testHttpMethodOnSuccess(HttpMethod method) throws JsonProcess String path = addRequestTestCaseAndGetPath(method, body, statusCode); - Item successResponse = doExecuteRequest(method, path, body, onError); + Item successResponse = + doExecuteRequest(method, path, body, onError, h -> assertThat(h).isNotEmpty()); if (method.usesRequestBody()) { Assert.assertEquals( @@ -157,7 +160,7 @@ public static void testHttpMethodOnFailure(HttpMethod method) throws JsonProcess RuntimeException.class, String.format( "Called error handler for method %s due to status code: %d", method, statusCode), - () -> doExecuteRequest(method, path, body, onError)); + () -> doExecuteRequest(method, path, body, onError, h -> {})); verify(onError).accept(any()); } @@ -208,11 +211,15 @@ private static String addRequestTestCaseAndGetPath(HttpMethod method, Item body, } private static Item doExecuteRequest( - HttpMethod method, String path, Item body, ErrorHandler onError) { + HttpMethod method, + String path, + Item body, + ErrorHandler onError, + Consumer> responseHeaders) { Map headers = ImmutableMap.of("Authorization", "Bearer " + BEARER_AUTH_TOKEN); switch (method) { case POST: - return restClient.post(path, body, Item.class, headers, onError); + return restClient.post(path, body, Item.class, headers, onError, responseHeaders); case GET: return restClient.get(path, Item.class, headers, onError); case HEAD: