use of com.amazonaws.athena.connector.lambda.data.BlockWriter 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.data.BlockWriter 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;
}
}
}
use of com.amazonaws.athena.connector.lambda.data.BlockWriter in project aws-athena-query-federation by awslabs.
the class JdbcMetadataHandlerTest method setup.
@Before
public void setup() {
this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
this.connection = Mockito.mock(Connection.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(this.jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(this.connection);
this.secretsManager = Mockito.mock(AWSSecretsManager.class);
this.athena = Mockito.mock(AmazonAthena.class);
Mockito.when(this.secretsManager.getSecretValue(Mockito.eq(new GetSecretValueRequest().withSecretId("testSecret")))).thenReturn(new GetSecretValueResult().withSecretString("{\"username\": \"testUser\", \"password\": \"testPassword\"}"));
DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", "fakedatabase", "fakedatabase://jdbc:fakedatabase://hostname/${testSecret}", "testSecret");
this.jdbcMetadataHandler = new JdbcMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory) {
@Override
public Schema getPartitionSchema(final String catalogName) {
return PARTITION_SCHEMA;
}
@Override
public void getPartitions(final BlockWriter blockWriter, final GetTableLayoutRequest getTableLayoutRequest, QueryStatusChecker queryStatusChecker) {
}
@Override
public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsRequest getSplitsRequest) {
return null;
}
};
this.federatedIdentity = Mockito.mock(FederatedIdentity.class);
this.blockAllocator = Mockito.mock(BlockAllocator.class);
}
Aggregations