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

Example 6 with QueryResults

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

the class StatementResource method getQueryResults.

private static Response getQueryResults(Query query, Optional<Long> token, UriInfo uriInfo, Duration wait) throws InterruptedException {
    QueryResults queryResults;
    if (token.isPresent()) {
        queryResults = query.getResults(token.get(), uriInfo, wait);
    } else {
        queryResults = query.getNextResults(uriInfo, wait);
    }
    ResponseBuilder response = Response.ok(queryResults);
    // add set session properties
    query.getSetSessionProperties().entrySet().forEach(entry -> response.header(PRESTO_SET_SESSION, entry.getKey() + '=' + entry.getValue()));
    // add clear session properties
    query.getResetSessionProperties().forEach(name -> response.header(PRESTO_CLEAR_SESSION, name));
    // add added prepare statements
    for (Entry<String, String> entry : query.getAddedPreparedStatements().entrySet()) {
        String encodedKey = urlEncode(entry.getKey());
        String encodedValue = urlEncode(entry.getValue());
        response.header(PRESTO_ADDED_PREPARE, encodedKey + '=' + encodedValue);
    }
    // add deallocated prepare statements
    for (String name : query.getDeallocatedPreparedStatements()) {
        response.header(PRESTO_DEALLOCATED_PREPARE, urlEncode(name));
    }
    // add new transaction ID
    query.getStartedTransactionId().ifPresent(transactionId -> response.header(PRESTO_STARTED_TRANSACTION_ID, transactionId));
    // add clear transaction ID directive
    if (query.isClearTransactionId()) {
        response.header(PRESTO_CLEAR_TRANSACTION_ID, true);
    }
    return response.build();
}

Example 7 with QueryResults

use of com.facebook.presto.client.QueryResults 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();
}

Example 8 with QueryResults

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

the class ColumnCache method queryPartitions.

private List<HivePartition> queryPartitions(String query) {
    final ImmutableList.Builder<HivePartition> cache = ImmutableList.builder();
    final Map<Column, List<Object>> objects = Maps.newHashMap();
    QueryRunner queryRunner = queryRunnerFactory.create();
    QueryClient queryClient = new QueryClient(queryRunner, io.dropwizard.util.Duration.seconds(60), query);
    try {
        queryClient.executeWith(new Function<StatementClient, Void>() {

            @Nullable
            @Override
            public Void apply(StatementClient client) {
                QueryResults results = client.current();
                if (results.getData() != null && results.getColumns() != null) {
                    final List<Column> columns = results.getColumns();
                    for (Column column : columns) {
                        objects.put(column, Lists.newArrayList());
                    }
                    for (List<Object> row : results.getData()) {
                        for (int i = 0; i < row.size(); i++) {
                            Column column = columns.get(i);
                            objects.get(column).add(row.get(i));
                        }
                    }
                }
                return null;
            }
        });
    } catch (QueryClient.QueryTimeOutException e) {
        log.error("Caught timeout loading columns", e);
    }
    for (Map.Entry<Column, List<Object>> entry : objects.entrySet()) {
        cache.add(HivePartition.fromColumn(entry.getKey(), entry.getValue()));
    }
    return cache.build();
}

Example 9 with QueryResults

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

the class ColumnCache method queryColumns.

private List<HiveColumn> queryColumns(String query) {
    final ImmutableList.Builder<HiveColumn> cache = ImmutableList.builder();
    QueryRunner queryRunner = queryRunnerFactory.create();
    QueryClient queryClient = new QueryClient(queryRunner, io.dropwizard.util.Duration.seconds(60), query);
    try {
        queryClient.executeWith(new Function<StatementClient, Void>() {

            @Nullable
            @Override
            public Void apply(StatementClient client) {
                QueryResults results = client.current();
                if (results.getData() != null) {
                    for (List<Object> row : results.getData()) {
                        Column column = new Column((String) row.get(0), (String) row.get(1), new ClientTypeSignature(TypeSignature.parseTypeSignature((String) row.get(1))));
                        boolean isNullable = (Boolean) row.get(2);
                        boolean isPartition = (Boolean) row.get(3);
                        cache.add(HiveColumn.fromColumn(column, isNullable, isPartition));
                    }
                }
                return null;
            }
        });
    } catch (QueryClient.QueryTimeOutException e) {
        log.error("Caught timeout loading columns", e);
    }
    return cache.build();
}

Example 10 with QueryResults

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

the class PreviewTableCache method queryRows.

private List<List<Object>> queryRows(String query) {
    final ImmutableList.Builder<List<Object>> cache = ImmutableList.builder();
    QueryRunner queryRunner = queryRunnerFactory.create();
    QueryClient queryClient = new QueryClient(queryRunner, io.dropwizard.util.Duration.seconds(60), query);
    try {
        queryClient.executeWith(new Function<StatementClient, Void>() {

            @Nullable
            @Override
            public Void apply(StatementClient client) {
                QueryResults results = client.current();
                if (results.getData() != null) {
                    cache.addAll(results.getData());
                }
                return null;
            }
        });
    } catch (QueryClient.QueryTimeOutException e) {
        log.error("Caught timeout loading columns", e);
    }
    return cache.build();
}