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

Example 1 with DataCenterClientSession

use of io.prestosql.client.DataCenterClientSession 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 DataCenterClientSession

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

the class DataCenterPageSourceProvider method createPageSource.

@Override
public ConnectorPageSource createPageSource(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<ColumnHandle> columns, Optional<DynamicFilterSupplier> dynamicFilterSupplier) {
    // Build the sql
    String query = buildSql((DataCenterTableHandle) table, columns, ((DataCenterTableHandle) table).getLimit());
    Map<String, String> properties = new HashMap<>();
    // Only set the session if there is any dynamic filter for this page source
    if (dynamicFilterSupplier != null && dynamicFilterSupplier.isPresent()) {
        properties.put(ENABLE_CROSS_REGION_DYNAMIC_FILTER, "true");
    }
    if (config.isCompressionEnabled()) {
        properties.put(EXCHANGE_COMPRESSION, "true");
    }
    // Create a new client session
    DataCenterClientSession clientSession = DataCenterStatementClientFactory.createClientSession(this.config, this.typeManager, properties);
    return new DataCenterPageSource(this.httpClient, clientSession, query, ((DataCenterSplit) split).getQueryId(), columns, dynamicFilterSupplier);
}

Example 3 with DataCenterClientSession

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