Examples with CohortStatus org.pmiops.workbench.model.CohortStatus used on opensource projects

Example 1 with CohortStatus

use of org.pmiops.workbench.model.CohortStatus in project workbench by all-of-us.

the class CohortsControllerTest method testMaterializeCohortWithEverything.

@Test
public void testMaterializeCohortWithEverything() throws Exception {
    Cohort cohort = createDefaultCohort();
    cohort = cohortsController.createCohort(workspace.getNamespace(), workspace.getId(), cohort).getBody();
    MaterializeCohortRequest request = new MaterializeCohortRequest();
    request.setCohortName(cohort.getName());
    request.setPageSize(123);
    request.setPageToken("token");
    request.setCdrVersionName(CDR_VERSION_NAME);
    List<CohortStatus> statuses = ImmutableList.of(CohortStatus.INCLUDED, CohortStatus.NOT_REVIEWED);
    request.setStatusFilter(statuses);
    MaterializeCohortResponse response = new MaterializeCohortResponse();
    when(cohortMaterializationService.materializeCohort(null, searchRequest, request)).thenReturn(response);
    assertThat(cohortsController.materializeCohort(WORKSPACE_NAMESPACE, WORKSPACE_NAME, request).getBody()).isEqualTo(response);
}

Example 2 with CohortStatus

use of org.pmiops.workbench.model.CohortStatus 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;
}