@Override
public URI handle(Request request, Response response) {
if (response.getStatusCode() != HttpStatus.CREATED.code()) {
throw new UnexpectedResponseException(
String.format(
"Expected response code to be 201 CREATED, but was %s %s",
response.getStatusCode(), response.getStatusMessage()),
request,
response);
}
String location = response.getHeader("Location");
if (location == null) {
throw new UnexpectedResponseException(
"Response does not contain a Location header", request, response);
}
return URI.create(location);
}
@Override
public PagesResponse handle(Request request, Response response) {
// no content means no content was created within the wait period, but query is still ok
// if job is finished, complete is set in the response
if (response.getStatusCode() == HttpStatus.NO_CONTENT.code()) {
return createEmptyPagesResponse(
getTaskInstanceId(response),
getToken(response),
getNextToken(response),
getComplete(response));
}
// otherwise we must have gotten an OK response, everything else is considered fatal
if (response.getStatusCode() != HttpStatus.OK.code()) {
StringBuilder body = new StringBuilder();
try (BufferedReader reader =
new BufferedReader(new InputStreamReader(response.getInputStream()))) {
// Get up to 1000 lines for debugging
for (int i = 0; i < 1000; i++) {
String line = reader.readLine();
// Don't output more than 100KB
if (line == null || body.length() + line.length() > 100 * 1024) {
break;
}
body.append(line + "\n");
}
} catch (RuntimeException | IOException e) {
// Ignored. Just return whatever message we were able to decode
}
throw new PageTransportErrorException(
format(
"Expected response code to be 200, but was %s %s: %s%n%s",
response.getStatusCode(),
response.getStatusMessage(),
request.getUri(),
body.toString()));
}
// invalid content type can happen when an error page is returned, but is unlikely given the
// above 200
String contentType = response.getHeader(CONTENT_TYPE);
if (contentType == null) {
throw new PageTransportErrorException(
format("%s header is not set: %s %s", CONTENT_TYPE, request.getUri(), response));
}
if (!mediaTypeMatches(contentType, PRESTO_PAGES_TYPE)) {
throw new PageTransportErrorException(
format(
"Expected %s response from server but got %s: %s",
PRESTO_PAGES_TYPE, contentType, request.getUri()));
}
String taskInstanceId = getTaskInstanceId(response);
long token = getToken(response);
long nextToken = getNextToken(response);
boolean complete = getComplete(response);
try (SliceInput input = new InputStreamSliceInput(response.getInputStream())) {
List<Page> pages = ImmutableList.copyOf(readPages(blockEncodingSerde, input));
return createPagesResponse(taskInstanceId, token, nextToken, pages, complete);
} catch (IOException e) {
throw Throwables.propagate(e);
}
}