private static boolean getComplete(Response response) {
String bufferComplete = response.getHeader(PRESTO_BUFFER_COMPLETE);
if (bufferComplete == null) {
throw new PageTransportErrorException(format("Expected %s header", PRESTO_BUFFER_COMPLETE));
}
return Boolean.parseBoolean(bufferComplete);
}
private static long getNextToken(Response response) {
String nextTokenHeader = response.getHeader(PRESTO_PAGE_NEXT_TOKEN);
if (nextTokenHeader == null) {
throw new PageTransportErrorException(format("Expected %s header", PRESTO_PAGE_NEXT_TOKEN));
}
return Long.parseLong(nextTokenHeader);
}
private static String getTaskInstanceId(Response response) {
String taskInstanceId = response.getHeader(PRESTO_TASK_INSTANCE_ID);
if (taskInstanceId == null) {
throw new PageTransportErrorException(
format("Expected %s header", PRESTO_TASK_INSTANCE_ID));
}
return taskInstanceId;
}
@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);
}
}