use of com.amazonaws.athena.connectors.timestream.query.SelectQueryBuilder in project aws-athena-query-federation by awslabs.
the class TimestreamRecordHandler method readWithConstraint.
/**
* Scans TimeStream.
*
* @see RecordHandler
*/
@Override
protected void readWithConstraint(BlockSpiller spiller, ReadRecordsRequest recordsRequest, QueryStatusChecker queryStatusChecker) {
TableName tableName = recordsRequest.getTableName();
SelectQueryBuilder queryBuilder = queryFactory.createSelectQueryBuilder(GlueMetadataHandler.VIEW_METADATA_FIELD);
String query = queryBuilder.withDatabaseName(tableName.getSchemaName()).withTableName(tableName.getTableName()).withProjection(recordsRequest.getSchema()).withConjucts(recordsRequest.getConstraints()).build();
logger.info("readWithConstraint: query[{}]", query);
GeneratedRowWriter rowWriter = buildRowWriter(recordsRequest);
String nextToken = null;
long numRows = 0;
do {
QueryResult queryResult = tsQuery.query(new QueryRequest().withQueryString(query).withNextToken(nextToken));
List<Row> data = queryResult.getRows();
if (data != null) {
numRows += data.size();
for (Row nextRow : data) {
spiller.writeRows((Block block, int rowNum) -> rowWriter.writeRow(block, rowNum, nextRow) ? 1 : 0);
}
}
nextToken = queryResult.getNextToken();
logger.info("readWithConstraint: numRows[{}]", numRows);
} while (nextToken != null && !nextToken.isEmpty());
}
Aggregations