@Override
public Response apply(Request request) {
ImmutableList<String> parts =
ImmutableList.copyOf(
Splitter.on("/").omitEmptyStrings().split(request.getUri().getPath()));
if (request.getMethod().equals("DELETE")) {
assertEquals(parts.size(), 1);
return new TestingResponse(HttpStatus.OK, ImmutableListMultimap.of(), new byte[0]);
}
assertEquals(parts.size(), 2);
String taskId = parts.get(0);
int pageToken = Integer.parseInt(parts.get(1));
Builder<String, String> headers = ImmutableListMultimap.builder();
headers.put(PRESTO_PAGE_TOKEN, String.valueOf(pageToken));
TaskBuffer taskBuffer = taskBuffers.getUnchecked(taskId);
Page page = taskBuffer.getPage(pageToken);
headers.put(CONTENT_TYPE, PRESTO_PAGES);
if (page != null) {
headers.put(PRESTO_PAGE_NEXT_TOKEN, String.valueOf(pageToken + 1));
headers.put(PRESTO_BUFFER_COMPLETE, String.valueOf(false));
DynamicSliceOutput output = new DynamicSliceOutput(256);
PagesSerde.writePages(blockEncodingSerde, output, page);
return new TestingResponse(HttpStatus.OK, headers.build(), output.slice().getInput());
} else if (taskBuffer.isFinished()) {
headers.put(PRESTO_PAGE_NEXT_TOKEN, String.valueOf(pageToken));
headers.put(PRESTO_BUFFER_COMPLETE, String.valueOf(true));
return new TestingResponse(HttpStatus.OK, headers.build(), new byte[0]);
} else {
headers.put(PRESTO_PAGE_NEXT_TOKEN, String.valueOf(pageToken));
headers.put(PRESTO_BUFFER_COMPLETE, String.valueOf(false));
return new TestingResponse(HttpStatus.NO_CONTENT, headers.build(), new byte[0]);
}
}
@Override
public void writeTo(
List<Page> pages,
Class<?> type,
Type genericType,
Annotation[] annotations,
MediaType mediaType,
MultivaluedMap<String, Object> httpHeaders,
OutputStream output)
throws IOException, WebApplicationException {
try {
SliceOutput sliceOutput = new OutputStreamSliceOutput(output);
PagesSerde.writePages(blockEncodingSerde, sliceOutput, pages);
// We use flush instead of close, because the underlying stream would be closed and that is
// not allowed.
sliceOutput.flush();
} catch (RuntimeIOException e) {
// EOF exception occurs when the client disconnects while writing data
// This is not a "server" problem so we don't want to log this
if (!(e.getCause() instanceof EOFException)) {
throw e;
}
}
}