Examples with PaginationToken org.pmiops.workbench.utils.PaginationToken used on opensource projects

Example 1 with PaginationToken

use of org.pmiops.workbench.utils.PaginationToken in project workbench by all-of-us.

the class CohortMaterializationServiceTest method testMaterializeCohortPaging.

@Test
public void testMaterializeCohortPaging() {
    MaterializeCohortRequest request = makeRequest(2);
    MaterializeCohortResponse response = cohortMaterializationService.materializeCohort(null, SearchRequests.allGenders(), request);
    assertPersonIds(response, 1L, 2L);
    assertThat(response.getNextPageToken()).isNotNull();
    request.setPageToken(response.getNextPageToken());
    MaterializeCohortResponse response2 = cohortMaterializationService.materializeCohort(null, SearchRequests.allGenders(), request);
    assertPersonIds(response2, 102246L);
    assertThat(response2.getNextPageToken()).isNull();
    try {
        // Pagination token doesn't match, this should fail.
        cohortMaterializationService.materializeCohort(null, SearchRequests.males(), request);
        fail("Exception expected");
    } catch (BadRequestException e) {
    // expected
    }
    PaginationToken token = PaginationToken.fromBase64(response.getNextPageToken());
    PaginationToken invalidToken = new PaginationToken(-1L, token.getParameterHash());
    request.setPageToken(invalidToken.toBase64());
    try {
        // Pagination token doesn't match, this should fail.
        cohortMaterializationService.materializeCohort(null, SearchRequests.males(), request);
        fail("Exception expected");
    } catch (BadRequestException e) {
    // expected
    }
}

Example 2 with PaginationToken

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