Examples with ListDatabasesRequest com.amazonaws.services.timestreamwrite.model.ListDatabasesRequest used on opensource projects

Example 1 with ListDatabasesRequest

use of com.amazonaws.services.timestreamwrite.model.ListDatabasesRequest 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);
}

Example 2 with ListDatabasesRequest

use of com.amazonaws.services.timestreamwrite.model.ListDatabasesRequest 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");
}