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

Example 16 with ListSchemasResponse

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

the class TimestreamMetadataHandler method doListSchemaNames.

@Override
public ListSchemasResponse doListSchemaNames(BlockAllocator blockAllocator, ListSchemasRequest request) throws Exception {
    List<String> schemaNames = new ArrayList<>();
    ListDatabasesRequest listDatabasesRequest = new ListDatabasesRequest();
    ListDatabasesResult nextResult = tsMeta.listDatabases(listDatabasesRequest);
    List<Database> nextDatabases = nextResult.getDatabases();
    while (!nextDatabases.isEmpty()) {
        nextDatabases.stream().forEach(next -> schemaNames.add(next.getDatabaseName()));
        if (nextResult.getNextToken() != null && !nextResult.getNextToken().isEmpty()) {
            listDatabasesRequest.setNextToken(nextResult.getNextToken());
            nextResult = tsMeta.listDatabases(listDatabasesRequest);
            nextDatabases = nextResult.getDatabases();
        } else {
            nextDatabases = Collections.EMPTY_LIST;
        }
    }
    return new ListSchemasResponse(request.getCatalogName(), schemaNames);
}

Example 17 with ListSchemasResponse

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

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

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

the class ConnectorValidator method showDatabases.

private static Collection<String> showDatabases(TestConfig testConfig) {
    ListSchemasResponse schemasResponse = LambdaMetadataProvider.listSchemas(testConfig.getCatalogId(), testConfig.getMetadataFunction(), testConfig.getIdentity());
    final Collection<String> schemas = schemasResponse.getSchemas();
    log.info("Found databases: " + schemas);
    requireNonNull(schemas, "Returned collection of schemas was null!");
    checkState(!schemas.isEmpty(), "No schemas were returned!");
    List<String> notLower = schemas.stream().filter(s -> !s.equals(s.toLowerCase())).collect(Collectors.toList());
    checkState(notLower.isEmpty(), "All returned schemas must be lowercase! Found these non-lowercase schemas: " + notLower);
    return schemas;
}

Example 20 with ListSchemasResponse

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