use of org.pmiops.workbench.cohortbuilder.TableQueryAndConfig in project workbench by all-of-us.
the class CohortMaterializationService method getTableQueryAndConfig.
private TableQueryAndConfig getTableQueryAndConfig(FieldSet fieldSet) {
TableQuery tableQuery;
if (fieldSet == null) {
tableQuery = new TableQuery();
tableQuery.setTableName(PERSON_TABLE);
tableQuery.setColumns(ImmutableList.of(PERSON_ID));
} else {
tableQuery = fieldSet.getTableQuery();
if (tableQuery == null) {
// TODO: support other kinds of field sets besides tableQuery
throw new BadRequestException("tableQuery must be specified in field sets");
}
String tableName = tableQuery.getTableName();
if (Strings.isNullOrEmpty(tableName)) {
throw new BadRequestException("Table name must be specified in field sets");
}
}
CdrBigQuerySchemaConfig cdrSchemaConfig = cdrSchemaConfigProvider.get();
TableConfig tableConfig = cdrSchemaConfig.cohortTables.get(tableQuery.getTableName());
if (tableConfig == null) {
throw new BadRequestException("Table " + tableQuery.getTableName() + " is not a valid " + "cohort table; valid tables are: " + cdrSchemaConfig.cohortTables.keySet().stream().sorted().collect(Collectors.joining(",")));
}
Map<String, ColumnConfig> columnMap = Maps.uniqueIndex(tableConfig.columns, columnConfig -> columnConfig.name);
List<String> columnNames = tableQuery.getColumns();
if (columnNames == null || columnNames.isEmpty()) {
// By default, return all columns on the table in question in our configuration.
tableQuery.setColumns(columnMap.keySet().stream().collect(Collectors.toList()));
} else {
for (String columnName : columnNames) {
// TODO: handle columns on foreign key tables
if (!columnMap.containsKey(columnName)) {
throw new BadRequestException("Unrecognized column name: " + columnName);
}
}
}
List<String> orderBy = tableQuery.getOrderBy();
if (orderBy == null || orderBy.isEmpty()) {
ColumnConfig primaryKey = findPrimaryKey(tableConfig);
if (PERSON_ID.equals(primaryKey)) {
tableQuery.setOrderBy(ImmutableList.of(PERSON_ID));
} else {
// TODO: consider having per-table default sort order based on e.g. timestamp
tableQuery.setOrderBy(ImmutableList.of(PERSON_ID, primaryKey.name));
}
} else {
for (String columnName : orderBy) {
if (columnName.toUpperCase().endsWith(DESCENDING_SUFFIX)) {
columnName = columnName.substring(0, columnName.length() - DESCENDING_SUFFIX.length());
}
if (!columnMap.containsKey(columnName)) {
throw new BadRequestException("Invalid column in orderBy: " + columnName);
}
}
}
return new TableQueryAndConfig(tableQuery, tableConfig, columnMap);
}
use of org.pmiops.workbench.cohortbuilder.TableQueryAndConfig 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;
}
Aggregations