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

Example 1 with StatementClient

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

the class BenchmarkQueryRunner method execute.

private StatementStats execute(ClientSession session, String name, String query) {
    // start query
    StatementClient client = new StatementClient(httpClient, queryResultsCodec, session, query);
    // read query output
    while (client.isValid() && client.advance()) {
    // we do not process the output
    }
    // verify final state
    if (client.isClosed()) {
        throw new IllegalStateException("Query aborted by user");
    }
    if (client.isGone()) {
        throw new IllegalStateException("Query is gone (server restarted?)");
    }
    QueryError resultsError = client.finalResults().getError();
    if (resultsError != null) {
        RuntimeException cause = null;
        if (resultsError.getFailureInfo() != null) {
            cause = resultsError.getFailureInfo().toException();
        }
        throw new BenchmarkDriverExecutionException(format("Query %s failed: %s", name, resultsError.getMessage()), cause);
    }
    return client.finalResults().getStats();
}

Example 2 with StatementClient

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

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 3 with StatementClient

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

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();
        }
    }
}

Example 4 with StatementClient

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

the class PrestoStatement method internalExecute.

final boolean internalExecute(String sql) throws SQLException {
    clearCurrentResults();
    checkOpen();
    StatementClient client = null;
    PrestoResultSet resultSet = null;
    boolean intercepted = false;
    try {
        WarningsManager warningsManager = new WarningsManager();
        currentWarningsManager.set(Optional.of(warningsManager));
        int statementDepth = this.statementDepth.incrementAndGet();
        boolean shouldIntercept = !connection().getQueryInterceptorInstances().isEmpty() && statementDepth == 1;
        if (shouldIntercept) {
            Optional<PrestoResultSet> newResultSet = connection().invokeQueryInterceptorsPre(sql, this);
            if (newResultSet.isPresent()) {
                resultSet = newResultSet.get();
            }
        }
        // Check if no resultSet is returned from an interceptor
        if (resultSet != null) {
            currentResult.set(resultSet);
            intercepted = true;
        } else {
            client = connection().startQuery(sql, getStatementSessionProperties());
            if (client.isFinished()) {
                QueryStatusInfo finalStatusInfo = client.finalStatusInfo();
                if (finalStatusInfo.getError() != null) {
                    throw resultsException(finalStatusInfo);
                }
            }
            executingClient.set(client);
            resultSet = new PrestoResultSet(this, client, maxRows.get(), progressConsumer, warningsManager);
            for (Map.Entry<String, SelectedRole> entry : client.getSetRoles().entrySet()) {
                connection.get().setRole(entry.getKey(), entry.getValue());
            }
        }
        // check if this is a query
        if (intercepted || client.currentStatusInfo().getUpdateType() == null) {
            currentResult.set(resultSet);
            if (shouldIntercept) {
                resultSet = connection().invokeQueryInterceptorsPost(sql, this, resultSet);
                verifyNotNull(resultSet, "invokeQueryInterceptorsPost should never return a null ResultSet");
                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 {
        this.statementDepth.decrementAndGet();
        executingClient.set(null);
        if (currentResult.get() == null) {
            if (resultSet != null) {
                resultSet.close();
            }
            if (client != null) {
                client.close();
            }
        }
    }
}

Example 5 with StatementClient

use of com.facebook.presto.client.StatementClient in project airpal by airbnb.

the class Execution method doExecute.

private Job doExecute() throws ExecutionFailureException {
    final String userQuery = QUERY_SPLITTER.splitToList(getJob().getQuery()).get(0);
    final JobOutputBuilder outputBuilder;
    job.setQueryStats(createNoOpQueryStats());
    try {
        outputBuilder = outputBuilderFactory.forJob(job);
    } catch (IOException e) {
        throw new ExecutionFailureException(job, "Could not create output builder for job", e);
    } catch (InvalidQueryException e) {
        throw new ExecutionFailureException(job, e.getMessage(), e);
    }
    final Persistor persistor = persistorFactory.getPersistor(job, job.getOutput());
    final String query = job.getOutput().processQuery(userQuery);
    if (!persistor.canPersist(authorizer)) {
        throw new ExecutionFailureException(job, "Not authorized to create tables", null);
    }
    final List<List<Object>> outputPreview = new ArrayList<>(maxRowsPreviewOutput);
    final Set<Table> tables = new HashSet<>();
    try {
        tables.addAll(authorizer.tablesUsedByQuery(query));
    } catch (ParsingException e) {
        job.setError(new QueryError(e.getMessage(), null, -1, null, null, new ErrorLocation(e.getLineNumber(), e.getColumnNumber()), null));
        throw new ExecutionFailureException(job, "Invalid query, could not parse", e);
    }
    if (!authorizer.isAuthorizedRead(tables)) {
        job.setQueryStats(createNoOpQueryStats());
        throw new ExecutionFailureException(job, "Cannot access tables", null);
    }
    QueryClient queryClient = new QueryClient(queryRunner, timeout, query);
    try {
        queryClient.executeWith(new Function<StatementClient, Void>() {

            @Nullable
            @Override
            public Void apply(@Nullable StatementClient client) {
                if (client == null) {
                    return null;
                }
                QueryResults results = client.current();
                List<Column> resultColumns = null;
                JobState jobState = null;
                QueryError queryError = null;
                QueryStats queryStats = null;
                if (isCancelled) {
                    throw new ExecutionFailureException(job, "Query was cancelled", null);
                }
                if (results.getError() != null) {
                    queryError = results.getError();
                    jobState = JobState.FAILED;
                }
                if ((results.getInfoUri() != null) && (jobState != JobState.FAILED)) {
                    BasicQueryInfo queryInfo = queryInfoClient.from(results.getInfoUri());
                    if (queryInfo != null) {
                        queryStats = queryInfo.getQueryStats();
                    }
                }
                if (results.getStats() != null) {
                    jobState = JobState.fromStatementState(results.getStats().getState());
                }
                try {
                    if (results.getColumns() != null) {
                        resultColumns = results.getColumns();
                        outputBuilder.addColumns(resultColumns);
                    }
                    if (results.getData() != null) {
                        List<List<Object>> resultsData = ImmutableList.copyOf(results.getData());
                        for (List<Object> row : resultsData) {
                            outputBuilder.addRow(row);
                        }
                    }
                } catch (FileTooLargeException e) {
                    throw new ExecutionFailureException(job, "Output file exceeded maximum configured filesize", e);
                }
                rlUpdateJobInfo(tables, resultColumns, queryStats, jobState, queryError, outputPreview);
                return null;
            }
        });
    } catch (QueryTimeOutException e) {
        throw new ExecutionFailureException(job, format("Query exceeded maximum execution time of %s minutes", Duration.millis(e.getElapsedMs()).getStandardMinutes()), e);
    }
    QueryResults finalResults = queryClient.finalResults();
    if (finalResults != null && finalResults.getInfoUri() != null) {
        BasicQueryInfo queryInfo = queryInfoClient.from(finalResults.getInfoUri());
        if (queryInfo != null) {
            updateJobInfo(null, null, queryInfo.getQueryStats(), JobState.fromStatementState(finalResults.getStats().getState()), finalResults.getError(), outputPreview, true);
        }
    }
    if (job.getState() != JobState.FAILED) {
        URI location = persistor.persist(outputBuilder, job);
        if (location != null) {
            job.getOutput().setLocation(location);
        }
    } else {
        throw new ExecutionFailureException(job, null, null);
    }
    return getJob();
}