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

Example 11 with StatementClient

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

the class BenchmarkQueryRunner method execute.

private StatementStats execute(ClientSession session, String query, Consumer<QueryData> queryDataConsumer, Consumer<QueryError> queryErrorConsumer) {
    // start query
    try (StatementClient client = newStatementClient(okHttpClient, session, query)) {
        // read query output
        while (client.isRunning()) {
            queryDataConsumer.accept(client.currentData());
            if (!client.advance()) {
                break;
            }
        }
        // verify final state
        if (client.isClientAborted()) {
            throw new IllegalStateException("Query aborted by user");
        }
        if (client.isClientError()) {
            throw new IllegalStateException("Query is gone (server restarted?)");
        }
        verify(client.isFinished());
        QueryError resultsError = client.finalStatusInfo().getError();
        if (resultsError != null) {
            queryErrorConsumer.accept(resultsError);
        }
        return client.finalStatusInfo().getStats();
    }
}

Example 12 with StatementClient

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

the class ColumnService method queryColumns.

private List<Column> queryColumns(String fqnTableName, String user) {
    String statement = format("SHOW COLUMNS FROM %s", fqnTableName);
    QueryRunner queryRunner = queryRunnerFactory.create(QueryEditorUIModule.UI_QUERY_SOURCE, user);
    QueryClient queryClient = new QueryClient(queryRunner, Duration.standardSeconds(60), statement);
    final ImmutableList.Builder<Column> cache = ImmutableList.builder();
    try {
        queryClient.executeWith(new Function<StatementClient, Void>() {

            @Nullable
            @Override
            public Void apply(StatementClient client) {
                QueryData results = client.currentData();
                if (results.getData() != null) {
                    for (List<Object> row : results.getData()) {
                        TypeSignature typeSignature = TypeSignature.parseTypeSignature((String) row.get(1));
                        Column column = new Column((String) row.get(0), (String) row.get(1), Query.toClientTypeSignature(typeSignature));
                        cache.add(column);
                    }
                }
                return null;
            }
        });
    } catch (QueryClient.QueryTimeOutException e) {
        log.error("Caught timeout loading columns", e);
    }
    return cache.build();
}

Example 13 with StatementClient

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

the class PreviewTableService method queryRows.

private List<List<Object>> queryRows(String fqnTableName, String user) {
    String statement = format("SELECT * FROM %s LIMIT %d", fqnTableName, PREVIEW_LIMIT);
    QueryRunner queryRunner = queryRunnerFactory.create(QueryEditorUIModule.UI_QUERY_SOURCE, user);
    QueryClient queryClient = new QueryClient(queryRunner, Duration.standardSeconds(60), statement);
    final ImmutableList.Builder<List<Object>> cache = ImmutableList.builder();
    try {
        queryClient.executeWith(new Function<StatementClient, Void>() {

            @Nullable
            @Override
            public Void apply(StatementClient client) {
                QueryData results = client.currentData();
                if (results.getData() != null) {
                    List<List<Object>> resultsData = ImmutableList.copyOf(results.getData());
                    for (List<Object> row : resultsData) {
                        final String[] values = new String[row.size()];
                        for (int i = 0; i < values.length; i++) {
                            // Display byte array as Hexadecimal in order to keep consistent with OpenLooKeng client
                            if (row.get(i) instanceof byte[]) {
                                byte[] bytes = (byte[]) row.get(i);
                                StringBuilder sb = new StringBuilder();
                                for (byte b : bytes) {
                                    String hex = Integer.toHexString(b & 0xFF);
                                    if (hex.length() < 2) {
                                        sb.append(0);
                                    }
                                    sb.append(hex);
                                }
                                values[i] = sb.toString();
                            } else {
                                final Object value = row.get(i);
                                values[i] = (value == null) ? "" : value.toString();
                            }
                        }
                        cache.add(Arrays.asList(values));
                    }
                }
                return null;
            }
        });
    } catch (QueryClient.QueryTimeOutException e) {
        log.error("Caught timeout loading columns", e);
    }
    return cache.build();
}

Example 14 with StatementClient

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

the class SchemaService method queryStatement.

private Set<String> queryStatement(QueryRunner queryRunner, String statement) {
    QueryClient queryClient = new QueryClient(queryRunner, Duration.standardSeconds(120), statement);
    final Set<String> resultSet = new HashSet();
    try {
        queryClient.executeWith(new Function<StatementClient, Void>() {

            @Nullable
            @Override
            public Void apply(StatementClient client) {
                QueryData results = client.currentData();
                if (results.getData() != null) {
                    for (List<Object> row : results.getData()) {
                        resultSet.add((String) row.get(0));
                    }
                }
                return null;
            }
        });
    } catch (QueryClient.QueryTimeOutException e) {
        LOG.error("Caught timeout loading data", e);
    }
    return resultSet;
}

Example 15 with StatementClient

use of io.prestosql.client.StatementClient 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)
    }
}