Reuse execute method in BenchmarkQueryRunner
package com.facebook.presto.benchmark.driver;
import com.facebook.presto.client.ClientSession;
+import com.facebook.presto.client.QueryData;
import com.facebook.presto.client.QueryError;
import com.facebook.presto.client.StatementClient;
import com.facebook.presto.client.StatementStats;
@@ -34,6 +35,8 @@ import java.net.URI;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
import static com.facebook.presto.benchmark.driver.BenchmarkQueryResult.failResult;
import static com.facebook.presto.benchmark.driver.BenchmarkQueryResult.passResult;
@@ -149,53 +152,55 @@ public class BenchmarkQueryRunner
{
failures = 0;
while (true) {
- // start query
- StatementClient client = new StatementClient(okHttpClient, session, "show schemas");
-
- // read query output
ImmutableList.Builder<String> schemas = ImmutableList.builder();
- while (client.isValid() && client.advance()) {
- // we do not process the output
- Iterable<List<Object>> data = client.currentData().getData();
- if (data != null) {
- for (List<Object> objects : data) {
+ AtomicBoolean success = new AtomicBoolean(true);
+ execute(
+ session,
+ "show schemas",
+ queryData -> {
+ if (queryData.getData() != null) {
+ for (List<Object> objects : queryData.getData()) {
schemas.add(objects.get(0).toString());
}
}
+ },
+ queryError -> {
+ success.set(false);
+ handleFailure(getCause(queryError));
+ });
+ if (success.get()) {
+ return schemas.build();
}
-
- // verify final state
- if (client.isClosed()) {
- throw new IllegalStateException("Query aborted by user");
- }
-
- if (client.isGone()) {
- throw new IllegalStateException("Query is gone (server restarted?)");
}
-
- QueryError resultsError = client.finalStatusInfo().getError();
- if (resultsError != null) {
- RuntimeException cause = null;
- if (resultsError.getFailureInfo() != null) {
- cause = resultsError.getFailureInfo().toException();
}
- handleFailure(cause);
- continue;
+ private StatementStats execute(ClientSession session, String name, String query)
+ {
+ return execute(
+ session,
+ query,
+ queryData -> {}, // we do not process the output
+ resultsError -> {
+ throw new BenchmarkDriverExecutionException(format("Query %s failed: %s", name, resultsError.getMessage()), getCause(resultsError));
+ });
}
- return schemas.build();
+ private static RuntimeException getCause(QueryError queryError)
+ {
+ if (queryError.getFailureInfo() != null) {
+ return queryError.getFailureInfo().toException();
}
+ return null;
}
- private StatementStats execute(ClientSession session, String name, String query)
+ private StatementStats execute(ClientSession session, String query, Consumer<QueryData> queryDataConsumer, Consumer<QueryError> queryErrorConsumer)
{
// start query
StatementClient client = new StatementClient(okHttpClient, session, query);
// read query output
while (client.isValid() && client.advance()) {
- // we do not process the output
+ queryDataConsumer.accept(client.currentData());
}
// verify final state
@@ -209,12 +214,7 @@ public class BenchmarkQueryRunner
QueryError resultsError = client.finalStatusInfo().getError();
if (resultsError != null) {
- RuntimeException cause = null;
- if (resultsError.getFailureInfo() != null) {
- cause = resultsError.getFailureInfo().toException();
- }
-
- throw new BenchmarkDriverExecutionException(format("Query %s failed: %s", name, resultsError.getMessage()), cause);
+ queryErrorConsumer.accept(resultsError);
}
return client.finalStatusInfo().getStats();
