com.amazonaws.services.athena.AmazonAthena Java Examples

/**
 * Submits an example query and returns a query execution ID of a running query to stop.
 */
public static String submitAthenaQuery(AmazonAthena athenaClient)
{
    QueryExecutionContext queryExecutionContext = new QueryExecutionContext().withDatabase(ExampleConstants.ATHENA_DEFAULT_DATABASE);

    ResultConfiguration resultConfiguration = new ResultConfiguration()
            .withOutputLocation(ExampleConstants.ATHENA_OUTPUT_BUCKET);

    StartQueryExecutionRequest startQueryExecutionRequest = new StartQueryExecutionRequest()
            .withQueryString(ExampleConstants.ATHENA_SAMPLE_QUERY)
            .withQueryExecutionContext(queryExecutionContext)
            .withResultConfiguration(resultConfiguration);

    StartQueryExecutionResult startQueryExecutionResult = athenaClient.startQueryExecution(startQueryExecutionRequest);

    return startQueryExecutionResult.getQueryExecutionId();
}
 

@Before
public void setup()
{
    //this.allocator = Mockito.mock(BlockAllocator.class);
    this.allocator = new BlockAllocatorImpl();
    //Mockito.when(this.allocator.createBlock(Mockito.any(Schema.class))).thenReturn(Mockito.mock(Block.class));
    this.mySqlMetadataHandler = Mockito.mock(MySqlMetadataHandler.class);
    this.metadataHandlerMap = Collections.singletonMap("mysql", this.mySqlMetadataHandler);
    this.secretsManager = Mockito.mock(AWSSecretsManager.class);
    this.athena = Mockito.mock(AmazonAthena.class);
    this.queryStatusChecker = Mockito.mock(QueryStatusChecker.class);
    this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
    DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", JdbcConnectionFactory.DatabaseEngine.MYSQL,
            "mysql://jdbc:mysql://hostname/${testSecret}", "testSecret");
    this.jdbcMetadataHandler = new MultiplexingJdbcMetadataHandler(this.secretsManager, this.athena, this.jdbcConnectionFactory, this.metadataHandlerMap, databaseConnectionConfig);
}
 

@VisibleForTesting
protected ElasticsearchMetadataHandler(AWSGlue awsGlue,
                                       EncryptionKeyFactory keyFactory,
                                       AWSSecretsManager awsSecretsManager,
                                       AmazonAthena athena,
                                       String spillBucket,
                                       String spillPrefix,
                                       ElasticsearchDomainMapProvider domainMapProvider,
                                       AwsRestHighLevelClientFactory clientFactory,
                                       long queryTimeout)
{
    super(awsGlue, keyFactory, awsSecretsManager, athena, SOURCE_TYPE, spillBucket, spillPrefix);
    this.awsGlue = awsGlue;
    this.domainMapProvider = domainMapProvider;
    this.domainMap = this.domainMapProvider.getDomainMap(null);
    this.clientFactory = clientFactory;
    this.glueTypeMapper = new ElasticsearchGlueTypeMapper();
    this.queryTimeout = queryTimeout;
}
 

/**
 * @param sourceType Used to aid in logging diagnostic info when raising a support case.
 */
public MetadataHandler(EncryptionKeyFactory encryptionKeyFactory,
        AWSSecretsManager secretsManager,
        AmazonAthena athena,
        String sourceType,
        String spillBucket,
        String spillPrefix)
{
    this.encryptionKeyFactory = encryptionKeyFactory;
    this.secretsManager = new CachableSecretsManager(secretsManager);
    this.athena = athena;
    this.sourceType = sourceType;
    this.spillBucket = spillBucket;
    this.spillPrefix = spillPrefix;
    this.verifier = new SpillLocationVerifier(AmazonS3ClientBuilder.standard().build());
}
 

