Examples with StatementClient com.facebook.presto.client.StatementClient used on opensource projects

Example 6 with StatementClient

use of com.facebook.presto.client.StatementClient in project presto by prestodb.

the class BenchmarkQueryRunner method getSchemas.

public List<String> getSchemas(ClientSession session) {
    failures = 0;
    while (true) {
        // start query
        StatementClient client = new StatementClient(httpClient, queryResultsCodec, 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.current().getData();
            if (data != null) {
                for (List<Object> objects : data) {
                    schemas.add(objects.get(0).toString());
                }
            }
        }
        // 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.finalResults().getError();
        if (resultsError != null) {
            RuntimeException cause = null;
            if (resultsError.getFailureInfo() != null) {
                cause = resultsError.getFailureInfo().toException();
            }
            handleFailure(cause);
            continue;
        }
        return schemas.build();
    }
}

Example 7 with StatementClient

use of com.facebook.presto.client.StatementClient in project presto by prestodb.

the class PrestoStatement method execute.

@Override
public boolean execute(String sql) throws SQLException {
    clearCurrentResults();
    checkOpen();
    StatementClient client = null;
    ResultSet resultSet = null;
    try {
        client = connection().startQuery(sql);
        if (client.isFailed()) {
            throw resultsException(client.finalResults());
        }
        resultSet = new PrestoResultSet(client, maxRows.get(), progressConsumer);
        checkSetOrResetSession(client);
        // check if this is a query
        if (client.current().getUpdateType() == null) {
            currentResult.set(resultSet);
            return true;
        }
        // this is an update, not a query
        while (resultSet.next()) {
        // ignore rows
        }
        Long updateCount = client.finalResults().getUpdateCount();
        currentUpdateCount.set((updateCount != null) ? updateCount : 0);
        return false;
    } catch (RuntimeException e) {
        throw new SQLException("Error executing query", e);
    } finally {
        if (currentResult.get() == null) {
            if (resultSet != null) {
                resultSet.close();
            }
            if (client != null) {
                client.close();
            }
        }
    }
}

Example 8 with StatementClient

use of com.facebook.presto.client.StatementClient in project presto by prestodb.

the class AbstractTestingPrestoClient method execute.

public T execute(Session session, @Language("SQL") String sql) {
    ResultsSession<T> resultsSession = getResultSession(session);
    ClientSession clientSession = toClientSession(session, prestoServer.getBaseUrl(), true, new Duration(2, TimeUnit.MINUTES));
    try (StatementClient client = new StatementClient(httpClient, QUERY_RESULTS_CODEC, clientSession, sql)) {
        while (client.isValid()) {
            QueryResults results = client.current();
            resultsSession.addResults(results);
            client.advance();
        }
        if (!client.isFailed()) {
            QueryResults results = client.finalResults();
            if (results.getUpdateType() != null) {
                resultsSession.setUpdateType(results.getUpdateType());
            }
            if (results.getUpdateCount() != null) {
                resultsSession.setUpdateCount(results.getUpdateCount());
            }
            return resultsSession.build(client.getSetSessionProperties(), client.getResetSessionProperties());
        }
        QueryError error = client.finalResults().getError();
        verify(error != null, "no error");
        if (error.getFailureInfo() != null) {
            throw error.getFailureInfo().toException();
        }
        throw new RuntimeException("Query failed: " + error.getMessage());
    // dump query info to console for debugging (NOTE: not pretty printed)
    // JsonCodec<QueryInfo> queryInfoJsonCodec = createCodecFactory().prettyPrint().jsonCodec(QueryInfo.class);
    // log.info("\n" + queryInfoJsonCodec.toJson(queryInfo));
    }
}

Example 9 with StatementClient

use of com.facebook.presto.client.StatementClient in project airpal by airbnb.

the class QueryClient method executeWith.

public <T> T executeWith(Function<StatementClient, T> function) throws QueryTimeOutException {
    final Stopwatch stopwatch = Stopwatch.createStarted();
    T t = null;
    try (StatementClient client = queryRunner.startInternalQuery(query)) {
        while (client.isValid() && !Thread.currentThread().isInterrupted()) {
            if (stopwatch.elapsed(TimeUnit.MILLISECONDS) > timeout.toMilliseconds()) {
                throw new QueryTimeOutException(stopwatch.elapsed(TimeUnit.MILLISECONDS));
            }
            t = function.apply(client);
            client.advance();
        }
        finalResults.set(client.finalResults());
    } catch (RuntimeException | QueryTimeOutException e) {
        stopwatch.stop();
        throw e;
    }
    return t;
}

Example 10 with StatementClient

use of com.facebook.presto.client.StatementClient in project airpal by airbnb.

the class PreviewTableCache method queryRows.

private List<List<Object>> queryRows(String query) {
    final ImmutableList.Builder<List<Object>> cache = ImmutableList.builder();
    QueryRunner queryRunner = queryRunnerFactory.create();
    QueryClient queryClient = new QueryClient(queryRunner, io.dropwizard.util.Duration.seconds(60), query);
    try {
        queryClient.executeWith(new Function<StatementClient, Void>() {

            @Nullable
            @Override
            public Void apply(StatementClient client) {
                QueryResults results = client.current();
                if (results.getData() != null) {
                    cache.addAll(results.getData());
                }
                return null;
            }
        });
    } catch (QueryClient.QueryTimeOutException e) {
        log.error("Caught timeout loading columns", e);
    }
    return cache.build();
}