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

Example 26 with ListTablesResponse

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

the class CloudwatchMetadataHandler method doListTables.

/**
 * List LogStreams within the requested schema (aka LogGroup) in your Cloudwatch account treating each as a 'table'.
 *
 * @see MetadataHandler
 */
@Override
public ListTablesResponse doListTables(BlockAllocator blockAllocator, ListTablesRequest listTablesRequest) throws TimeoutException {
    String nextToken = null;
    String logGroupName = tableResolver.validateSchema(listTablesRequest.getSchemaName());
    DescribeLogStreamsRequest request = new DescribeLogStreamsRequest(logGroupName);
    DescribeLogStreamsResult result;
    List<TableName> tables = new ArrayList<>();
    if (listTablesRequest.getPageSize() == UNLIMITED_PAGE_SIZE_VALUE) {
        do {
            if (tables.size() > MAX_RESULTS) {
                throw new RuntimeException("Too many log streams, exceeded max metadata results for table count.");
            }
            result = invoker.invoke(() -> awsLogs.describeLogStreams(request));
            result.getLogStreams().forEach(next -> tables.add(toTableName(listTablesRequest, next)));
            request.setNextToken(result.getNextToken());
            logger.info("doListTables: Listing log streams  with token {} and size {}", result.getNextToken(), tables.size());
        } while (result.getNextToken() != null);
    } else {
        request.setNextToken(listTablesRequest.getNextToken());
        request.setLimit(listTablesRequest.getPageSize());
        result = invoker.invoke(() -> awsLogs.describeLogStreams(request));
        result.getLogStreams().forEach(next -> tables.add(toTableName(listTablesRequest, next)));
        nextToken = result.getNextToken();
        logger.info("doListTables: Listing log streams with token {} and size {}", result.getNextToken(), tables.size());
    }
    // We add a special table that represents all log streams. This is helpful depending on how
    // you have your logs organized.
    tables.add(new TableName(listTablesRequest.getSchemaName(), ALL_LOG_STREAMS_TABLE));
    return new ListTablesResponse(listTablesRequest.getCatalogName(), tables, nextToken);
}

Example 27 with ListTablesResponse

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

the class CloudwatchMetadataHandlerTest method doListTables.

