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

Example 26 with ListSchemasRequest

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

the class TimestreamMetadataHandlerTest method doListSchemaNames.

@Test
public void doListSchemaNames() throws Exception {
    logger.info("doListSchemaNames - enter");
    when(mockTsMeta.listDatabases(any(ListDatabasesRequest.class))).thenAnswer((InvocationOnMock invocation) -> {
        ListDatabasesRequest request = invocation.getArgumentAt(0, ListDatabasesRequest.class);
        String newNextToken = null;
        List<Database> databases = new ArrayList<>();
        if (request.getNextToken() == null) {
            for (int i = 0; i < 10; i++) {
                databases.add(new Database().withDatabaseName("database_" + i));
            }
            newNextToken = "1";
        } else if (request.getNextToken().equals("1")) {
            for (int i = 10; i < 100; i++) {
                databases.add(new Database().withDatabaseName("database_" + i));
            }
            newNextToken = "2";
        } else if (request.getNextToken().equals("2")) {
            for (int i = 100; i < 1000; i++) {
                databases.add(new Database().withDatabaseName("database_" + i));
            }
            newNextToken = null;
        }
        return new ListDatabasesResult().withDatabases(databases).withNextToken(newNextToken);
    });
    ListSchemasRequest req = new ListSchemasRequest(identity, "queryId", "default");
    ListSchemasResponse res = handler.doListSchemaNames(allocator, req);
    logger.info("doListSchemaNames - {}", res.getSchemas());
    assertEquals(1000, res.getSchemas().size());
    verify(mockTsMeta, times(3)).listDatabases(any(ListDatabasesRequest.class));
    Iterator<String> schemaItr = res.getSchemas().iterator();
    for (int i = 0; i < 1000; i++) {
        assertEquals("database_" + i, schemaItr.next());
    }
    logger.info("doListSchemaNames - exit");
}

Example 27 with ListSchemasRequest

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

the class ExampleMetadataHandlerTest method doListSchemaNames.

@Test
public void doListSchemaNames() {
    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("doListSchemaNames: Tests are disabled, to enable them set the 'publishing' environment variable " + "using maven clean install -Dpublishing=true");
        return;
    }
    logger.info("doListSchemas - enter");
    ListSchemasRequest req = new ListSchemasRequest(fakeIdentity(), "queryId", "default");
    ListSchemasResponse res = handler.doListSchemaNames(allocator, req);
    logger.info("doListSchemas - {}", res.getSchemas());
    assertFalse(res.getSchemas().isEmpty());
    logger.info("doListSchemas - exit");
}

Example 28 with ListSchemasRequest

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

the class LambdaMetadataProvider method listSchemas.

/**
 * This method builds and executes a ListSchemasRequest against the specified Lambda function.
 *
 * @param catalog the catalog name to be passed to Lambda
 * @param metadataFunction the name of the Lambda function to call
 * @param identity the identity of the caller
 * @return the response
 */
public static ListSchemasResponse listSchemas(String catalog, String metadataFunction, FederatedIdentity identity) {
    String queryId = generateQueryId();
    log.info("Submitting ListSchemasRequest with ID " + queryId);
    try (ListSchemasRequest request = new ListSchemasRequest(identity, queryId, catalog)) {
        log.info("Submitting request: {}", request);
        ListSchemasResponse response = (ListSchemasResponse) getService(metadataFunction, identity, catalog).call(request);
        log.info("Received response: {}", response);
        return response;
    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}

Example 29 with ListSchemasRequest

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

the class AwsCmdbMetadataHandlerTest method doListSchemaNames.

@Test
public void doListSchemaNames() {
    ListSchemasRequest request = new ListSchemasRequest(identity, queryId, catalog);
    ListSchemasResponse response = handler.doListSchemaNames(blockAllocator, request);
    assertEquals(2, response.getSchemas().size());
    assertTrue(response.getSchemas().contains("schema1"));
    assertTrue(response.getSchemas().contains("schema2"));
}

Example 30 with ListSchemasRequest

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

the class ImpalaMuxMetadataHandlerTest method doListSchemaNames.

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