public static void main(String[] args) throws Exception
{
    // Build an Athena client
    AthenaClientFactory factory = new AthenaClientFactory();
    AmazonAthena athenaClient = factory.createClient();

    // Create the named query request.
    CreateNamedQueryRequest createNamedQueryRequest = new CreateNamedQueryRequest()
            .withDatabase(ExampleConstants.ATHENA_DEFAULT_DATABASE)
            .withQueryString(ExampleConstants.ATHENA_SAMPLE_QUERY)
            .withDescription("Sample Description")
            .withName("SampleQuery2");

    // Call Athena to create the named query. If it fails, an exception is thrown.
    CreateNamedQueryResult createNamedQueryResult = athenaClient.createNamedQuery(createNamedQueryRequest);
}
 

public static void main(String[] args) throws Exception
{
    // Build an Athena client
    AthenaClientFactory factory = new AthenaClientFactory();
    AmazonAthena athenaClient = factory.createClient();

    String sampleQueryExecutionId = submitAthenaQuery(athenaClient);

    // Submit the stop query Request
    StopQueryExecutionRequest stopQueryExecutionRequest = new StopQueryExecutionRequest()
            .withQueryExecutionId(sampleQueryExecutionId);

    StopQueryExecutionResult stopQueryExecutionResult = athenaClient.stopQueryExecution(stopQueryExecutionRequest);

    // Ensure that the query was stopped
    GetQueryExecutionRequest getQueryExecutionRequest = new GetQueryExecutionRequest()
            .withQueryExecutionId(sampleQueryExecutionId);

    GetQueryExecutionResult getQueryExecutionResult = athenaClient.getQueryExecution(getQueryExecutionRequest);
    if (getQueryExecutionResult.getQueryExecution().getStatus().getState().equals(QueryExecutionState.CANCELLED)) {
        // Query was cancelled.
        System.out.println("Query has been cancelled");
    }
}
 

/**
 * Submits a sample query to Athena and returns the execution ID of the query.
 */
private static String submitAthenaQuery(AmazonAthena athenaClient)
{
    // The QueryExecutionContext allows us to set the Database.
    QueryExecutionContext queryExecutionContext = new QueryExecutionContext().withDatabase(ExampleConstants.ATHENA_DEFAULT_DATABASE);

    // The result configuration specifies where the results of the query should go in S3 and encryption options
    ResultConfiguration resultConfiguration = new ResultConfiguration()
            // You can provide encryption options for the output that is written.
            // .withEncryptionConfiguration(encryptionConfiguration)
            .withOutputLocation(ExampleConstants.ATHENA_OUTPUT_BUCKET);

    // Create the StartQueryExecutionRequest to send to Athena which will start the query.
    StartQueryExecutionRequest startQueryExecutionRequest = new StartQueryExecutionRequest()
            .withQueryString(ExampleConstants.ATHENA_SAMPLE_QUERY)
            .withQueryExecutionContext(queryExecutionContext)
            .withResultConfiguration(resultConfiguration);

    StartQueryExecutionResult startQueryExecutionResult = athenaClient.startQueryExecution(startQueryExecutionRequest);
    return startQueryExecutionResult.getQueryExecutionId();
}
 

@VisibleForTesting
protected RedisRecordHandler(AmazonS3 amazonS3,
        AWSSecretsManager secretsManager,
        AmazonAthena athena,
        JedisPoolFactory jedisPoolFactory)
{
    super(amazonS3, secretsManager, athena, SOURCE_TYPE);
    this.amazonS3 = amazonS3;
    this.jedisPoolFactory = jedisPoolFactory;
}
 

