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

Example 31 with GetTableRequest

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

the class GlueMetadataHandlerTest method doGetTable.

@Test
public void doGetTable() throws Exception {
    String sourceTable = "My-Table";
    Map<String, String> expectedParams = new HashMap<>();
    expectedParams.put(SOURCE_TABLE_PROPERTY, sourceTable);
    expectedParams.put(COLUMN_NAME_MAPPING_PROPERTY, "col2=Col2,col3=Col3, col4=Col4");
    expectedParams.put(DATETIME_FORMAT_MAPPING_PROPERTY, "col2=someformat2, col1=someformat1 ");
    List<Column> columns = new ArrayList<>();
    columns.add(new Column().withName("col1").withType("int").withComment("comment"));
    columns.add(new Column().withName("col2").withType("bigint").withComment("comment"));
    columns.add(new Column().withName("col3").withType("string").withComment("comment"));
    columns.add(new Column().withName("col4").withType("timestamp").withComment("comment"));
    columns.add(new Column().withName("col5").withType("date").withComment("comment"));
    columns.add(new Column().withName("col6").withType("timestamptz").withComment("comment"));
    columns.add(new Column().withName("col7").withType("timestamptz").withComment("comment"));
    Table mockTable = mock(Table.class);
    StorageDescriptor mockSd = mock(StorageDescriptor.class);
    when(mockTable.getName()).thenReturn(table);
    when(mockTable.getStorageDescriptor()).thenReturn(mockSd);
    when(mockTable.getParameters()).thenReturn(expectedParams);
    when(mockSd.getColumns()).thenReturn(columns);
    when(mockGlue.getTable(any(com.amazonaws.services.glue.model.GetTableRequest.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
        com.amazonaws.services.glue.model.GetTableRequest request = (com.amazonaws.services.glue.model.GetTableRequest) invocationOnMock.getArguments()[0];
        assertEquals(accountId, request.getCatalogId());
        assertEquals(schema, request.getDatabaseName());
        assertEquals(table, request.getName());
        GetTableResult mockResult = mock(GetTableResult.class);
        when(mockResult.getTable()).thenReturn(mockTable);
        return mockResult;
    });
    GetTableRequest req = new GetTableRequest(IdentityUtil.fakeIdentity(), queryId, catalog, new TableName(schema, table));
    GetTableResponse res = handler.doGetTable(allocator, req);
    logger.info("doGetTable - {}", res);
    assertTrue(res.getSchema().getFields().size() == 7);
    assertTrue(res.getSchema().getCustomMetadata().size() > 0);
    assertTrue(res.getSchema().getCustomMetadata().containsKey(DATETIME_FORMAT_MAPPING_PROPERTY));
    assertEquals(res.getSchema().getCustomMetadata().get(DATETIME_FORMAT_MAPPING_PROPERTY_NORMALIZED), "Col2=someformat2,col1=someformat1");
    assertEquals(sourceTable, getSourceTableName(res.getSchema()));
    // Verify column name mapping works
    assertNotNull(res.getSchema().findField("col1"));
    assertNotNull(res.getSchema().findField("Col2"));
    assertNotNull(res.getSchema().findField("Col3"));
    assertNotNull(res.getSchema().findField("Col4"));
    assertNotNull(res.getSchema().findField("col5"));
    assertNotNull(res.getSchema().findField("col6"));
    assertNotNull(res.getSchema().findField("col7"));
    // Verify types
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("col1").getType()).equals(Types.MinorType.INT));
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("Col2").getType()).equals(Types.MinorType.BIGINT));
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("Col3").getType()).equals(Types.MinorType.VARCHAR));
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("Col4").getType()).equals(Types.MinorType.DATEMILLI));
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("col5").getType()).equals(Types.MinorType.DATEDAY));
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("col6").getType()).equals(Types.MinorType.TIMESTAMPMILLITZ));
    assertTrue(Types.getMinorTypeForArrowType(res.getSchema().findField("col7").getType()).equals(Types.MinorType.TIMESTAMPMILLITZ));
}

Example 32 with GetTableRequest

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

the class HbaseMetadataHandlerTest method doGetTable.

/**
 * TODO: Add more types.
 */
@Test
public void doGetTable() throws Exception {
    List<Result> results = TestUtils.makeResults();
    ResultScanner mockScanner = mock(ResultScanner.class);
    when(mockScanner.iterator()).thenReturn(results.iterator());
    when(mockClient.scanTable(anyObject(), any(Scan.class), anyObject())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        ResultProcessor processor = (ResultProcessor) invocationOnMock.getArguments()[2];
        return processor.scan(mockScanner);
    });
    GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME);
    GetTableResponse res = handler.doGetTable(allocator, req);
    logger.info("doGetTable - {}", res);
    Schema expectedSchema = TestUtils.makeSchema().addField(HbaseSchemaUtils.ROW_COLUMN_NAME, Types.MinorType.VARCHAR.getType()).build();
    assertEquals(expectedSchema.getFields().size(), res.getSchema().getFields().size());
}

Example 33 with GetTableRequest

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

the class GetTableRequestSerDeTest method beforeTest.

@Before
public void beforeTest() throws IOException {
    expected = new GetTableRequest(federatedIdentity, "test-query-id", "test-catalog", new TableName("test-schema", "test-table"));
    String expectedSerDeFile = utils.getResourceOrFail("serde/v2", "GetTableRequest.json");
    expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim();
}

Example 34 with GetTableRequest

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

the class HiveMuxMetadataHandlerTest method doGetTable.

@Test
public void doGetTable() throws Exception {
    GetTableRequest getTableRequest = Mockito.mock(GetTableRequest.class);
    Mockito.when(getTableRequest.getCatalogName()).thenReturn("metaHive");
    this.jdbcMetadataHandler.doGetTable(this.allocator, getTableRequest);
    Mockito.verify(this.hiveMetadataHandler, Mockito.times(1)).doGetTable(Mockito.eq(this.allocator), Mockito.eq(getTableRequest));
}

Example 35 with GetTableRequest

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

the class MySqlMuxJdbcMetadataHandlerTest method doGetTable.

@Test
public void doGetTable() {
    GetTableRequest getTableRequest = Mockito.mock(GetTableRequest.class);
    Mockito.when(getTableRequest.getCatalogName()).thenReturn("fakedatabase");
    this.jdbcMetadataHandler.doGetTable(this.allocator, getTableRequest);
    Mockito.verify(this.mySqlMetadataHandler, Mockito.times(1)).doGetTable(Mockito.eq(this.allocator), Mockito.eq(getTableRequest));
}