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

Example 6 with StatementClient

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

the class ExecutionClient method scheduleExecution.

private UUID scheduleExecution(Duration timeout, QueryRunner queryRunner, QueryExecutionAuthorizer authorizer, BlockingQueue<Job> jobs, URI requestUri) {
    final Job job;
    try {
        job = jobs.take();
    } catch (InterruptedException e) {
        return null;
    }
    final Execution execution = new Execution(job, queryRunner, queryInfoClient, authorizer, timeout, outputBuilderFactory, persistorFactory, requestUri);
    executionMap.put(job.getUuid(), execution);
    activeJobsStore.jobStarted(job);
    ListenableFuture<Job> result = executor.submit(execution);
    Futures.addCallback(result, new FutureCallback<Job>() {

        @Override
        public void onSuccess(@Nullable Job result) {
            if (result != null) {
                result.setState(JobState.FINISHED);
            }
            // Add Active Job
            if (jobs.peek() != null) {
                QueryRunner nextQueryRunner = getNextQueryRunner();
                scheduleExecution(timeout, nextQueryRunner, authorizer, jobs, requestUri);
            }
            jobFinished(result);
        }

        // Re-Use session level fields among multi statement queries.
        private QueryRunner getNextQueryRunner() {
            StatementClient client = queryRunner.getCurrentClient();
            ClientSession session = queryRunner.getSession();
            ClientSession.Builder builder = ClientSession.builder(session).withoutTransactionId();
            if (client.getSetCatalog().isPresent()) {
                builder.withCatalog(client.getSetCatalog().get());
            }
            if (client.getSetSchema().isPresent()) {
                builder.withSchema(client.getSetSchema().get());
            }
            if (client.getStartedTransactionId() != null) {
                builder = builder.withTransactionId(client.getStartedTransactionId());
            }
            if (client.getSetPath().isPresent()) {
                builder = builder.withPath(client.getSetPath().get());
            }
            if (!client.getSetSessionProperties().isEmpty() || !client.getResetSessionProperties().isEmpty()) {
                Map<String, String> sessionProperties = new HashMap<>(session.getProperties());
                sessionProperties.putAll(client.getSetSessionProperties());
                sessionProperties.keySet().removeAll(client.getResetSessionProperties());
                builder = builder.withProperties(sessionProperties);
            }
            if (!client.getSetRoles().isEmpty()) {
                Map<String, ClientSelectedRole> roles = new HashMap<>(session.getRoles());
                roles.putAll(client.getSetRoles());
                builder = builder.withRoles(roles);
            }
            if (!client.getAddedPreparedStatements().isEmpty() || !client.getDeallocatedPreparedStatements().isEmpty()) {
                Map<String, String> preparedStatements = new HashMap<>(session.getPreparedStatements());
                preparedStatements.putAll(client.getAddedPreparedStatements());
                preparedStatements.keySet().removeAll(client.getDeallocatedPreparedStatements());
                builder = builder.withPreparedStatements(preparedStatements);
            }
            return queryRunnerFactory.create(builder.build());
        }

        @Override
        public void onFailure(@NotNull Throwable t) {
            job.setState(JobState.FAILED);
            if (job.getError() == null) {
                job.setError(new QueryError(t.getMessage(), null, -1, null, null, null, null, null));
            }
            jobFinished(job);
        }
    }, MoreExecutors.directExecutor());
    return job.getUuid();
}

Example 7 with StatementClient

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

the class QueryClient method executeWith.

public <T> T executeWith(Function<StatementClient, T> function) throws QueryTimeOutException {
    final Stopwatch stopwatch = Stopwatch.createStarted();
    T t = null;
    try (StatementClient client = queryRunner.startInternalQuery(query)) {
        while (client.isRunning() && !Thread.currentThread().isInterrupted()) {
            if (stopwatch.elapsed(TimeUnit.MILLISECONDS) > timeout.toMillis()) {
                throw new QueryTimeOutException(stopwatch.elapsed(TimeUnit.MILLISECONDS));
            }
            t = function.apply(client);
            client.advance();
        }
        finalResults.set(client.finalStatusInfo());
    } catch (RuntimeException | QueryTimeOutException e) {
        stopwatch.stop();
        throw e;
    }
    return t;
}

Example 8 with StatementClient

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

the class DataCenterClient method getResults.

private Iterable<List<Object>> getResults(DataCenterClientSession session, String query) throws SQLException {
    try (StatementClient client = DataCenterStatementClient.newStatementClient(this.httpClient, session, query, this.globalQueryIdGenerator.createId())) {
        List<Iterable<List<Object>>> list = new LinkedList<>();
        while (client.isRunning()) {
            if (client.currentData() != null) {
                Iterable<List<Object>> data = client.currentData().getData();
                if (data != null) {
                    list.add(data);
                }
            }
            client.advance();
        }
        verify(client.isFinished());
        QueryStatusInfo results = client.finalStatusInfo();
        if (results.getError() != null) {
            throw resultsException(results);
        }
        return Iterables.concat(list);
    } catch (RuntimeException ex) {
        throw new SQLException(ex.getMessage(), ex);
    }
}

Example 9 with StatementClient

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

the class DataCenterClient method execute.

/**
 * Execute a query by client session.
 *
 * @param query statement.
 * @return a statement client instance.
 * @throws SQLException when new a statement client failed.
 */
private StatementClient execute(String query) throws SQLException {
    try (StatementClient client = DataCenterStatementClient.newStatementClient(this.httpClient, this.clientSession, query, this.globalQueryIdGenerator.createId())) {
        while (client.isRunning()) {
            client.advance();
        }
        verify(client.isFinished());
        QueryStatusInfo results = client.finalStatusInfo();
        if (results.getError() != null) {
            throw resultsException(results);
        }
        return client;
    } catch (RuntimeException ex) {
        throw new SQLException(ex.getMessage(), ex);
    }
}

Example 10 with StatementClient

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

the class PrestoStatement method partialCancel.

public void partialCancel() throws SQLException {
    checkOpen();
    StatementClient client = executingClient.get();
    if (client != null) {
        client.cancelLeafStage();
    } else {
        PrestoResultSet resultSet = currentResult.get();
        if (resultSet != null) {
            resultSet.partialCancel();
        }
    }
}