예제 #1
0
    @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;
     }
   }
 }