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

Example 1 with ListTablesResult

use of com.amazonaws.services.timestreamwrite.model.ListTablesResult in project aws-athena-query-federation by awslabs.

the class TimestreamMetadataHandler method doListTables.

@Override
public ListTablesResponse doListTables(BlockAllocator blockAllocator, ListTablesRequest request) throws Exception {
    List<TableName> tableNames = new ArrayList<>();
    com.amazonaws.services.timestreamwrite.model.ListTablesRequest listTablesRequest = new com.amazonaws.services.timestreamwrite.model.ListTablesRequest().withDatabaseName(request.getSchemaName());
    ListTablesResult nextResult = tsMeta.listTables(listTablesRequest);
    List<com.amazonaws.services.timestreamwrite.model.Table> nextTables = nextResult.getTables();
    while (!nextTables.isEmpty()) {
        nextTables.stream().forEach(next -> tableNames.add(new TableName(request.getSchemaName(), next.getTableName())));
        if (nextResult.getNextToken() != null && !nextResult.getNextToken().isEmpty()) {
            listTablesRequest.setNextToken(nextResult.getNextToken());
            nextResult = tsMeta.listTables(listTablesRequest);
            nextTables = nextResult.getTables();
        } else {
            nextTables = Collections.EMPTY_LIST;
        }
    }
    return new ListTablesResponse(request.getCatalogName(), tableNames, null);
}

Example 2 with ListTablesResult

use of com.amazonaws.services.timestreamwrite.model.ListTablesResult in project aws-athena-query-federation by awslabs.

the class TimestreamMetadataHandlerTest method doListTables.

@Test
public void doListTables() throws Exception {
    logger.info("doListTables - enter");
    when(mockTsMeta.listTables(any(com.amazonaws.services.timestreamwrite.model.ListTablesRequest.class))).thenAnswer((InvocationOnMock invocation) -> {
        com.amazonaws.services.timestreamwrite.model.ListTablesRequest request = invocation.getArgumentAt(0, com.amazonaws.services.timestreamwrite.model.ListTablesRequest.class);
        String newNextToken = null;
        List<Table> tables = new ArrayList<>();
        if (request.getNextToken() == null) {
            for (int i = 0; i < 10; i++) {
                tables.add(new Table().withDatabaseName(request.getDatabaseName()).withTableName("table_" + i));
            }
            newNextToken = "1";
        } else if (request.getNextToken().equals("1")) {
            for (int i = 10; i < 100; i++) {
                tables.add(new Table().withDatabaseName(request.getDatabaseName()).withTableName("table_" + i));
            }
            newNextToken = "2";
        } else if (request.getNextToken().equals("2")) {
            for (int i = 100; i < 1000; i++) {
                tables.add(new Table().withDatabaseName(request.getDatabaseName()).withTableName("table_" + i));
            }
            newNextToken = null;
        }
        return new ListTablesResult().withTables(tables).withNextToken(newNextToken);
    });
    ListTablesRequest req = new ListTablesRequest(identity, "queryId", "default", defaultSchema, null, UNLIMITED_PAGE_SIZE_VALUE);
    ListTablesResponse res = handler.doListTables(allocator, req);
    logger.info("doListTables - {}", res.getTables());
    assertEquals(1000, res.getTables().size());
    verify(mockTsMeta, times(3)).listTables(any(com.amazonaws.services.timestreamwrite.model.ListTablesRequest.class));
    Iterator<TableName> schemaItr = res.getTables().iterator();
    for (int i = 0; i < 1000; i++) {
        TableName tableName = schemaItr.next();
        assertEquals(defaultSchema, tableName.getSchemaName());
        assertEquals("table_" + i, tableName.getTableName());
    }
    logger.info("doListTables - exit");
}