Examples with ListTablesResponse com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse used on opensource projects

Example 11 with ListTablesResponse

use of com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse in project aws-athena-query-federation by awslabs.

the class ElasticsearchMetadataHandler method doListTables.

/**
 * Used to get the list of indices contained in the specified domain.
 * @param allocator Tool for creating and managing Apache Arrow Blocks.
 * @param request Provides details on who made the request and which Athena catalog and database they are querying.
 * @return A ListTablesResponse which primarily contains a List<TableName> enumerating the tables in this
 * catalog, database tuple. It also contains the catalog name corresponding the Athena catalog that was queried.
 * @throws RuntimeException when the domain does not exist in the map, or the client is unable to retrieve the
 * indices from the Elasticsearch instance.
 */
@Override
public ListTablesResponse doListTables(BlockAllocator allocator, ListTablesRequest request) throws RuntimeException {
    logger.debug("doListTables: enter - " + request);
    List<TableName> indices = new ArrayList<>();
    try {
        String endpoint = getDomainEndpoint(request.getSchemaName());
        AwsRestHighLevelClient client = clientFactory.getOrCreateClient(endpoint);
        try {
            for (String index : client.getAliases()) {
                // Ignore all system indices starting with period `.` (e.g. .kibana, .tasks, etc...)
                if (index.startsWith(".")) {
                    logger.info("Ignoring system index: {}", index);
                    continue;
                }
                indices.add(new TableName(request.getSchemaName(), index));
            }
        } catch (IOException error) {
            throw new RuntimeException("Error retrieving indices: " + error.getMessage(), error);
        }
    } catch (RuntimeException error) {
        throw new RuntimeException("Error processing request to list indices: " + error.getMessage(), error);
    }
    return new ListTablesResponse(request.getCatalogName(), indices, null);
}

Example 12 with ListTablesResponse

use of com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse in project aws-athena-query-federation by awslabs.

the class ConnectorValidator method showTables.

private static Collection<TableName> showTables(TestConfig testConfig, String db) {
    ListTablesResponse tablesResponse = LambdaMetadataProvider.listTables(testConfig.getCatalogId(), db, testConfig.getMetadataFunction(), testConfig.getIdentity());
    final Collection<TableName> tables = tablesResponse.getTables();
    log.info("Found tables: " + tables.stream().map(t -> toQualifiedTableName(t)).collect(Collectors.toList()));
    requireNonNull(tables, "Returned collection of tables was null!");
    if (!testConfig.isAllowEmptyTables()) {
        checkState(!tables.isEmpty(), "No tables were returned!");
        List<String> notLower = tables.stream().filter(t -> !t.equals(new TableName(t.getSchemaName().toLowerCase(), t.getTableName().toLowerCase()))).limit(5).map(t -> toQualifiedTableName(t)).collect(Collectors.toList());
        checkState(notLower.isEmpty(), "All returned tables must be lowercase! Found these non-lowercase tables: " + notLower);
    } else {
        log.info("showTables is allowing empty table list, actual size is " + tables.size());
    }
    return tables;
}

Example 13 with ListTablesResponse

use of com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse in project aws-athena-query-federation by awslabs.

the class ExampleMetadataHandler method doListTables.

/**
 * Returns a paginated list of TableNames.
 *
 * @param allocator Tool for creating and managing Apache Arrow Blocks.
 * @param request Provides details on who made the request and which Athena catalog and database they are querying.
 * @return A ListTablesResponse containing the list of available TableNames and a nextToken String for the start of
 * the next pagination (may be null if there are no more tables to paginate).
 * @implNote A complete (un-paginated) list of tables should be returned if the request's pageSize is set to
 * ListTablesRequest.UNLIMITED_PAGE_SIZE_VALUE.
 */
