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

Example 16 with ClientSession

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

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)
    // JsonCodec<QueryInfo> queryInfoJsonCodec = createCodecFactory().prettyPrint().jsonCodec(QueryInfo.class);
    // log.info("\n" + queryInfoJsonCodec.toJson(queryInfo));
    }
}

Example 17 with ClientSession

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

the class AbstractTestingPrestoClient method toClientSession.

private static ClientSession toClientSession(Session session, URI server, Duration clientRequestTimeout) {
    ImmutableMap.Builder<String, String> properties = ImmutableMap.builder();
    properties.putAll(session.getSystemProperties());
    for (Entry<String, Map<String, String>> connectorProperties : session.getUnprocessedCatalogProperties().entrySet()) {
        for (Entry<String, String> entry : connectorProperties.getValue().entrySet()) {
            properties.put(connectorProperties.getKey() + "." + entry.getKey(), entry.getValue());
        }
    }
    ImmutableMap.Builder<String, String> resourceEstimates = ImmutableMap.builder();
    ResourceEstimates estimates = session.getResourceEstimates();
    estimates.getExecutionTime().ifPresent(e -> resourceEstimates.put(EXECUTION_TIME, e.toString()));
    estimates.getCpuTime().ifPresent(e -> resourceEstimates.put(CPU_TIME, e.toString()));
    estimates.getPeakMemory().ifPresent(e -> resourceEstimates.put(PEAK_MEMORY, e.toString()));
    Map<String, String> serializedSessionFunctions = session.getSessionFunctions().entrySet().stream().collect(collectingAndThen(toMap(e -> SQL_FUNCTION_ID_JSON_CODEC.toJson(e.getKey()), e -> SQL_INVOKED_FUNCTION_JSON_CODEC.toJson(e.getValue())), ImmutableMap::copyOf));
    return new ClientSession(server, session.getIdentity().getUser(), session.getSource().orElse(null), session.getTraceToken(), session.getClientTags(), session.getClientInfo().orElse(null), session.getCatalog().orElse(null), session.getSchema().orElse(null), session.getTimeZoneKey().getId(), session.getLocale(), resourceEstimates.build(), properties.build(), session.getPreparedStatements(), session.getIdentity().getRoles(), session.getIdentity().getExtraCredentials(), session.getTransactionId().map(Object::toString).orElse(null), clientRequestTimeout, true, serializedSessionFunctions, ImmutableMap.of());
}

Example 18 with ClientSession

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

the class TestFinalQueryInfo method startQuery.

private static QueryId startQuery(String sql, DistributedQueryRunner queryRunner) {
    OkHttpClient httpClient = new OkHttpClient();
    try {
        ClientSession clientSession = new ClientSession(queryRunner.getCoordinator().getBaseUrl(), "user", "source", Optional.empty(), ImmutableSet.of(), null, null, null, "America/Los_Angeles", Locale.ENGLISH, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(), null, new Duration(2, MINUTES), true, ImmutableMap.of(), ImmutableMap.of());
        // start query
        StatementClient client = newStatementClient(httpClient, clientSession, sql);
        // wait for query to be fully scheduled
        while (client.isRunning() && !client.currentStatusInfo().getStats().isScheduled()) {
            client.advance();
        }
        return new QueryId(client.currentStatusInfo().getId());
    } finally {
        // close the client since, query is not managed by the client protocol
        httpClient.dispatcher().executorService().shutdown();
        httpClient.connectionPool().evictAll();
    }
}