Search in sources :

Example 46 with GetTableRequest

use of com.amazonaws.athena.connector.lambda.metadata.GetTableRequest 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));
}
Also used : Table(com.amazonaws.services.glue.model.Table) ArrayList(java.util.ArrayList) StorageDescriptor(com.amazonaws.services.glue.model.StorageDescriptor) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) GetTableRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableRequest) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) GetTableLayoutResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse) Column(com.amazonaws.services.glue.model.Column) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) GetTableResult(com.amazonaws.services.glue.model.GetTableResult) Test(org.junit.Test)

Example 47 with GetTableRequest

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

the class DynamoDBMetadataHandlerTest method doGetEmptyTable.

@Test
public void doGetEmptyTable() throws Exception {
    when(glueClient.getTable(any())).thenThrow(new AmazonServiceException(""));
    GetTableRequest req = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, TEST_TABLE_2_NAME);
    GetTableResponse res = handler.doGetTable(allocator, req);
    logger.info("doGetEmptyTable - {}", res.getSchema());
    assertThat(res.getTableName(), equalTo(TEST_TABLE_2_NAME));
    assertThat(res.getSchema().getFields().size(), equalTo(2));
}
Also used : GetTableRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableRequest) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) AmazonServiceException(com.amazonaws.AmazonServiceException) Test(org.junit.Test)

Example 48 with GetTableRequest

use of com.amazonaws.athena.connector.lambda.metadata.GetTableRequest 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)));
}
Also used : AttributeValue(com.amazonaws.services.dynamodbv2.model.AttributeValue) Table(com.amazonaws.services.glue.model.Table) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) StorageDescriptor(com.amazonaws.services.glue.model.StorageDescriptor) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) GetTableRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableRequest) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) GetTableLayoutResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse) Column(com.amazonaws.services.glue.model.Column) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) ValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet) EquatableValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.EquatableValueSet) GetTableResult(com.amazonaws.services.glue.model.GetTableResult) Test(org.junit.Test)

Example 49 with GetTableRequest

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

the class DynamoDBRecordHandlerTest method testStructWithNullFromDdbTable.

@Test
public void testStructWithNullFromDdbTable() throws Exception {
    when(glueClient.getTable(any())).thenThrow(new EntityNotFoundException(""));
    TableName tableName = new TableName(DEFAULT_SCHEMA, TEST_TABLE4);
    GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName);
    GetTableResponse getTableResponse = metadataHandler.doGetTable(allocator, getTableRequest);
    logger.info("testStructWithNullFromGlueTable: GetTableResponse[{}]", getTableResponse);
    logger.info("testStructWithNullFromGlueTable: GetTableResponse Schema[{}]", getTableResponse.getSchema());
    Schema schema4 = getTableResponse.getSchema();
    for (Field f : schema4.getFields()) {
        if (f.getName().equals("Col2")) {
            assertEquals(1, f.getChildren().size());
            assertTrue(f.getType() instanceof ArrowType.Struct);
        }
    }
    Split split = Split.newBuilder(SPILL_LOCATION, keyFactory.create()).add(TABLE_METADATA, TEST_TABLE4).add(SEGMENT_ID_PROPERTY, "0").add(SEGMENT_COUNT_METADATA, "1").build();
    ReadRecordsRequest request = new ReadRecordsRequest(TEST_IDENTITY, TEST_CATALOG_NAME, TEST_QUERY_ID, TEST_TABLE_4_NAME, schema4, split, new Constraints(ImmutableMap.of()), // too big to spill
    100_000_000_000L, 100_000_000_000L);
    RecordResponse rawResponse = handler.doReadRecords(allocator, request);
    assertTrue(rawResponse instanceof ReadRecordsResponse);
    ReadRecordsResponse response = (ReadRecordsResponse) rawResponse;
    logger.info("testStructWithNullFromGlueTable: {}", BlockUtils.rowToString(response.getRecords(), 0));
    Block result = response.getRecords();
    assertEquals(1, result.getRowCount());
    assertEquals(schema4, result.getSchema());
    assertEquals("[Col0 : hashVal], [Col1 : {[field1 : someField1]}]", BlockUtils.rowToString(response.getRecords(), 0));
}
Also used : ReadRecordsResponse(com.amazonaws.athena.connector.lambda.records.ReadRecordsResponse) Schema(org.apache.arrow.vector.types.pojo.Schema) ArrowType(org.apache.arrow.vector.types.pojo.ArrowType) EntityNotFoundException(com.amazonaws.services.glue.model.EntityNotFoundException) RecordResponse(com.amazonaws.athena.connector.lambda.records.RecordResponse) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) GetTableRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableRequest) Field(org.apache.arrow.vector.types.pojo.Field) ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) Block(com.amazonaws.athena.connector.lambda.data.Block) Split(com.amazonaws.athena.connector.lambda.domain.Split) Test(org.junit.Test)