@Override
public ListTablesResponse doListTables(BlockAllocator allocator, ListTablesRequest request) {
    logCaller(request);
    String nextToken = null;
    int pageSize = request.getPageSize();
    // The following list is purposely unordered to demonstrate that the pagination logic does not consider the
    // order of the tables deterministic (i.e. pagination will work irrespective of the order that the tables are
    // returned from the source).
    List<TableName> tables = new ImmutableList.Builder<TableName>().add(new TableName("schema", "table4")).add(new TableName("schema", "table3")).add(new TableName("schema", "table5")).add(new TableName("schema", "table1")).add(new TableName("schema", "table2")).build();
    // Check if request has unlimited page size. If not, then list of tables will be paginated.
    if (pageSize != UNLIMITED_PAGE_SIZE_VALUE) {
        // Get the stating table for this page (if null, then this is the first page).
        String startToken = request.getNextToken();
        // Sort the list. Include all tables if the startToken is null, or only the tables whose names are equal to
        // or higher in value than startToken. Limit the number of tables in the list to the pageSize + 1 (the
        // nextToken).
        List<TableName> paginatedTables = tables.stream().sorted(Comparator.comparing(TableName::getTableName)).filter(table -> startToken == null || table.getTableName().compareTo(startToken) >= 0).limit(pageSize + 1).collect(Collectors.toList());
        if (paginatedTables.size() > pageSize) {
            // Paginated list contains full page of results + nextToken.
            // nextToken is the last element in the paginated list.
            // In an actual connector, the nextToken's value should be obfuscated.
            nextToken = paginatedTables.get(pageSize).getTableName();
            // nextToken is removed to include only the paginated results.
            tables = paginatedTables.subList(0, pageSize);
        } else {
            // Paginated list contains all remaining tables - end of the pagination.
            tables = paginatedTables;
        }
    }
    // that is exercised in our unit test suite.
    return new ListTablesResponse(request.getCatalogName(), tables.stream().filter(table -> request.getSchemaName() == null || request.getSchemaName().equals(table.getSchemaName())).collect(Collectors.toList()), nextToken);
}

Example 14 with ListTablesResponse

use of com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse in project aws-athena-query-federation by awslabs.

the class ExampleMetadataHandlerTest method doListTables.

@Test
public void doListTables() {
    logger.info("doListTables - enter");
    // Test request with unlimited page size
    logger.info("doListTables - Test unlimited page size");
    ListTablesRequest req = new ListTablesRequest(IdentityUtil.fakeIdentity(), "queryId", "default", "schema", null, UNLIMITED_PAGE_SIZE_VALUE);
    ListTablesResponse expectedResponse = new ListTablesResponse("default", new ImmutableList.Builder<TableName>().add(new TableName("schema", "table1")).add(new TableName("schema", "table2")).add(new TableName("schema", "table3")).add(new TableName("schema", "table4")).add(new TableName("schema", "table5")).build(), null);
    ObjectMapperUtil.assertSerialization(req);
    ListTablesResponse res = metadataHandler.doListTables(allocator, req);
    ObjectMapperUtil.assertSerialization(res);
    logger.info("doListTables - {}", res);
    assertEquals("Expecting a different response", expectedResponse, res);
    // Test first paginated request with pageSize: 3, nextToken: null
    logger.info("doListTables - Test first pagination request");
    req = new ListTablesRequest(IdentityUtil.fakeIdentity(), "queryId", "default", "schema", null, 3);
    expectedResponse = new ListTablesResponse("default", new ImmutableList.Builder<TableName>().add(new TableName("schema", "table1")).add(new TableName("schema", "table2")).add(new TableName("schema", "table3")).build(), "table4");
    ObjectMapperUtil.assertSerialization(req);
    res = metadataHandler.doListTables(allocator, req);
    ObjectMapperUtil.assertSerialization(res);
    logger.info("doListTables - {}", res);
    assertEquals("Expecting a different response", expectedResponse, res);
    // Test second paginated request with pageSize: 3, nextToken: res.getNextToken()
    logger.info("doListTables - Test second pagination request");
    req = new ListTablesRequest(IdentityUtil.fakeIdentity(), "queryId", "default", "schema", res.getNextToken(), 3);
    expectedResponse = new ListTablesResponse("default", new ImmutableList.Builder<TableName>().add(new TableName("schema", "table4")).add(new TableName("schema", "table5")).build(), null);
    ObjectMapperUtil.assertSerialization(req);
    res = metadataHandler.doListTables(allocator, req);
    ObjectMapperUtil.assertSerialization(res);
    logger.info("doListTables - {}", res);
    assertEquals("Expecting a different response", expectedResponse, res);
    logger.info("doListTables - exit");
}

Example 15 with ListTablesResponse

use of com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse in project aws-athena-query-federation by awslabs.

the class HbaseMetadataHandler method doListTables.

/**
 * List tables in the requested schema in your HBase instance treating the requested schema as an HBase
 * namespace.
 *
 * @see GlueMetadataHandler
 */
@Override
public ListTablesResponse doListTables(BlockAllocator blockAllocator, ListTablesRequest request) throws IOException {
    List<com.amazonaws.athena.connector.lambda.domain.TableName> tableNames = new ArrayList<>();
    TableName[] tables = getOrCreateConn(request).listTableNamesByNamespace(request.getSchemaName());
    for (int i = 0; i < tables.length; i++) {
        TableName tableName = tables[i];
        tableNames.add(new com.amazonaws.athena.connector.lambda.domain.TableName(request.getSchemaName(), tableName.getNameAsString().replace(request.getSchemaName() + ":", "")));
    }
    return new ListTablesResponse(request.getCatalogName(), tableNames, null);
}