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

Example 31 with ListTablesRequest

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

the class MetricsMetadataHandlerTest method doListTables.

@Test
public void doListTables() {
    logger.info("doListTables - enter");
    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(2, res.getTables().size());
    assertTrue(res.getTables().contains(new TableName(defaultSchema, "metrics")));
    assertTrue(res.getTables().contains(new TableName(defaultSchema, "metric_samples")));
    logger.info("doListTables - exit");
}

Example 32 with ListTablesRequest

use of com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest 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 33 with ListTablesRequest

use of com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest 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());
}