@Test
public void doListTables() throws TimeoutException {
    logger.info("doListTables - enter");
    when(mockAwsLogs.describeLogStreams(any(DescribeLogStreamsRequest.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
        DescribeLogStreamsRequest request = (DescribeLogStreamsRequest) invocationOnMock.getArguments()[0];
        DescribeLogStreamsResult result = new DescribeLogStreamsResult();
        Integer nextToken;
        if (request.getNextToken() == null) {
            nextToken = 1;
        } else if (Integer.valueOf(request.getNextToken()) < 3) {
            nextToken = Integer.valueOf(request.getNextToken()) + 1;
        } else {
            nextToken = null;
        }
        List<LogStream> logStreams = new ArrayList<>();
        if (request.getNextToken() == null || Integer.valueOf(request.getNextToken()) < 3) {
            for (int i = 0; i < 10; i++) {
                LogStream nextLogStream = new LogStream();
                nextLogStream.setLogStreamName("table-" + String.valueOf(i));
                logStreams.add(nextLogStream);
            }
        }
        result.withLogStreams(logStreams);
        if (nextToken != null) {
            result.setNextToken(String.valueOf(nextToken));
        }
        return result;
    });
    ListTablesRequest req = new ListTablesRequest(identity, "queryId", "default", "schema-1", null, UNLIMITED_PAGE_SIZE_VALUE);
    ListTablesResponse res = handler.doListTables(allocator, req);
    logger.info("doListTables - {}", res.getTables());
    assertTrue(res.getTables().contains(new TableName("schema-1", "all_log_streams")));
    assertTrue(res.getTables().size() == 31);
    verify(mockAwsLogs, times(4)).describeLogStreams(any(DescribeLogStreamsRequest.class));
    verify(mockAwsLogs, times(1)).describeLogGroups(any(DescribeLogGroupsRequest.class));
    verifyNoMoreInteractions(mockAwsLogs);
    logger.info("doListTables - exit");
}

Example 28 with ListTablesResponse

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

the class DocDBMetadataHandler method doListTables.

/**
 * List collections in the requested schema in your DocumentDB instance treating the requested schema as an DocumentDB
 * database.
 *
 * @see GlueMetadataHandler
 */
@Override
public ListTablesResponse doListTables(BlockAllocator blockAllocator, ListTablesRequest request) {
    MongoClient client = getOrCreateConn(request);
    List<TableName> tables = new ArrayList<>();
    try (MongoCursor<String> itr = client.getDatabase(request.getSchemaName()).listCollectionNames().iterator()) {
        while (itr.hasNext()) {
            tables.add(new TableName(request.getSchemaName(), itr.next()));
        }
        return new ListTablesResponse(request.getCatalogName(), tables, null);
    }
}

Example 29 with ListTablesResponse

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

the class DocDBMetadataHandlerTest method doListTables.

@Test
public void doListTables() {
    List<String> tableNames = new ArrayList<>();
    tableNames.add("table1");
    tableNames.add("table2");
    tableNames.add("table3");
    MongoDatabase mockDatabase = mock(MongoDatabase.class);
    when(mockClient.getDatabase(eq(DEFAULT_SCHEMA))).thenReturn(mockDatabase);
    when(mockDatabase.listCollectionNames()).thenReturn(StubbingCursor.iterate(tableNames));
    ListTablesRequest req = new ListTablesRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, DEFAULT_SCHEMA, null, UNLIMITED_PAGE_SIZE_VALUE);
    ListTablesResponse res = handler.doListTables(allocator, req);
    logger.info("doListTables - {}", res.getTables());
    for (TableName next : res.getTables()) {
        assertEquals(DEFAULT_SCHEMA, next.getSchemaName());
        assertTrue(tableNames.contains(next.getTableName()));
    }
    assertEquals(tableNames.size(), res.getTables().size());
}

Example 30 with ListTablesResponse

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

the class GlueMetadataHandler method doListTables.

/**
 * Returns a paginated list of tables from AWS Glue DataCatalog with optional filtering for the requested schema
 * (aka database).
 *
 * @param blockAllocator Tool for creating and managing Apache Arrow Blocks.
 * @param request Provides details on who made the request and which Athena catalog they are querying.
 * @param filter The TableFilter to apply to all tables before adding them to the results list.
 * @return The ListTablesResponse which mostly contains the list of table names.
 * @implNote A complete (un-paginated) list of tables should be returned if the request's pageSize is set to
 * ListTablesRequest.UNLIMITED_PAGE_SIZE_VALUE.
 */
protected ListTablesResponse doListTables(BlockAllocator blockAllocator, ListTablesRequest request, TableFilter filter) throws Exception {
    GetTablesRequest getTablesRequest = new GetTablesRequest();
    getTablesRequest.setCatalogId(getCatalog(request));
    getTablesRequest.setDatabaseName(request.getSchemaName());
    Set<TableName> tables = new HashSet<>();
    String nextToken = request.getNextToken();
    int pageSize = request.getPageSize();
    do {
        getTablesRequest.setNextToken(nextToken);
        if (pageSize != UNLIMITED_PAGE_SIZE_VALUE) {
            // Paginated requests will include the maxResults argument determined by the minimum value between the
            // pageSize and the maximum results supported by Glue (as defined in the Glue API docs).
            int maxResults = Math.min(pageSize, GET_TABLES_REQUEST_MAX_RESULTS);
            getTablesRequest.setMaxResults(maxResults);
            pageSize -= maxResults;
        }
        GetTablesResult result = awsGlue.getTables(getTablesRequest);
        for (Table next : result.getTableList()) {
            if (filter == null || filter.filter(next)) {
                tables.add(new TableName(request.getSchemaName(), next.getName()));
            }
        }
        nextToken = result.getNextToken();
    } while (nextToken != null && (pageSize == UNLIMITED_PAGE_SIZE_VALUE || pageSize > 0));
    return new ListTablesResponse(request.getCatalogName(), tables, nextToken);
}