Examples with BigQueryException com.google.cloud.bigquery.BigQueryException used on opensource projects

Example 6 with BigQueryException

use of com.google.cloud.bigquery.BigQueryException in project workbench by all-of-us.

the class CohortMaterializationService method materializeCohort.

public MaterializeCohortResponse materializeCohort(@Nullable CohortReview cohortReview, SearchRequest searchRequest, MaterializeCohortRequest request) {
    long offset = 0L;
    FieldSet fieldSet = request.getFieldSet();
    List<CohortStatus> statusFilter = request.getStatusFilter();
    String paginationToken = request.getPageToken();
    int pageSize = request.getPageSize();
    // TODO: add CDR version ID here
    Object[] paginationParameters = new Object[] { searchRequest, statusFilter };
    if (paginationToken != null) {
        PaginationToken token = PaginationToken.fromBase64(paginationToken);
        if (token.matchesParameters(paginationParameters)) {
            offset = token.getOffset();
        } else {
            throw new BadRequestException(String.format("Use of pagination token %s with new parameter values", paginationToken));
        }
    }
    int limit = pageSize + 1;
    if (statusFilter == null) {
        statusFilter = ALL_STATUSES;
    }
    ParticipantCriteria criteria;
    MaterializeCohortResponse response = new MaterializeCohortResponse();
    if (statusFilter.contains(CohortStatus.NOT_REVIEWED)) {
        Set<Long> participantIdsToExclude;
        if (statusFilter.size() < CohortStatus.values().length) {
            // Find the participant IDs that have statuses which *aren't* in the filter.
            Set<CohortStatus> statusesToExclude = Sets.difference(ImmutableSet.copyOf(CohortStatus.values()), ImmutableSet.copyOf(statusFilter));
            participantIdsToExclude = getParticipantIdsWithStatus(cohortReview, ImmutableList.copyOf(statusesToExclude));
        } else {
            participantIdsToExclude = ImmutableSet.of();
        }
        criteria = new ParticipantCriteria(searchRequest, participantIdsToExclude);
    } else {
        Set<Long> participantIds = getParticipantIdsWithStatus(cohortReview, statusFilter);
        if (participantIds.isEmpty()) {
            // return an empty response.
            return response;
        }
        criteria = new ParticipantCriteria(participantIds);
    }
    TableQueryAndConfig tableQueryAndConfig = getTableQueryAndConfig(fieldSet);
    QueryJobConfiguration jobConfiguration = fieldSetQueryBuilder.buildQuery(criteria, tableQueryAndConfig, limit, offset);
    QueryResult result;
    try {
        result = bigQueryService.executeQuery(bigQueryService.filterBigQueryConfig(jobConfiguration));
    } catch (BigQueryException e) {
        if (e.getCode() == HttpServletResponse.SC_SERVICE_UNAVAILABLE) {
            throw new ServerUnavailableException("BigQuery was temporarily unavailable, try again later", e);
        } else if (e.getCode() == HttpServletResponse.SC_FORBIDDEN) {
            throw new ForbiddenException("Access to the CDR is denied", e);
        } else {
            throw new ServerErrorException(String.format("An unexpected error occurred materializing the cohort with " + "query = (%s), params = (%s)", jobConfiguration.getQuery(), jobConfiguration.getNamedParameters()), e);
        }
    }
    Map<String, Integer> rm = bigQueryService.getResultMapper(result);
    int numResults = 0;
    boolean hasMoreResults = false;
    ArrayList<Object> results = new ArrayList<>();
    for (List<FieldValue> row : result.iterateAll()) {
        if (numResults == pageSize) {
            hasMoreResults = true;
            break;
        }
        Map<String, Object> resultMap = fieldSetQueryBuilder.extractResults(tableQueryAndConfig, row);
        results.add(resultMap);
        numResults++;
    }
    response.setResults(results);
    if (hasMoreResults) {
        // TODO: consider pagination based on cursor / values rather than offset
        PaginationToken token = PaginationToken.of(offset + pageSize, paginationParameters);
        response.setNextPageToken(token.toBase64());
    }
    return response;
}

Example 7 with BigQueryException

use of com.google.cloud.bigquery.BigQueryException in project presto by prestodb.

the class BigQuerySplitManager method createEmptyProjection.

private List<BigQuerySplit> createEmptyProjection(TableId tableId, int actualParallelism, Optional<String> filter) {
    try {
        long numberOfRows;
        if (filter.isPresent()) {
            // count the rows based on the filter
            String sql = bigQueryClient.createFormatSql(tableId, "COUNT(*)", new String[] { filter.get() });
            TableResult result = bigQueryClient.query(sql);
            numberOfRows = result.iterateAll().iterator().next().get(0).getLongValue();
        } else {
            // no filters, so we can take the value from the table info
            numberOfRows = bigQueryClient.getTable(tableId).getNumRows().longValue();
        }
        long rowsPerSplit = numberOfRows / actualParallelism;
        // need to be added to one of the split due to integer division
        long remainingRows = numberOfRows - (rowsPerSplit * actualParallelism);
        List<BigQuerySplit> splits = range(0, actualParallelism).mapToObj(ignored -> BigQuerySplit.emptyProjection(rowsPerSplit)).collect(toList());
        splits.set(0, BigQuerySplit.emptyProjection(rowsPerSplit + remainingRows));
        return splits;
    } catch (BigQueryException e) {
        throw new PrestoException(BIGQUERY_FAILED_TO_EXECUTE_QUERY, format("Failed to compute empty projection"), e);
    }
}