private Response proxyJsonResponse(String nodeId, String workerPath) { Set<Node> nodes = nodeManager.getNodes(NodeState.ACTIVE); Node node = nodes.stream() .filter(n -> n.getNodeIdentifier().equals(nodeId)) .findFirst() .orElseThrow(() -> new WebApplicationException(NOT_FOUND)); Request request = prepareGet() .setUri(uriBuilderFrom(node.getHttpUri()) .appendPath(workerPath) .build()) .build(); InputStream responseStream = httpClient.execute(request, new StreamingJsonResponseHandler()); return Response.ok(responseStream, APPLICATION_JSON_TYPE).build(); }
private long getTotalCpuTime() { long totalCpuTime = 0; for (URI server : nodes) { URI addressUri = uriBuilderFrom(server).replacePath("/v1/jmx/mbean/java.lang:type=OperatingSystem/ProcessCpuTime").build(); String data = httpClient.execute(prepareGet().setUri(addressUri).build(), createStringResponseHandler()).getBody(); totalCpuTime += parseLong(data.trim()); } return TimeUnit.NANOSECONDS.toNanos(totalCpuTime); }
private List<URI> getAllNodes(URI server) { Request request = prepareGet().setUri(uriBuilderFrom(server).replacePath("/v1/service/presto").build()).build(); JsonResponseHandler<ServiceDescriptorsRepresentation> responseHandler = createJsonResponseHandler(jsonCodec(ServiceDescriptorsRepresentation.class)); ServiceDescriptorsRepresentation serviceDescriptors = httpClient.execute(request, responseHandler); ImmutableList.Builder<URI> addresses = ImmutableList.builder(); for (ServiceDescriptor serviceDescriptor : serviceDescriptors.getServiceDescriptors()) { String httpUri = serviceDescriptor.getProperties().get("http"); if (httpUri != null) { addresses.add(URI.create(httpUri)); } } return addresses.build(); } }
public Response cancel(String queryId) { requireNonNull(queryId, "queryId is null"); URI cancelUri = uriBuilderFrom(uri).appendPath("/v1/query").appendPath(queryId).build(); Request request = prepareDelete().setUri(cancelUri).build(); return httpClient.execute(request, new ResponseHandler<Response, RuntimeException>() { @Override public Response handleException(Request request, Exception exception) { throw propagate(request, exception); } @Override public Response handle(Request request, Response response) { return response; } }); }
@Test public void testGetFailedNodes() { List<Stats> nodes = client.execute( prepareGet().setUri(server.resolve("/v1/node/failed")).build(), createJsonResponseHandler(listJsonCodec(Stats.class))); assertTrue(nodes.isEmpty()); } }
@Test public void testServerStarts() { StatusResponseHandler.StatusResponse response = client.execute( prepareGet().setUri(server.resolve("/v1/query")).build(), createStatusResponseHandler()); assertEquals(response.getStatusCode(), OK.getStatusCode()); }
@Test public void testGetAllNodes() { List<Stats> nodes = client.execute( prepareGet().setUri(server.resolve("/v1/node")).build(), createJsonResponseHandler(listJsonCodec(Stats.class))); // we only have one node and the list never contains the current node assertTrue(nodes.isEmpty()); }
@Test public void testGetQueryStateInfosForUserNoResult() { List<QueryStateInfo> infos = client.execute( prepareGet().setUri(server.resolve("/v1/queryState?user=user3")).build(), createJsonResponseHandler(listJsonCodec(QueryStateInfo.class))); assertTrue(infos.isEmpty()); }
@Override public Optional<QueryStats> getQueryStats(String queryId) { URI uri = uriBuilderFrom(baseUri).appendPath("/v1/query").appendPath(queryId).build(); Request request = prepareGet().setUri(uri).build(); return httpClient.execute(request, new GetQueryStatsResponseHandler()); }
@Test public void testGetQueryStateInfosForUser() { List<QueryStateInfo> infos = client.execute( prepareGet().setUri(server.resolve("/v1/queryState?user=user2")).build(), createJsonResponseHandler(listJsonCodec(QueryStateInfo.class))); assertEquals(infos.size(), 1); }
@Test public void testGetQueryStateInfo() { QueryStateInfo info = client.execute( prepareGet().setUri(server.resolve("/v1/queryState/" + queryResults.getId())).build(), createJsonResponseHandler(jsonCodec(QueryStateInfo.class))); assertNotNull(info); }
@Test public void testGetAllQueryStateInfos() { List<QueryStateInfo> infos = client.execute( prepareGet().setUri(server.resolve("/v1/queryState")).build(), createJsonResponseHandler(listJsonCodec(QueryStateInfo.class))); assertEquals(infos.size(), 2); }
@Test(expectedExceptions = {UnexpectedResponseException.class}, expectedExceptionsMessageRegExp = ".*404: Not Found") public void testGetQueryStateInfoNo() { client.execute( prepareGet().setUri(server.resolve("/v1/queryState/123")).build(), createJsonResponseHandler(jsonCodec(QueryStateInfo.class))); } }
private List<BasicQueryInfo> getQueryInfos(String path) { Request request = prepareGet().setUri(server.resolve(path)).build(); return client.execute(request, createJsonResponseHandler(listJsonCodec(BasicQueryInfo.class))); }
@Test public void testNoTransactionSupport() { Request request = preparePost() .setUri(uriFor("/v1/statement")) .setBodyGenerator(createStaticBodyGenerator("start transaction", UTF_8)) .setHeader(PRESTO_USER, "user") .setHeader(PRESTO_SOURCE, "source") .build(); QueryResults queryResults = client.execute(request, createJsonResponseHandler(QUERY_RESULTS_CODEC)); while (queryResults.getNextUri() != null) { queryResults = client.execute(prepareGet().setUri(queryResults.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_CODEC)); } assertNotNull(queryResults.getError()); assertEquals(queryResults.getError().getErrorCode(), INCOMPATIBLE_CLIENT.toErrorCode().getCode()); }
@Test public void testTransactionSupport() { Request request = preparePost() .setUri(uriFor("/v1/statement")) .setBodyGenerator(createStaticBodyGenerator("start transaction", UTF_8)) .setHeader(PRESTO_USER, "user") .setHeader(PRESTO_SOURCE, "source") .setHeader(PRESTO_TRANSACTION_ID, "none") .build(); JsonResponse<QueryResults> queryResults = client.execute(request, createFullJsonResponseHandler(QUERY_RESULTS_CODEC)); ImmutableList.Builder<List<Object>> data = ImmutableList.builder(); while (true) { if (queryResults.getValue().getData() != null) { data.addAll(queryResults.getValue().getData()); } if (queryResults.getValue().getNextUri() == null) { break; } queryResults = client.execute(prepareGet().setUri(queryResults.getValue().getNextUri()).build(), createFullJsonResponseHandler(QUERY_RESULTS_CODEC)); } assertNull(queryResults.getValue().getError()); assertNotNull(queryResults.getHeader(PRESTO_STARTED_TRANSACTION_ID)); }
private void runToCompletion(String sql) { URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/statement")).build(); Request request = preparePost() .setHeader(PRESTO_USER, "user") .setUri(uri) .setBodyGenerator(createStaticBodyGenerator(sql, UTF_8)) .build(); QueryResults queryResults = client.execute(request, createJsonResponseHandler(jsonCodec(QueryResults.class))); while (queryResults.getNextUri() != null) { request = prepareGet() .setHeader(PRESTO_USER, "user") .setUri(queryResults.getNextUri()) .build(); queryResults = client.execute(request, createJsonResponseHandler(jsonCodec(QueryResults.class))); } }
@BeforeClass public void setup() { Request request1 = preparePost() .setUri(uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/statement").build()) .setBodyGenerator(createStaticBodyGenerator(LONG_LASTING_QUERY, UTF_8)) .setHeader(PRESTO_USER, "user1") .build(); queryResults = client.execute(request1, createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC)); client.execute(prepareGet().setUri(queryResults.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC)); Request request2 = preparePost() .setUri(uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/statement").build()) .setBodyGenerator(createStaticBodyGenerator(LONG_LASTING_QUERY, UTF_8)) .setHeader(PRESTO_USER, "user2") .build(); QueryResults queryResults2 = client.execute(request2, createJsonResponseHandler(jsonCodec(QueryResults.class))); client.execute(prepareGet().setUri(queryResults2.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC)); // queries are started in the background, so they may not all be immediately visible while (true) { List<BasicQueryInfo> queryInfos = client.execute( prepareGet().setUri(uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/query").build()).build(), createJsonResponseHandler(listJsonCodec(BasicQueryInfo.class))); if ((queryInfos.size() == 2) && queryInfos.stream().allMatch(info -> info.getState() == RUNNING)) { break; } } }
@Test public void testInvalidSessionError() { String invalidTimeZone = "this_is_an_invalid_time_zone"; Request request = preparePost().setHeader(PRESTO_USER, "user") .setUri(uriFor("/v1/statement")) .setBodyGenerator(createStaticBodyGenerator("show catalogs", UTF_8)) .setHeader(PRESTO_SOURCE, "source") .setHeader(PRESTO_CATALOG, "catalog") .setHeader(PRESTO_SCHEMA, "schema") .setHeader(PRESTO_PATH, "path") .setHeader(PRESTO_TIME_ZONE, invalidTimeZone) .build(); QueryResults queryResults = client.execute(request, createJsonResponseHandler(QUERY_RESULTS_CODEC)); while (queryResults.getNextUri() != null) { queryResults = client.execute(prepareGet().setUri(queryResults.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_CODEC)); } QueryError queryError = queryResults.getError(); assertNotNull(queryError); TimeZoneNotSupportedException expected = new TimeZoneNotSupportedException(invalidTimeZone); assertEquals(queryError.getErrorCode(), expected.getErrorCode().getCode()); assertEquals(queryError.getErrorName(), expected.getErrorCode().getName()); assertEquals(queryError.getErrorType(), expected.getErrorCode().getType().name()); assertEquals(queryError.getMessage(), expected.getMessage()); }