Examples with ValueSet com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet used on opensource projects

Example 26 with ValueSet

use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.

the class RdsTableProvider method readWithConstraint.

/**
 * Calls DescribeDBInstances on the AWS RDS Client returning all DB Instances that match the supplied predicate and attempting
 * to push down certain predicates (namely queries for specific DB Instance) to EC2.
 *
 * @See TableProvider
 */
@Override
public void readWithConstraint(BlockSpiller spiller, ReadRecordsRequest recordsRequest, QueryStatusChecker queryStatusChecker) {
    boolean done = false;
    DescribeDBInstancesRequest request = new DescribeDBInstancesRequest();
    ValueSet idConstraint = recordsRequest.getConstraints().getSummary().get("instance_id");
    if (idConstraint != null && idConstraint.isSingleValue()) {
        request.setDBInstanceIdentifier(idConstraint.getSingleValue().toString());
    }
    while (!done) {
        DescribeDBInstancesResult response = rds.describeDBInstances(request);
        for (DBInstance instance : response.getDBInstances()) {
            instanceToRow(instance, spiller);
        }
        request.setMarker(response.getMarker());
        if (response.getMarker() == null || !queryStatusChecker.isQueryRunning()) {
            done = true;
        }
    }
}

Example 27 with ValueSet

use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.

the class EbsTableProvider method readWithConstraint.

/**
 * Calls DescribeVolumes on the AWS EC2 Client returning all volumes that match the supplied predicate and attempting
 * to push down certain predicates (namely queries for specific volumes) to EC2.
 *
 * @See TableProvider
 */
@Override
public void readWithConstraint(BlockSpiller spiller, ReadRecordsRequest recordsRequest, QueryStatusChecker queryStatusChecker) {
    boolean done = false;
    DescribeVolumesRequest request = new DescribeVolumesRequest();
    ValueSet idConstraint = recordsRequest.getConstraints().getSummary().get("id");
    if (idConstraint != null && idConstraint.isSingleValue()) {
        request.setVolumeIds(Collections.singletonList(idConstraint.getSingleValue().toString()));
    }
    while (!done) {
        DescribeVolumesResult response = ec2.describeVolumes(request);
        for (Volume volume : response.getVolumes()) {
            logger.info("readWithConstraint: {}", response);
            instanceToRow(volume, spiller);
        }
        request.setNextToken(response.getNextToken());
        if (response.getNextToken() == null || !queryStatusChecker.isQueryRunning()) {
            done = true;
        }
    }
}

Example 28 with ValueSet

use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.

the class AbstractTableProviderTest method readTableTest.

@Test
public void readTableTest() {
    GetTableRequest request = new GetTableRequest(identity, expectedQuery, expectedCatalog, expectedTableName);
    GetTableResponse response = provider.getTable(allocator, request);
    assertTrue(response.getSchema().getFields().size() > 1);
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put(idField, EquatableValueSet.newBuilder(allocator, Types.MinorType.VARCHAR.getType(), true, false).add(idValue).build());
    Constraints constraints = new Constraints(constraintsMap);
    ConstraintEvaluator evaluator = new ConstraintEvaluator(allocator, response.getSchema(), constraints);
    S3SpillLocation spillLocation = S3SpillLocation.newBuilder().withBucket("bucket").withPrefix("prefix").withSplitId(UUID.randomUUID().toString()).withQueryId(UUID.randomUUID().toString()).withIsDirectory(true).build();
    ReadRecordsRequest readRequest = new ReadRecordsRequest(identity, expectedCatalog, "queryId", expectedTableName, response.getSchema(), Split.newBuilder(spillLocation, keyFactory.create()).build(), constraints, 100_000_000, 100_000_000);
    SpillConfig spillConfig = SpillConfig.newBuilder().withSpillLocation(spillLocation).withMaxBlockBytes(3_000_000).withMaxInlineBlockBytes(0).withRequestId("queryid").withEncryptionKey(keyFactory.create()).build();
    setUpRead();
    BlockSpiller spiller = new S3BlockSpiller(amazonS3, spillConfig, allocator, response.getSchema(), evaluator);
    provider.readWithConstraint(spiller, readRequest, queryStatusChecker);
    validateRead(response.getSchema(), blockSpillReader, spiller.getSpillLocations(), spillConfig.getEncryptionKey());
}

Example 29 with ValueSet

use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.

the class HiveRecordHandlerTest method getSingleValueSet.

private ValueSet getSingleValueSet(Object value) {
    Range range = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(range.isSingleValue()).thenReturn(true);
    Mockito.when(range.getLow().getValue()).thenReturn(value);
    ValueSet valueSet = Mockito.mock(SortedRangeSet.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(valueSet.getRanges().getOrderedRanges()).thenReturn(Collections.singletonList(range));
    return valueSet;
}

Example 30 with ValueSet

use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.

the class ImpalaRecordHandlerTest method getSingleValueSet.

private ValueSet getSingleValueSet(Object value) {
    Range range = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(range.isSingleValue()).thenReturn(true);
    Mockito.when(range.getLow().getValue()).thenReturn(value);
    ValueSet valueSet = Mockito.mock(SortedRangeSet.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(valueSet.getRanges().getOrderedRanges()).thenReturn(Collections.singletonList(range));
    return valueSet;
}