Search in sources :

Example 56 with MaterializeCohortResponse

use of org.pmiops.workbench.model.MaterializeCohortResponse 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;
}
Also used : ServerUnavailableException(org.pmiops.workbench.exceptions.ServerUnavailableException) ArrayList(java.util.ArrayList) MaterializeCohortResponse(org.pmiops.workbench.model.MaterializeCohortResponse) FieldSet(org.pmiops.workbench.model.FieldSet) ParticipantIdAndCohortStatus(org.pmiops.workbench.db.model.ParticipantIdAndCohortStatus) CohortStatus(org.pmiops.workbench.model.CohortStatus) QueryResult(com.google.cloud.bigquery.QueryResult) FieldValue(com.google.cloud.bigquery.FieldValue) QueryJobConfiguration(com.google.cloud.bigquery.QueryJobConfiguration) PaginationToken(org.pmiops.workbench.utils.PaginationToken) ForbiddenException(org.pmiops.workbench.exceptions.ForbiddenException) TableQueryAndConfig(org.pmiops.workbench.cohortbuilder.TableQueryAndConfig) BadRequestException(org.pmiops.workbench.exceptions.BadRequestException) ParticipantCriteria(org.pmiops.workbench.cohortbuilder.ParticipantCriteria) BigQueryException(com.google.cloud.bigquery.BigQueryException) ServerErrorException(org.pmiops.workbench.exceptions.ServerErrorException)

Aggregations

MaterializeCohortResponse (org.pmiops.workbench.model.MaterializeCohortResponse)56 Test (org.junit.Test)52 BigQueryBaseTest (org.pmiops.workbench.api.BigQueryBaseTest)48 FieldSet (org.pmiops.workbench.model.FieldSet)36 TableQuery (org.pmiops.workbench.model.TableQuery)35 ColumnFilter (org.pmiops.workbench.model.ColumnFilter)29 MaterializeCohortRequest (org.pmiops.workbench.model.MaterializeCohortRequest)17 BigDecimal (java.math.BigDecimal)8 Cohort (org.pmiops.workbench.model.Cohort)6 DataJpaTest (org.springframework.boot.test.autoconfigure.orm.jpa.DataJpaTest)4 BadRequestException (org.pmiops.workbench.exceptions.BadRequestException)3 CohortReview (org.pmiops.workbench.db.model.CohortReview)2 CohortStatus (org.pmiops.workbench.model.CohortStatus)2 PaginationToken (org.pmiops.workbench.utils.PaginationToken)2 BigQueryException (com.google.cloud.bigquery.BigQueryException)1 FieldValue (com.google.cloud.bigquery.FieldValue)1 QueryJobConfiguration (com.google.cloud.bigquery.QueryJobConfiguration)1 QueryResult (com.google.cloud.bigquery.QueryResult)1 Gson (com.google.gson.Gson)1 JsonSyntaxException (com.google.gson.JsonSyntaxException)1