use of com.amazonaws.services.timestreamwrite.model.Database in project aws-athena-query-federation by awslabs.
the class TimestreamMetadataHandler method doListSchemaNames.
@Override
public ListSchemasResponse doListSchemaNames(BlockAllocator blockAllocator, ListSchemasRequest request) throws Exception {
List<String> schemaNames = new ArrayList<>();
ListDatabasesRequest listDatabasesRequest = new ListDatabasesRequest();
ListDatabasesResult nextResult = tsMeta.listDatabases(listDatabasesRequest);
List<Database> nextDatabases = nextResult.getDatabases();
while (!nextDatabases.isEmpty()) {
nextDatabases.stream().forEach(next -> schemaNames.add(next.getDatabaseName()));
if (nextResult.getNextToken() != null && !nextResult.getNextToken().isEmpty()) {
listDatabasesRequest.setNextToken(nextResult.getNextToken());
nextResult = tsMeta.listDatabases(listDatabasesRequest);
nextDatabases = nextResult.getDatabases();
} else {
nextDatabases = Collections.EMPTY_LIST;
}
}
return new ListSchemasResponse(request.getCatalogName(), schemaNames);
}
use of com.amazonaws.services.timestreamwrite.model.Database in project aws-athena-query-federation by awslabs.
the class TimestreamMetadataHandlerTest method doListSchemaNames.
@Test
public void doListSchemaNames() throws Exception {
logger.info("doListSchemaNames - enter");
when(mockTsMeta.listDatabases(any(ListDatabasesRequest.class))).thenAnswer((InvocationOnMock invocation) -> {
ListDatabasesRequest request = invocation.getArgumentAt(0, ListDatabasesRequest.class);
String newNextToken = null;
List<Database> databases = new ArrayList<>();
if (request.getNextToken() == null) {
for (int i = 0; i < 10; i++) {
databases.add(new Database().withDatabaseName("database_" + i));
}
newNextToken = "1";
} else if (request.getNextToken().equals("1")) {
for (int i = 10; i < 100; i++) {
databases.add(new Database().withDatabaseName("database_" + i));
}
newNextToken = "2";
} else if (request.getNextToken().equals("2")) {
for (int i = 100; i < 1000; i++) {
databases.add(new Database().withDatabaseName("database_" + i));
}
newNextToken = null;
}
return new ListDatabasesResult().withDatabases(databases).withNextToken(newNextToken);
});
ListSchemasRequest req = new ListSchemasRequest(identity, "queryId", "default");
ListSchemasResponse res = handler.doListSchemaNames(allocator, req);
logger.info("doListSchemaNames - {}", res.getSchemas());
assertEquals(1000, res.getSchemas().size());
verify(mockTsMeta, times(3)).listDatabases(any(ListDatabasesRequest.class));
Iterator<String> schemaItr = res.getSchemas().iterator();
for (int i = 0; i < 1000; i++) {
assertEquals("database_" + i, schemaItr.next());
}
logger.info("doListSchemaNames - exit");
}
Aggregations