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

Example 26 with ListTablesRequest

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

Example 27 with ListTablesRequest

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

the class ExampleMetadataHandlerTest method doListTables.

@Test
public void doListTables() {
    if (!enableTests) {
        // We do this because until you complete the tutorial these tests will fail. When you attempt to publis
        // using ../toos/publish.sh ...  it will set the publishing flag and force these tests. This is how we
        // avoid breaking the build but still have a useful tutorial. We are also duplicateing this block
        // on purpose since this is a somewhat odd pattern.
        logger.info("doListTables: Tests are disabled, to enable them set the 'publishing' environment variable " + "using maven clean install -Dpublishing=true");
        return;
    }
    logger.info("doListTables - enter");
    // Test request with unlimited page size
    logger.info("doListTables - Test unlimited page size");
    ListTablesRequest req = new ListTablesRequest(fakeIdentity(), "queryId", "default", "schema1", null, UNLIMITED_PAGE_SIZE_VALUE);
    ListTablesResponse res = handler.doListTables(allocator, req);
    ListTablesResponse expectedResponse = new ListTablesResponse("default", new ImmutableList.Builder<TableName>().add(new TableName("schema1", "table1")).add(new TableName("schema1", "table2")).add(new TableName("schema1", "table3")).build(), null);
    logger.info("doListTables - {}", res);
    assertEquals("Expecting a different response", expectedResponse, res);
    // Test first paginated request with pageSize: 2, nextToken: null
    logger.info("doListTables - Test first pagination request");
    req = new ListTablesRequest(fakeIdentity(), "queryId", "default", "schema1", null, 2);
    expectedResponse = new ListTablesResponse("default", new ImmutableList.Builder<TableName>().add(new TableName("schema1", "table1")).add(new TableName("schema1", "table2")).build(), "table3");
    res = handler.doListTables(allocator, req);
    logger.info("doListTables - {}", res);
    assertEquals("Expecting a different response", expectedResponse, res);
    // Test second paginated request with pageSize: 2, nextToken: res.getNextToken()
    logger.info("doListTables - Test second pagination request");
    req = new ListTablesRequest(fakeIdentity(), "queryId", "default", "schema1", res.getNextToken(), 2);
    expectedResponse = new ListTablesResponse("default", new ImmutableList.Builder<TableName>().add(new TableName("schema1", "table3")).build(), null);
    res = handler.doListTables(allocator, req);
    logger.info("doListTables - {}", res);
    assertEquals("Expecting a different response", expectedResponse, res);
    logger.info("doListTables - exit");
}

Example 28 with ListTablesRequest

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

the class LambdaMetadataProvider method listTables.

/**
 * This method builds and executes a ListTablesRequest against the specified Lambda function.
 *
 * @param catalog the catalog name to be passed to Lambda
 * @param schema the name of the contextual schema for the request
 * @param metadataFunction the name of the Lambda function to call
 * @param identity the identity of the caller
 * @return the response
 */
public static ListTablesResponse listTables(String catalog, String schema, String metadataFunction, FederatedIdentity identity) {
    String queryId = generateQueryId();
    log.info("Submitting ListTablesRequest with ID " + queryId);
    /**
     * TODO: Add logic to ensure that the connector supports pagination.
     */
    try (ListTablesRequest request = new ListTablesRequest(identity, queryId, catalog, schema, null, UNLIMITED_PAGE_SIZE_VALUE)) {
        log.info("Submitting request: {}", request);
        ListTablesResponse response = (ListTablesResponse) getService(metadataFunction, identity, catalog).call(request);
        log.info("Received response: {}", response);
        return response;
    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}

Example 29 with ListTablesRequest

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

the class AwsCmdbMetadataHandlerTest method doListTables.

@Test
public void doListTables() {
    ListTablesRequest request = new ListTablesRequest(identity, queryId, catalog, "schema1", null, UNLIMITED_PAGE_SIZE_VALUE);
    ListTablesResponse response = handler.doListTables(blockAllocator, request);
    assertEquals(2, response.getTables().size());
    assertTrue(response.getTables().contains(new TableName("schema1", "table1")));
    assertTrue(response.getTables().contains(new TableName("schema1", "table2")));
}

Example 30 with ListTablesRequest

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

the class ImpalaMuxMetadataHandlerTest method doListTables.

@Test
public void doListTables() {
    ListTablesRequest listTablesRequest = Mockito.mock(ListTablesRequest.class);
    Mockito.when(listTablesRequest.getCatalogName()).thenReturn("metaImpala");
    this.jdbcMetadataHandler.doListTables(this.allocator, listTablesRequest);
    Mockito.verify(this.impalaMetadataHandler, Mockito.times(1)).doListTables(Mockito.eq(this.allocator), Mockito.eq(listTablesRequest));
}