Examples with GetTableLayoutRequest com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest used on opensource projects

Example 36 with GetTableLayoutRequest

use of com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest in project aws-athena-query-federation by awslabs.

the class ExampleMetadataHandlerTest method getPartitions.

@Test
public void getPartitions() throws Exception {
    if (!enableTests) {
        // We do this because until you complete the tutorial these tests will fail. When you attempt to publis
        // using ../toos/publish.sh ...  it will set the publishing flag and force these tests. This is how we
        // avoid breaking the build but still have a useful tutorial. We are also duplicateing this block
        // on purpose since this is a somewhat odd pattern.
        logger.info("getPartitions: Tests are disabled, to enable them set the 'publishing' environment variable " + "using maven clean install -Dpublishing=true");
        return;
    }
    logger.info("doGetTableLayout - enter");
    Schema tableSchema = SchemaBuilder.newBuilder().addIntField("day").addIntField("month").addIntField("year").build();
    Set<String> partitionCols = new HashSet<>();
    partitionCols.add("day");
    partitionCols.add("month");
    partitionCols.add("year");
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("day", SortedRangeSet.copyOf(Types.MinorType.INT.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.INT.getType(), 0)), false));
    constraintsMap.put("month", SortedRangeSet.copyOf(Types.MinorType.INT.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.INT.getType(), 0)), false));
    constraintsMap.put("year", SortedRangeSet.copyOf(Types.MinorType.INT.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.INT.getType(), 2000)), false));
    GetTableLayoutRequest req = null;
    GetTableLayoutResponse res = null;
    try {
        req = new GetTableLayoutRequest(fakeIdentity(), "queryId", "default", new TableName("schema1", "table1"), new Constraints(constraintsMap), tableSchema, partitionCols);
        res = handler.doGetTableLayout(allocator, req);
        logger.info("doGetTableLayout - {}", res);
        Block partitions = res.getPartitions();
        for (int row = 0; row < partitions.getRowCount() && row < 10; row++) {
            logger.info("doGetTableLayout:{} {}", row, BlockUtils.rowToString(partitions, row));
        }
        assertTrue(partitions.getRowCount() > 0);
        logger.info("doGetTableLayout: partitions[{}]", partitions.getRowCount());
    } finally {
        try {
            req.close();
            res.close();
        } catch (Exception ex) {
            logger.error("doGetTableLayout: ", ex);
        }
    }
    logger.info("doGetTableLayout - exit");
}

Example 37 with GetTableLayoutRequest

use of com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest in project aws-athena-query-federation by awslabs.

the class LambdaMetadataProvider method getTableLayout.

/**
 * This method builds and executes a GetTableLayoutRequest against the specified Lambda function.
 *
 * @param catalog the catalog name to be passed to Lambda
 * @param tableName the schema-qualified table name indicating the table whose layout should be retrieved
 * @param constraints the constraints to be applied to the request
 * @param schema the schema of the table in question
 * @param partitionCols the partition column names for the table in question
 * @param metadataFunction the name of the Lambda function to call
 * @param identity the identity of the caller
 * @return the response
 */
public static GetTableLayoutResponse getTableLayout(String catalog, TableName tableName, Constraints constraints, Schema schema, Set<String> partitionCols, String metadataFunction, FederatedIdentity identity) {
    String queryId = generateQueryId();
    log.info("Submitting GetTableLayoutRequest with ID " + queryId);
    try (GetTableLayoutRequest request = new GetTableLayoutRequest(identity, queryId, catalog, tableName, constraints, schema, partitionCols)) {
        log.info("Submitting request: {}", request);
        GetTableLayoutResponse response = (GetTableLayoutResponse) getService(metadataFunction, identity, catalog).call(request);
        log.info("Received response: {}", response);
        return response;
    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}

Example 38 with GetTableLayoutRequest

use of com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest in project foundry-athena-query-federation-connector by palantir.

the class PartitionFetcher method getAndWritePartitions.

private void getAndWritePartitions(BlockWriter blockWriter, GetTableLayoutRequest request, QueryStatusChecker queryStatusChecker) {
    CatalogLocator locator = FoundryAthenaObjectMapper.objectMapper().convertValue(request.getSchema().getCustomMetadata(), CatalogLocator.class);
    Optional<String> pageToken = Optional.empty();
    while (queryStatusChecker.isQueryRunning()) {
        GetPartitionsResponsePage page = metadataService.getPartitions(authProvider.getAuthHeader(), GetPartitionsRequest.builder().locator(locator).limit(PAGE_SIZE).pageToken(pageToken).build());
        page.getPartitions().forEach(partition -> blockWriter.writeRows((block, rowNum) -> {
            boolean matched = partition.get().entrySet().stream().map(fieldName -> fieldName.getValue().accept(new PartitionValueWriter(block, fieldName.getKey(), rowNum))).reduce(true, Boolean::logicalAnd);
            // if all fields passed then we wrote 1 row
            return matched ? 1 : 0;
        }));
        if (page.getNextPageToken().isPresent()) {
            pageToken = page.getNextPageToken();
        } else {
            return;
        }
    }
}

Example 39 with GetTableLayoutRequest

use of com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest in project aws-athena-query-federation by awslabs.

the class ConstraintSerializationTest method serializationTest.

@Test
public void serializationTest() throws Exception {
    logger.info("serializationTest - enter");
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col2", SortedRangeSet.copyOf(Types.MinorType.BIGINT.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.BIGINT.getType(), 950L)), false));
    constraintsMap.put("col3", SortedRangeSet.copyOf(Types.MinorType.BIT.getType(), ImmutableList.of(Range.equal(allocator, Types.MinorType.BIT.getType(), false)), false));
    constraintsMap.put("col4", SortedRangeSet.copyOf(Types.MinorType.FLOAT8.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.FLOAT8.getType(), 950.0D)), false));
    constraintsMap.put("col5", SortedRangeSet.copyOf(Types.MinorType.VARCHAR.getType(), ImmutableList.of(Range.equal(allocator, Types.MinorType.VARCHAR.getType(), "8"), Range.equal(allocator, Types.MinorType.VARCHAR.getType(), "9")), false));
    try (GetTableLayoutRequest req = new GetTableLayoutRequest(IdentityUtil.fakeIdentity(), "queryId", "default", new TableName("schema1", "table1"), new Constraints(constraintsMap), SchemaBuilder.newBuilder().build(), new HashSet<>())) {
        ObjectMapperUtil.assertSerialization(req);
    }
    logger.info("serializationTest - exit");
}

Example 40 with GetTableLayoutRequest

use of com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest in project aws-athena-query-federation by awslabs.

the class HbaseMetadataHandlerTest method doGetTableLayout.

@Test
public void doGetTableLayout() throws Exception {
    GetTableLayoutRequest req = new GetTableLayoutRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME, new Constraints(new HashMap<>()), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET);
    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, req);
    logger.info("doGetTableLayout - {}", res);
    Block partitions = res.getPartitions();
    for (int row = 0; row < partitions.getRowCount() && row < 10; row++) {
        logger.info("doGetTableLayout:{} {}", row, BlockUtils.rowToString(partitions, row));
    }
    assertTrue(partitions.getRowCount() > 0);
}