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

Example 71 with GetTableLayoutRequest

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

the class DataLakeGen2MetadataHandlerTest method doGetTableLayoutWithNoPartitions.

@Test
public void doGetTableLayoutWithNoPartitions() throws Exception {
    BlockAllocator blockAllocator = new BlockAllocatorImpl();
    Constraints constraints = Mockito.mock(Constraints.class);
    TableName tableName = new TableName("testSchema", "testTable");
    Schema partitionSchema = this.dataLakeGen2MetadataHandler.getPartitionSchema("testCatalogName");
    Set<String> partitionCols = partitionSchema.getFields().stream().map(Field::getName).collect(Collectors.toSet());
    GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, constraints, partitionSchema, partitionCols);
    GetTableLayoutResponse getTableLayoutResponse = this.dataLakeGen2MetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
    List<String> actualValues = new ArrayList<>();
    for (int i = 0; i < getTableLayoutResponse.getPartitions().getRowCount(); i++) {
        actualValues.add(BlockUtils.rowToString(getTableLayoutResponse.getPartitions(), i));
    }
    Assert.assertEquals(Collections.singletonList("[PARTITION_NUMBER : 0]"), actualValues);
    SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
    expectedSchemaBuilder.addField(FieldBuilder.newBuilder(DataLakeGen2MetadataHandler.PARTITION_NUMBER, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
    Schema expectedSchema = expectedSchemaBuilder.build();
    Assert.assertEquals(expectedSchema, getTableLayoutResponse.getPartitions().getSchema());
    Assert.assertEquals(tableName, getTableLayoutResponse.getTableName());
}

Example 72 with GetTableLayoutRequest

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

the class DataLakeGen2MuxMetadataHandlerTest method doGetTableLayout.

@Test
public void doGetTableLayout() throws Exception {
    GetTableLayoutRequest getTableLayoutRequest = Mockito.mock(GetTableLayoutRequest.class);
    Mockito.when(getTableLayoutRequest.getTableName()).thenReturn(new TableName("testSchema", "testTable"));
    Mockito.when(getTableLayoutRequest.getCatalogName()).thenReturn("fakedatabase");
    this.jdbcMetadataHandler.doGetTableLayout(this.allocator, getTableLayoutRequest);
    Mockito.verify(this.dataLakeGen2MetadataHandler, Mockito.times(1)).doGetTableLayout(Mockito.eq(this.allocator), Mockito.eq(getTableLayoutRequest));
}

Example 73 with GetTableLayoutRequest

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

the class DynamoDBMetadataHandlerTest method validateSourceTableNamePropagation.

@Test
public void validateSourceTableNamePropagation() throws Exception {
    List<Column> columns = new ArrayList<>();
    columns.add(new Column().withName("col1").withType("int"));
    columns.add(new Column().withName("col2").withType("bigint"));
    columns.add(new Column().withName("col3").withType("string"));
    Map<String, String> param = ImmutableMap.of(SOURCE_TABLE_PROPERTY, TEST_TABLE, COLUMN_NAME_MAPPING_PROPERTY, "col1=Col1 , col2=Col2 ,col3=Col3", DATETIME_FORMAT_MAPPING_PROPERTY, "col1=datetime1,col3=datetime3 ");
    Table table = new Table().withParameters(param).withPartitionKeys().withStorageDescriptor(new StorageDescriptor().withColumns(columns));
    GetTableResult mockResult = new GetTableResult().withTable(table);
    when(glueClient.getTable(any())).thenReturn(mockResult);
    TableName tableName = new TableName(DEFAULT_SCHEMA, "glueTableForTestTable");
    GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName);
    GetTableResponse getTableResponse = handler.doGetTable(allocator, getTableRequest);
    logger.info("validateSourceTableNamePropagation: GetTableResponse[{}]", getTableResponse);
    Map<String, String> customMetadata = getTableResponse.getSchema().getCustomMetadata();
    assertThat(customMetadata.get(SOURCE_TABLE_PROPERTY), equalTo(TEST_TABLE));
    assertThat(customMetadata.get(DATETIME_FORMAT_MAPPING_PROPERTY_NORMALIZED), equalTo("Col1=datetime1,Col3=datetime3"));
    GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, new Constraints(ImmutableMap.of()), getTableResponse.getSchema(), Collections.EMPTY_SET);
    GetTableLayoutResponse getTableLayoutResponse = handler.doGetTableLayout(allocator, getTableLayoutRequest);
    logger.info("validateSourceTableNamePropagation: GetTableLayoutResponse[{}]", getTableLayoutResponse);
    assertThat(getTableLayoutResponse.getPartitions().getSchema().getCustomMetadata().get(TABLE_METADATA), equalTo(TEST_TABLE));
}

Example 74 with GetTableLayoutRequest

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

