use of com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse in project aws-athena-query-federation by awslabs.
the class CompositeHandlerTest method setUp.
@Before
public void setUp() throws Exception {
logger.info("{}: enter", testName.getMethodName());
allocator = new BlockAllocatorImpl();
objectMapper = ObjectMapperFactory.create(allocator);
mockMetadataHandler = mock(MetadataHandler.class);
mockRecordHandler = mock(RecordHandler.class);
schemaForRead = SchemaBuilder.newBuilder().addField("col1", new ArrowType.Int(32, true)).build();
when(mockMetadataHandler.doGetTableLayout(any(BlockAllocatorImpl.class), any(GetTableLayoutRequest.class))).thenReturn(new GetTableLayoutResponse("catalog", new TableName("schema", "table"), BlockUtils.newBlock(allocator, "col1", Types.MinorType.BIGINT.getType(), 1L)));
when(mockMetadataHandler.doListTables(any(BlockAllocatorImpl.class), any(ListTablesRequest.class))).thenReturn(new ListTablesResponse("catalog", Collections.singletonList(new TableName("schema", "table")), null));
when(mockMetadataHandler.doGetTable(any(BlockAllocatorImpl.class), any(GetTableRequest.class))).thenReturn(new GetTableResponse("catalog", new TableName("schema", "table"), SchemaBuilder.newBuilder().addStringField("col1").build()));
when(mockMetadataHandler.doListSchemaNames(any(BlockAllocatorImpl.class), any(ListSchemasRequest.class))).thenReturn(new ListSchemasResponse("catalog", Collections.singleton("schema1")));
when(mockMetadataHandler.doGetSplits(any(BlockAllocatorImpl.class), any(GetSplitsRequest.class))).thenReturn(new GetSplitsResponse("catalog", Split.newBuilder(null, null).build()));
when(mockMetadataHandler.doPing(any(PingRequest.class))).thenReturn(new PingResponse("catalog", "queryId", "type", 23, 2));
when(mockRecordHandler.doReadRecords(any(BlockAllocatorImpl.class), any(ReadRecordsRequest.class))).thenReturn(new ReadRecordsResponse("catalog", BlockUtils.newEmptyBlock(allocator, "col", new ArrowType.Int(32, true))));
compositeHandler = new CompositeHandler(mockMetadataHandler, mockRecordHandler);
}
use of com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse in project aws-athena-query-federation by awslabs.
the class GlueMetadataHandlerTest method setUp.
@Before
public void setUp() throws Exception {
logger.info("{}: enter", testName.getMethodName());
handler = new GlueMetadataHandler(mockGlue, new LocalKeyFactory(), mock(AWSSecretsManager.class), mock(AmazonAthena.class), "glue-test", "spill-bucket", "spill-prefix") {
@Override
public GetTableLayoutResponse doGetTableLayout(BlockAllocator blockAllocator, GetTableLayoutRequest request) {
throw new UnsupportedOperationException();
}
@Override
public void getPartitions(BlockWriter blockWriter, GetTableLayoutRequest request, QueryStatusChecker queryStatusChecker) throws Exception {
throw new UnsupportedOperationException();
}
@Override
public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsRequest request) {
throw new UnsupportedOperationException();
}
};
allocator = new BlockAllocatorImpl();
// doListTables pagination.
when(mockGlue.getTables(any(GetTablesRequest.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
GetTablesRequest request = (GetTablesRequest) invocationOnMock.getArguments()[0];
String nextToken = request.getNextToken();
int pageSize = request.getMaxResults() == null ? UNLIMITED_PAGE_SIZE_VALUE : request.getMaxResults();
assertEquals(accountId, request.getCatalogId());
assertEquals(schema, request.getDatabaseName());
GetTablesResult mockResult = mock(GetTablesResult.class);
if (pageSize == UNLIMITED_PAGE_SIZE_VALUE) {
// Simulate full list of tables returned from Glue.
when(mockResult.getTableList()).thenReturn(unPaginatedTables);
when(mockResult.getNextToken()).thenReturn(null);
} else {
// Simulate paginated list of tables returned from Glue.
List<Table> paginatedTables = unPaginatedTables.stream().sorted(Comparator.comparing(Table::getName)).filter(table -> nextToken == null || table.getName().compareTo(nextToken) >= 0).limit(pageSize + 1).collect(Collectors.toList());
if (paginatedTables.size() > pageSize) {
when(mockResult.getNextToken()).thenReturn(paginatedTables.get(pageSize).getName());
when(mockResult.getTableList()).thenReturn(paginatedTables.subList(0, pageSize));
} else {
when(mockResult.getNextToken()).thenReturn(null);
when(mockResult.getTableList()).thenReturn(paginatedTables);
}
}
return mockResult;
});
}
use of com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse in project aws-athena-query-federation by awslabs.
the class ExampleMetadataHandlerTest method doGetTableLayout.
/**
* 200,000,000 million partitions pruned down to 38,000 and transmitted in 25 seconds
*
* @throws Exception
*/
@Test
public void doGetTableLayout() throws Exception {
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(), 20)), false));
constraintsMap.put("month", SortedRangeSet.copyOf(Types.MinorType.INT.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.INT.getType(), 2)), false));
constraintsMap.put("year", SortedRangeSet.copyOf(Types.MinorType.INT.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.INT.getType(), 1900)), false));
GetTableLayoutRequest req = null;
GetTableLayoutResponse res = null;
try {
req = new GetTableLayoutRequest(IdentityUtil.fakeIdentity(), "queryId", "default", new TableName("schema1", "table1"), new Constraints(constraintsMap), tableSchema, partitionCols);
ObjectMapperUtil.assertSerialization(req);
res = metadataHandler.doGetTableLayout(allocator, req);
ObjectMapperUtil.assertSerialization(res);
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");
}
use of com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse in project aws-athena-query-federation by awslabs.
the class GetTableLayoutResponseSerDeTest method deserialize.
@Test
public void deserialize() throws IOException {
logger.info("deserialize: enter");
InputStream input = new ByteArrayInputStream(expectedSerDeText.getBytes());
GetTableLayoutResponse actual = (GetTableLayoutResponse) mapper.readValue(input, FederationResponse.class);
logger.info("deserialize: deserialized[{}]", actual);
assertEquals(expected, actual);
logger.info("deserialize: exit");
}
use of com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse in project aws-athena-query-federation by awslabs.
the class GetTableLayoutResponseSerDeTest method beforeTest.
@Before
public void beforeTest() throws IOException {
String yearCol = "year";
String monthCol = "month";
String dayCol = "day";
Schema schema = SchemaBuilder.newBuilder().addField(yearCol, new ArrowType.Int(32, true)).addField(monthCol, new ArrowType.Int(32, true)).addField(dayCol, new ArrowType.Int(32, true)).addField("col3", new ArrowType.Utf8()).build();
Block partitions = allocator.createBlock(schema);
int num_partitions = 10;
for (int i = 0; i < num_partitions; i++) {
BlockUtils.setValue(partitions.getFieldVector(yearCol), i, 2016 + i);
BlockUtils.setValue(partitions.getFieldVector(monthCol), i, (i % 12) + 1);
BlockUtils.setValue(partitions.getFieldVector(dayCol), i, (i % 28) + 1);
}
partitions.setRowCount(num_partitions);
expected = new GetTableLayoutResponse("test-catalog", new TableName("test-schema", "test-table"), partitions);
String expectedSerDeFile = utils.getResourceOrFail("serde/v2", "GetTableLayoutResponse.json");
expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim();
}
Aggregations