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

Example 1 with StatementClient

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

the class TestDataCenterHTTPClientV1 method assertQuery.

private void assertQuery(String sql) {
    OkHttpClient httpClient = new OkHttpClient();
    try {
        DataCenterClientSession clientSession = DataCenterClientSession.builder(queryRunner.getCoordinator().getBaseUrl(), "user").withClientTimeout(new Duration(2, MINUTES)).withTypeManager(typeManager).build();
        StatementClient client1 = DataCenterStatementClient.newStatementClient(httpClient, clientSession, sql, UUID.randomUUID().toString());
        long client1Count = 0;
        // wait for query to be fully scheduled
        while (client1.isRunning()) {
            if (client1.currentData().getData() != null) {
                for (List<Object> row : client1.currentData().getData()) {
                    System.out.println(row);
                    client1Count++;
                }
            }
            client1.advance();
        }
        System.out.println("ROWS: " + client1Count);
        StatementClient client2 = newStatementClient(httpClient, clientSession, sql);
        long client2Count = 0;
        // wait for query to be fully scheduled
        while (client2.isRunning()) {
            if (client2.currentData().getData() != null) {
                for (List<Object> row : client2.currentData().getData()) {
                    System.out.println(row);
                    client2Count++;
                }
            }
            client2.advance();
        }
        System.out.println("ROWS: " + client2Count);
        assertEquals(client1Count, client2Count);
    } finally {
        // close the client since, query is not managed by the client protocol
        httpClient.dispatcher().executorService().shutdown();
        httpClient.connectionPool().evictAll();
    }
}

Example 2 with StatementClient

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

the class TableNameCompleter method queryMetadata.

private List<String> queryMetadata(String query) {
    ImmutableList.Builder<String> cache = ImmutableList.builder();
    try (StatementClient client = queryRunner.startInternalQuery(query)) {
        while (client.isRunning() && !Thread.currentThread().isInterrupted()) {
            QueryData results = client.currentData();
            if (results.getData() != null) {
                for (List<Object> row : results.getData()) {
                    cache.add((String) row.get(0));
                }
            }
            client.advance();
        }
    }
    return cache.build();
}

Example 3 with StatementClient

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

the class DataCenterClient method getColumns.

/**
 * Get the columns information of this sql.
 *
 * @param sql statement.
 * @return the list of data center column that this sql contains.
 */
public List<DataCenterColumn> getColumns(String sql) {
    String query = "PREPARE subStatement FROM " + sql;
    String describe = "DESCRIBE OUTPUT subStatement";
    try (StatementClient preparedClient = execute(query)) {
        DataCenterClientSession newClientSession = DataCenterStatementClientFactory.createClientSession(preparedClient, this.config, this.typeManager);
        ImmutableList.Builder<DataCenterColumn> builder = new ImmutableList.Builder<>();
        Iterable<List<Object>> data = getResults(newClientSession, describe);
        for (List<Object> row : data) {
            String name = row.get(0).toString();
            String type = row.get(TYPE_POSITION).toString();
            if ("unknown".equalsIgnoreCase(type)) {
                // Cannot support unknown types
                return Collections.emptyList();
            }
            try {
                builder.add(new DataCenterColumn(name, parseType(typeManager, type)));
            } catch (IllegalArgumentException ex) {
                return Collections.emptyList();
            }
        }
        return builder.build();
    } catch (SQLException ex) {
        return Collections.emptyList();
    }
}

Example 4 with StatementClient

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

the class PrestoStatement method internalExecute.

final boolean internalExecute(String sql) throws SQLException {
    clearCurrentResults();
    checkOpen();
    StatementClient client = null;
    PrestoResultSet resultSet = null;
    try {
        client = connection().startQuery(sql, getStatementSessionProperties());
        if (client.isFinished()) {
            QueryStatusInfo finalStatusInfo = client.finalStatusInfo();
            if (finalStatusInfo.getError() != null) {
                throw resultsException(finalStatusInfo);
            }
        }
        executingClient.set(client);
        WarningsManager warningsManager = new WarningsManager();
        currentWarningsManager.set(Optional.of(warningsManager));
        resultSet = new PrestoResultSet(client, maxRows.get(), progressConsumer, warningsManager);
        // check if this is a query
        if (client.currentStatusInfo().getUpdateType() == null) {
            currentResult.set(resultSet);
            return true;
        }
        // this is an update, not a query
        while (resultSet.next()) {
        // ignore rows
        }
        connection().updateSession(client);
        Long updateCount = client.finalStatusInfo().getUpdateCount();
        currentUpdateCount.set((updateCount != null) ? updateCount : 0);
        currentUpdateType.set(client.finalStatusInfo().getUpdateType());
        warningsManager.addWarnings(client.finalStatusInfo().getWarnings());
        return false;
    } catch (ClientException e) {
        throw new SQLException(e.getMessage(), e);
    } catch (RuntimeException e) {
        throw new SQLException("Error executing query", e);
    } finally {
        executingClient.set(null);
        if (currentResult.get() == null) {
            if (resultSet != null) {
                resultSet.close();
            }
            if (client != null) {
                client.close();
            }
        }
    }
}

Example 5 with StatementClient

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

the class PrestoStatement method cancel.

@Override
public void cancel() throws SQLException {
    checkOpen();
    StatementClient client = executingClient.get();
    if (client != null) {
        client.close();
    }
    closeResultSet();
}