the class DynamoDBMetadataHandlerTest method doGetTableLayoutScanWithTypeOverride.

@Test
public void doGetTableLayoutScanWithTypeOverride() throws Exception {
    List<Column> columns = new ArrayList<>();
    columns.add(new Column().withName("col1").withType("int"));
    columns.add(new Column().withName("col2").withType("timestamptz"));
    columns.add(new Column().withName("col3").withType("string"));
    Map<String, String> param = ImmutableMap.of(SOURCE_TABLE_PROPERTY, TEST_TABLE, COLUMN_NAME_MAPPING_PROPERTY, "col1=Col1", DATETIME_FORMAT_MAPPING_PROPERTY, "col1=datetime1,col3=datetime3 ");
    Table table = new Table().withParameters(param).withPartitionKeys().withStorageDescriptor(new StorageDescriptor().withColumns(columns));
    GetTableResult mockResult = new GetTableResult().withTable(table);
    when(glueClient.getTable(any())).thenReturn(mockResult);
    TableName tableName = new TableName(DEFAULT_SCHEMA, "glueTableForTestTable");
    GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName);
    GetTableResponse getTableResponse = handler.doGetTable(allocator, getTableRequest);
    logger.info("validateSourceTableNamePropagation: GetTableResponse[{}]", getTableResponse);
    Map<String, String> customMetadata = getTableResponse.getSchema().getCustomMetadata();
    assertThat(customMetadata.get(SOURCE_TABLE_PROPERTY), equalTo(TEST_TABLE));
    assertThat(customMetadata.get(DATETIME_FORMAT_MAPPING_PROPERTY_NORMALIZED), equalTo("Col1=datetime1,col3=datetime3"));
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col3", EquatableValueSet.newBuilder(allocator, new ArrowType.Bool(), true, true).add(true).build());
    constraintsMap.put("col2", EquatableValueSet.newBuilder(allocator, new ArrowType.Bool(), true, true).add(true).build());
    GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName, new Constraints(constraintsMap), getTableResponse.getSchema(), Collections.EMPTY_SET);
    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, getTableLayoutRequest);
    logger.info("doGetTableLayoutScanWithTypeOverride schema - {}", res.getPartitions().getSchema());
    logger.info("doGetTableLayoutScanWithTypeOverride partitions - {}", res.getPartitions());
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(PARTITION_TYPE_METADATA), equalTo(SCAN_PARTITION_TYPE));
    // no hash key constraints, so look for segment count column
    assertThat(res.getPartitions().getSchema().findField(SEGMENT_COUNT_METADATA) != null, is(true));
    assertThat(res.getPartitions().getRowCount(), equalTo(1));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(NON_KEY_FILTER_METADATA), equalTo("(#col3 = :v0 OR attribute_not_exists(#col3) OR #col3 = :v1)"));
    ImmutableMap<String, String> expressionNames = ImmutableMap.of("#col3", "col3", "#col2", "col2");
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_NAMES_METADATA), equalTo(Jackson.toJsonString(expressionNames)));
    ImmutableMap<String, AttributeValue> expressionValues = ImmutableMap.of(":v0", ItemUtils.toAttributeValue(true), ":v1", ItemUtils.toAttributeValue(null));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(Jackson.toJsonString(expressionValues)));
}

Example 75 with GetTableLayoutRequest

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

the class DynamoDBMetadataHandlerTest method doGetTableLayoutScan.

@Test
public void doGetTableLayoutScan() throws Exception {
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col_3", EquatableValueSet.newBuilder(allocator, new ArrowType.Bool(), true, true).add(true).build());
    GetTableLayoutRequest req = new GetTableLayoutRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, new TableName(TEST_CATALOG_NAME, TEST_TABLE), new Constraints(constraintsMap), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET);
    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, req);
    logger.info("doGetTableLayout schema - {}", res.getPartitions().getSchema());
    logger.info("doGetTableLayout partitions - {}", res.getPartitions());
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(PARTITION_TYPE_METADATA), equalTo(SCAN_PARTITION_TYPE));
    // no hash key constraints, so look for segment count column
    assertThat(res.getPartitions().getSchema().findField(SEGMENT_COUNT_METADATA) != null, is(true));
    assertThat(res.getPartitions().getRowCount(), equalTo(1));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(NON_KEY_FILTER_METADATA), equalTo("(#col_3 = :v0 OR attribute_not_exists(#col_3) OR #col_3 = :v1)"));
    ImmutableMap<String, String> expressionNames = ImmutableMap.of("#col_3", "col_3");
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_NAMES_METADATA), equalTo(Jackson.toJsonString(expressionNames)));
    ImmutableMap<String, AttributeValue> expressionValues = ImmutableMap.of(":v0", ItemUtils.toAttributeValue(true), ":v1", ItemUtils.toAttributeValue(null));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(Jackson.toJsonString(expressionValues)));
}