Example 50 with GetTableRequest

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

the class DynamoDBRecordHandlerTest method testStructWithNullFromGlueTable.

@Test
public void testStructWithNullFromGlueTable() throws Exception {
    List<Column> columns = new ArrayList<>();
    columns.add(new Column().withName("col0").withType("string"));
    columns.add(new Column().withName("col1").withType("struct"));
    Map<String, String> param = ImmutableMap.of(SOURCE_TABLE_PROPERTY, TEST_TABLE4, COLUMN_NAME_MAPPING_PROPERTY, "col1=Col1,col2=Col2");
    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, TEST_TABLE4);
    GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName);
    GetTableResponse getTableResponse = metadataHandler.doGetTable(allocator, getTableRequest);
    logger.info("testStructWithNullFromGlueTable: GetTableResponse[{}]", getTableResponse);
    logger.info("testStructWithNullFromGlueTable: GetTableResponse Schema[{}]", getTableResponse.getSchema());
    Schema schema4 = getTableResponse.getSchema();
    for (Field f : schema4.getFields()) {
        if (f.getName().equals("Col2")) {
            assertEquals(2, f.getChildren().size());
            assertTrue(f.getType() instanceof ArrowType.Struct);
        }
    }
    Split split = Split.newBuilder(SPILL_LOCATION, keyFactory.create()).add(TABLE_METADATA, TEST_TABLE4).add(SEGMENT_ID_PROPERTY, "0").add(SEGMENT_COUNT_METADATA, "1").build();
    ReadRecordsRequest request = new ReadRecordsRequest(TEST_IDENTITY, TEST_CATALOG_NAME, TEST_QUERY_ID, TEST_TABLE_4_NAME, schema4, split, new Constraints(ImmutableMap.of()), // too big to spill
    100_000_000_000L, 100_000_000_000L);
    RecordResponse rawResponse = handler.doReadRecords(allocator, request);
    assertTrue(rawResponse instanceof ReadRecordsResponse);
    ReadRecordsResponse response = (ReadRecordsResponse) rawResponse;
    logger.info("testStructWithNullFromGlueTable: {}", BlockUtils.rowToString(response.getRecords(), 0));
    Block result = response.getRecords();
    assertEquals(1, result.getRowCount());
    assertEquals(schema4, result.getSchema());
    assertEquals("[Col0 : hashVal], [Col1 : {[field1 : someField1]}]", BlockUtils.rowToString(response.getRecords(), 0));
}
Also used : Table(com.amazonaws.services.glue.model.Table) ReadRecordsResponse(com.amazonaws.athena.connector.lambda.records.ReadRecordsResponse) Schema(org.apache.arrow.vector.types.pojo.Schema) ArrayList(java.util.ArrayList) StorageDescriptor(com.amazonaws.services.glue.model.StorageDescriptor) ArrowType(org.apache.arrow.vector.types.pojo.ArrowType) Jackson.toJsonString(com.amazonaws.util.json.Jackson.toJsonString) RecordResponse(com.amazonaws.athena.connector.lambda.records.RecordResponse) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) GetTableRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableRequest) Field(org.apache.arrow.vector.types.pojo.Field) ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) Column(com.amazonaws.services.glue.model.Column) GetTableResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableResponse) Block(com.amazonaws.athena.connector.lambda.data.Block) Split(com.amazonaws.athena.connector.lambda.domain.Split) GetTableResult(com.amazonaws.services.glue.model.GetTableResult) Test(org.junit.Test)

Aggregations

GetTableRequest (com.amazonaws.athena.connector.lambda.metadata.GetTableRequest)51 Test (org.junit.Test)48 GetTableResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableResponse)33 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)30 ArrayList (java.util.ArrayList)13 GetTableResult (com.amazonaws.services.glue.model.GetTableResult)11 Column (com.amazonaws.services.glue.model.Column)10 Schema (org.apache.arrow.vector.types.pojo.Schema)10 StorageDescriptor (com.amazonaws.services.glue.model.StorageDescriptor)9 Table (com.amazonaws.services.glue.model.Table)8 InvocationOnMock (org.mockito.invocation.InvocationOnMock)8 SchemaBuilder (com.amazonaws.athena.connector.lambda.data.SchemaBuilder)7 Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)7 HashMap (java.util.HashMap)7 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)6 ResultSet (java.sql.ResultSet)6 AtomicInteger (java.util.concurrent.atomic.AtomicInteger)6 Field (org.apache.arrow.vector.types.pojo.Field)6 ReadRecordsRequest (com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest)5 ReadRecordsResponse (com.amazonaws.athena.connector.lambda.records.ReadRecordsResponse)4