Examples with QueryStatusInfo io.prestosql.client.QueryStatusInfo used on opensource projects

Example 11 with QueryStatusInfo

use of io.prestosql.client.QueryStatusInfo in project hetu-core by openlookeng.

the class AbstractTestingPrestoClient method execute.

public ResultWithQueryId<T> execute(Session session, @Language("SQL") String sql) {
    ResultsSession<T> resultsSession = getResultSession(session);
    ClientSession clientSession = toClientSession(session, prestoServer.getBaseUrl(), new Duration(2, TimeUnit.MINUTES));
    try (StatementClient client = newStatementClient(httpClient, clientSession, sql)) {
        while (client.isRunning()) {
            resultsSession.addResults(client.currentStatusInfo(), client.currentData());
            client.advance();
        }
        checkState(client.isFinished());
        QueryError error = client.finalStatusInfo().getError();
        if (error == null) {
            QueryStatusInfo results = client.finalStatusInfo();
            if (results.getUpdateType() != null) {
                resultsSession.setUpdateType(results.getUpdateType());
            }
            if (results.getUpdateCount() != null) {
                resultsSession.setUpdateCount(results.getUpdateCount());
            }
            resultsSession.setWarnings(results.getWarnings());
            T result = resultsSession.build(client.getSetSessionProperties(), client.getResetSessionProperties());
            return new ResultWithQueryId<>(new QueryId(results.getId()), result);
        }
        if (error.getFailureInfo() != null) {
            RuntimeException remoteException = error.getFailureInfo().toException();
            throw new RuntimeException(Optional.ofNullable(remoteException.getMessage()).orElseGet(remoteException::toString), remoteException);
        }
        throw new RuntimeException("Query failed: " + error.getMessage());
    // dump query info to console for debugging (NOTE: not pretty printed)
    }
}