@Before
public void setup()
{
    this.amazonS3 = Mockito.mock(AmazonS3.class);
    this.secretsManager = Mockito.mock(AWSSecretsManager.class);
    this.athena = Mockito.mock(AmazonAthena.class);
    this.connection = Mockito.mock(Connection.class);
    this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
    Mockito.when(this.jdbcConnectionFactory.getConnection(Mockito.mock(JdbcCredentialProvider.class))).thenReturn(this.connection);
    jdbcSplitQueryBuilder = new MySqlQueryStringBuilder("`");
    final DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", JdbcConnectionFactory.DatabaseEngine.MYSQL,
            "mysql://jdbc:mysql://hostname/user=A&password=B");

    this.mySqlRecordHandler = new MySqlRecordHandler(databaseConnectionConfig, amazonS3, secretsManager, athena, jdbcConnectionFactory, jdbcSplitQueryBuilder);
}
 

@Before
public void setup()
{
    this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
    this.connection = Mockito.mock(Connection.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(this.jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(this.connection);
    this.secretsManager = Mockito.mock(AWSSecretsManager.class);
    this.athena = Mockito.mock(AmazonAthena.class);
    Mockito.when(this.secretsManager.getSecretValue(Mockito.eq(new GetSecretValueRequest().withSecretId("testSecret")))).thenReturn(new GetSecretValueResult().withSecretString("{\"username\": \"testUser\", \"password\": \"testPassword\"}"));
    this.mySqlMetadataHandler = new MySqlMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, this.jdbcConnectionFactory);
    this.federatedIdentity = Mockito.mock(FederatedIdentity.class);
}
 

@VisibleForTesting
protected MetricsMetadataHandler(AmazonCloudWatch metrics,
        EncryptionKeyFactory keyFactory,
        AWSSecretsManager secretsManager,
        AmazonAthena athena,
        String spillBucket,
        String spillPrefix)
{
    super(keyFactory, secretsManager, athena, SOURCE_TYPE, spillBucket, spillPrefix);
    this.metrics = metrics;
}
 

@VisibleForTesting
protected CloudwatchMetadataHandler(AWSLogs awsLogs,
        EncryptionKeyFactory keyFactory,
        AWSSecretsManager secretsManager,
        AmazonAthena athena,
        String spillBucket,
        String spillPrefix)
{
    super(keyFactory, secretsManager, athena, SOURCE_TYPE, spillBucket, spillPrefix);
    this.awsLogs = awsLogs;
    tableResolver = new CloudwatchTableResolver(invoker, awsLogs, MAX_RESULTS, MAX_RESULTS);
}
 

@VisibleForTesting
DynamoDBMetadataHandler(EncryptionKeyFactory keyFactory,
        AWSSecretsManager secretsManager,
        AmazonAthena athena,
        String spillBucket,
        String spillPrefix,
        AmazonDynamoDB ddbClient,
        AWSGlue glueClient)
{
    super(glueClient, keyFactory, secretsManager, athena, SOURCE_TYPE, spillBucket, spillPrefix);
    this.glueClient = glueClient;
    this.ddbClient = ddbClient;
    this.tableResolver = new DynamoDBTableResolver(invoker, ddbClient);
}
 

@Before
public void setup()
{
    logger.info("{}: enter", testName.getMethodName());

    allocator = new BlockAllocatorImpl();
    handler = new DynamoDBRecordHandler(ddbClient, mock(AmazonS3.class), mock(AWSSecretsManager.class), mock(AmazonAthena.class), "source_type");
    metadataHandler = new DynamoDBMetadataHandler(new LocalKeyFactory(), secretsManager, athena, "spillBucket", "spillPrefix", ddbClient, glueClient);
}
 

@VisibleForTesting
protected DocDBMetadataHandler(AWSGlue glue,
        DocDBConnectionFactory connectionFactory,
        EncryptionKeyFactory keyFactory,
        AWSSecretsManager secretsManager,
        AmazonAthena athena,
        String spillBucket,
        String spillPrefix)
{
    super(glue, keyFactory, secretsManager, athena, SOURCE_TYPE, spillBucket, spillPrefix);
    this.glue = glue;
    this.connectionFactory = connectionFactory;
}
 

@VisibleForTesting
protected ElasticsearchRecordHandler(AmazonS3 amazonS3, AWSSecretsManager secretsManager, AmazonAthena amazonAthena,
                                     AwsRestHighLevelClientFactory clientFactory, long queryTimeout)
{
    super(amazonS3, secretsManager, amazonAthena, SOURCE_TYPE);

    this.typeUtils = new ElasticsearchTypeUtils();
    this.clientFactory = clientFactory;
    this.queryTimeout = queryTimeout;
}
 

@VisibleForTesting
protected AwsCmdbMetadataHandler(TableProviderFactory tableProviderFactory,
        EncryptionKeyFactory keyFactory,
        AWSSecretsManager secretsManager,
        AmazonAthena athena,
        String spillBucket,
        String spillPrefix)
{
    super(keyFactory, secretsManager, athena, SOURCE_TYPE, spillBucket, spillPrefix);
    schemas = tableProviderFactory.getSchemas();
    tableProviders = tableProviderFactory.getTableProviders();
}
 

@VisibleForTesting
protected RedisMetadataHandler(AWSGlue awsGlue,
        EncryptionKeyFactory keyFactory,
        AWSSecretsManager secretsManager,
        AmazonAthena athena,
        JedisPoolFactory jedisPoolFactory,
        String spillBucket,
        String spillPrefix)
{
    super(awsGlue, keyFactory, secretsManager, athena, SOURCE_TYPE, spillBucket, spillPrefix);
    this.awsGlue = awsGlue;
    this.jedisPoolFactory = jedisPoolFactory;
}
 

@VisibleForTesting
protected HbaseMetadataHandler(AWSGlue awsGlue,
        EncryptionKeyFactory keyFactory,
        AWSSecretsManager secretsManager,
        AmazonAthena athena,
        HbaseConnectionFactory connectionFactory,
        String spillBucket,
        String spillPrefix)
{
    super(awsGlue, keyFactory, secretsManager, athena, SOURCE_TYPE, spillBucket, spillPrefix);
    this.awsGlue = awsGlue;
    this.connectionFactory = connectionFactory;
}
 

@VisibleForTesting
protected HbaseRecordHandler(AmazonS3 amazonS3, AWSSecretsManager secretsManager, AmazonAthena athena, HbaseConnectionFactory connectionFactory)
{
    super(amazonS3, secretsManager, athena, SOURCE_TYPE);
    this.amazonS3 = amazonS3;
    this.connectionFactory = connectionFactory;
}
 

public void addPartitionIfMissing(
    String region, String bucket, String year, String month, String day, String hour) {
  String partition = String.format("year=%s/month=%s/day=%s/hour=%s", year, month, day, hour);
  String table = String.format(TABLE_TEMPLATE, environmentName);
  if (!partitions.contains(partition)) {
    try {
      String query =
          String.format(
              "ALTER TABLE %s ADD PARTITION (year='%s', month='%s', day='%s', hour='%s') "
                  + "LOCATION 's3://%s/audit-logs/partitioned/year=%s/month=%s/day=%s/hour=%s'",
              table, year, month, day, hour, bucket, year, month, day, hour);

      AmazonAthena athena = athenaClientFactory.getClient(region);

      StartQueryExecutionResult result =
          athena.startQueryExecution(
              new StartQueryExecutionRequest()
                  .withQueryString(query)
                  .withResultConfiguration(
                      new ResultConfiguration()
                          .withOutputLocation(String.format("s3://%s/results/", bucket))));
      log.debug(
          "Started query: '{}' to add partition: '{}' to table: '{}'",
          result.getQueryExecutionId(),
          partition,
          table);
      partitions.add(partition);
    } catch (AmazonClientException e) {
      log.error(
          "Failed to start add partition query for year={}/month={}/day={}/hour={}",
          year,
          month,
          day,
          hour,
          e);
    }
  }
}
 

public AmazonAthena getClient(String region) {
  AmazonAthena client = athenaClientMap.get(region);

  if (client == null) {
    client = AmazonAthenaClient.builder().withRegion(region).build();
    athenaClientMap.put(region, client);
  }

  return client;
}
 

/**
 * This code calls Athena and retrieves the results of a query.
 * The query must be in a completed state before the results can be retrieved and
 * paginated. The first row of results are the column headers.
 */
private static void processResultRows(AmazonAthena athenaClient, String queryExecutionId)
{
    GetQueryResultsRequest getQueryResultsRequest = new GetQueryResultsRequest()
            // Max Results can be set but if its not set,
            // it will choose the maximum page size
            // As of the writing of this code, the maximum value is 1000
            // .withMaxResults(1000)
            .withQueryExecutionId(queryExecutionId);

    GetQueryResultsResult getQueryResultsResult = athenaClient.getQueryResults(getQueryResultsRequest);
    List<ColumnInfo> columnInfoList = getQueryResultsResult.getResultSet().getResultSetMetadata().getColumnInfo();

    while (true) {
        List<Row> results = getQueryResultsResult.getResultSet().getRows();
        for (Row row : results) {
            // Process the row. The first row of the first page holds the column names.
            processRow(row, columnInfoList);
        }
        // If nextToken is null, there are no more pages to read. Break out of the loop.
        if (getQueryResultsResult.getNextToken() == null) {
            break;
        }
        getQueryResultsResult = athenaClient.getQueryResults(
                getQueryResultsRequest.withNextToken(getQueryResultsResult.getNextToken()));
    }
}
 

public static void main(String[] args) throws Exception
{
    // Build an Athena client
    AthenaClientFactory factory = new AthenaClientFactory();
    AmazonAthena athenaClient = factory.createClient();

    // Build the request
    ListQueryExecutionsRequest listQueryExecutionsRequest = new ListQueryExecutionsRequest();

    // Get the list results.
    ListQueryExecutionsResult listQueryExecutionsResult = athenaClient.listQueryExecutions(listQueryExecutionsRequest);

    // Process the results.
    boolean hasMoreResults = true;
    while (hasMoreResults) {
        List<String> queryExecutionIds = listQueryExecutionsResult.getQueryExecutionIds();
        // process queryExecutionIds.

        System.out.println(queryExecutionIds);

        //If nextToken is not null, then there are more results. Get the next page of results.
        if (listQueryExecutionsResult.getNextToken() != null) {
            listQueryExecutionsResult = athenaClient.listQueryExecutions(
                    listQueryExecutionsRequest.withNextToken(listQueryExecutionsResult.getNextToken()));
        }
        else {
            hasMoreResults = false;
        }
    }
}
 

public static void main(String[] args) throws Exception
{
    // Build an Athena client
    AthenaClientFactory factory = new AthenaClientFactory();
    AmazonAthena athenaClient = factory.createClient();

    // Build the request
    ListNamedQueriesRequest listNamedQueriesRequest = new ListNamedQueriesRequest();

    // Get the list results.
    ListNamedQueriesResult listNamedQueriesResult = athenaClient.listNamedQueries(listNamedQueriesRequest);

    // Process the results.
    boolean hasMoreResults = true;

    while (hasMoreResults) {
        List<String> namedQueryIds = listNamedQueriesResult.getNamedQueryIds();
        // process named query IDs

        // If nextToken is not null,  there are more results. Get the next page of results.
        if (listNamedQueriesResult.getNextToken() != null) {
            listNamedQueriesResult = athenaClient.listNamedQueries(
                    listNamedQueriesRequest.withNextToken(listNamedQueriesResult.getNextToken()));
        }
        else {
            hasMoreResults = false;
        }
    }
}
 

/**
 * Wait for an Athena query to complete, fail or to be cancelled. This is done by polling Athena over an
 * interval of time. If a query fails or is cancelled, then it will throw an exception.
 */

     private static void waitForQueryToComplete(AmazonAthena athenaClient, String queryExecutionId) throws InterruptedException
{
    GetQueryExecutionRequest getQueryExecutionRequest = new GetQueryExecutionRequest()
            .withQueryExecutionId(queryExecutionId);

    GetQueryExecutionResult getQueryExecutionResult = null;
    boolean isQueryStillRunning = true;
    while (isQueryStillRunning) {
        getQueryExecutionResult = athenaClient.getQueryExecution(getQueryExecutionRequest);
        String queryState = getQueryExecutionResult.getQueryExecution().getStatus().getState();
        if (queryState.equals(QueryExecutionState.FAILED.toString())) {
            throw new RuntimeException("Query Failed to run with Error Message: " + getQueryExecutionResult.getQueryExecution().getStatus().getStateChangeReason());
        }
        else if (queryState.equals(QueryExecutionState.CANCELLED.toString())) {
            throw new RuntimeException("Query was cancelled.");
        }
        else if (queryState.equals(QueryExecutionState.SUCCEEDED.toString())) {
            isQueryStillRunning = false;
        }
        else {
            // Sleep an amount of time before retrying again.
            Thread.sleep(ExampleConstants.SLEEP_AMOUNT_IN_MS);
        }
        System.out.println("Current Status is: " + queryState);
    }
}
 

private static String getNamedQueryId(AmazonAthena athenaClient)
{
    // Create the NameQuery Request.
    CreateNamedQueryRequest createNamedQueryRequest = new CreateNamedQueryRequest()
            .withDatabase(ExampleConstants.ATHENA_DEFAULT_DATABASE)
            .withQueryString(ExampleConstants.ATHENA_SAMPLE_QUERY)
            .withName("SampleQueryName")
            .withDescription("Sample Description");

    // Create the named query. If it fails, an exception is thrown.
    CreateNamedQueryResult createNamedQueryResult = athenaClient.createNamedQuery(createNamedQueryRequest);
    return createNamedQueryResult.getNamedQueryId();
}
 

public static void main(String[] args) throws Exception
{
    // Build an Athena client
    AthenaClientFactory factory = new AthenaClientFactory();
    AmazonAthena athenaClient = factory.createClient();

    String sampleNamedQueryId = getNamedQueryId(athenaClient);

    // Create the delete named query request
    DeleteNamedQueryRequest deleteNamedQueryRequest = new DeleteNamedQueryRequest()
            .withNamedQueryId(sampleNamedQueryId);

    // Delete the named query
    DeleteNamedQueryResult deleteNamedQueryResult = athenaClient.deleteNamedQuery(deleteNamedQueryRequest);
}
 

public static void main(String[] args) throws InterruptedException
{
    // Build an AmazonAthena client
    AthenaClientFactory factory = new AthenaClientFactory();
    AmazonAthena athenaClient = factory.createClient();

    String queryExecutionId = submitAthenaQuery(athenaClient);

    waitForQueryToComplete(athenaClient, queryExecutionId);

    processResultRows(athenaClient, queryExecutionId);
}
 

/**
 * @param metadataHandlerMap catalog -> JdbcMetadataHandler
 */
@VisibleForTesting
MultiplexingJdbcMetadataHandler(final AWSSecretsManager secretsManager, final AmazonAthena athena, final JdbcConnectionFactory jdbcConnectionFactory,
        final Map<String, JdbcMetadataHandler> metadataHandlerMap, final DatabaseConnectionConfig databaseConnectionConfig)
{
    super(databaseConnectionConfig, secretsManager, athena, jdbcConnectionFactory);
    this.metadataHandlerMap = Validate.notEmpty(metadataHandlerMap, "metadataHandlerMap must not be empty");

    if (this.metadataHandlerMap.size() > MAX_CATALOGS_TO_MULTIPLEX) {
        throw new RuntimeException("Max 100 catalogs supported in multiplexer.");
    }
}