com.amazonaws.services.secretsmanager.AWSSecretsManager Java Examples

protected GetSecretValueResult getSecretValue(String secretRegion, String secretName) {
  AWSSecretsManager client =
      AWSSecretsManagerClientBuilder.standard().withRegion(secretRegion).build();

  GetSecretValueRequest getSecretValueRequest =
      new GetSecretValueRequest().withSecretId(secretName);

  try {
    return client.getSecretValue(getSecretValueRequest);
  } catch (AWSSecretsManagerException e) {
    throw new SecretException(
        String.format(
            "An error occurred when using AWS Secrets Manager to fetch: [secretName: %s, secretRegion: %s]",
            secretName, secretRegion),
        e);
  }
}
 

/**
 * Gets the username failure invalid content test.
 *
 * @return the username failure invalid content test
 * @throws Exception the exception
 */
@Test(expected=RuntimeException.class)
public void getUsernameFailureInvalidContentTest() throws Exception {
	assertNotNull(new SecretCredentialsManagerImpl(null, "true", null, null) {
		
		@Override
		protected SecretCache getSecretCache() {
			return new SecretCache(Mockito.mock(AWSSecretsManager.class)) {
								 
				@Override
				public String getSecretString(final String secretId) {
					return "{ \"wrongusernamefield\" : \"username\" }";
				}
			};
		}
	}.getUsername());
}
 

/**
 * Gets the username success test.
 *
 * @return the username success test
 * @throws Exception the exception
 */
@Test
public void getUsernameSuccessTest() throws Exception {
	assertNotNull(new SecretCredentialsManagerImpl(null, "true", null, null) {
		
		@Override
		protected SecretCache getSecretCache() {
			return new SecretCache(Mockito.mock(AWSSecretsManager.class)) {
								 
				@Override
				public String getSecretString(final String secretId) {
					return "{ \"username\" : \"username\" }";
				}
			};
		}
	}.getUsername());
}
 

/**
 * Gets the password success test.
 *
 * @return the password success test
 * @throws Exception the exception
 */
@Test
public void getPasswordSuccessTest() throws Exception {
	final SecretCredentialsManagerImpl secretCredentialsManagerImpl = new SecretCredentialsManagerImpl(null, "true", null, null) {
		
		@Override
		protected SecretCache getSecretCache() {
			return new SecretCache(Mockito.mock(AWSSecretsManager.class)) {
								 
				@Override
				public String getSecretString(final String secretId) {
					return "{ \"password\" : \"password\", \"username\" : \"username\" }";
				}
			};
		}
	};
	assertNotNull(secretCredentialsManagerImpl.getUsername());
	assertNotNull(secretCredentialsManagerImpl.getPassword());
}
 

@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());
}
 

@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\"}"));
    DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", JdbcConnectionFactory.DatabaseEngine.MYSQL,
            "mysql://jdbc:mysql://hostname/${testSecret}", "testSecret");
    this.jdbcMetadataHandler = new JdbcMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory)
    {
        @Override
        public Schema getPartitionSchema(final String catalogName)
        {
            return PARTITION_SCHEMA;
        }

        @Override
        public void getPartitions(final BlockWriter blockWriter, final GetTableLayoutRequest getTableLayoutRequest, QueryStatusChecker queryStatusChecker)
        {
        }

        @Override
        public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsRequest getSplitsRequest)
        {
            return null;
        }
    };
    this.federatedIdentity = Mockito.mock(FederatedIdentity.class);
    this.blockAllocator = Mockito.mock(BlockAllocator.class);
}
 

@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);
}
 

@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);
    Mockito.when(this.secretsManager.getSecretValue(Mockito.eq(new GetSecretValueRequest().withSecretId("testSecret")))).thenReturn(new GetSecretValueResult().withSecretString("{\"username\": \"testUser\", \"password\": \"testPassword\"}"));
    this.postGreSqlMetadataHandler = new PostGreSqlMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, this.jdbcConnectionFactory);
    this.federatedIdentity = Mockito.mock(FederatedIdentity.class);
}
 

@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 PostGreSqlQueryStringBuilder("\"");
    final DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", JdbcConnectionFactory.DatabaseEngine.MYSQL,
            "mysql://jdbc:mysql://hostname/user=A&password=B");

    this.postGreSqlRecordHandler = new PostGreSqlRecordHandler(databaseConnectionConfig, amazonS3, secretsManager, athena, jdbcConnectionFactory, jdbcSplitQueryBuilder);
}
 

@Before
public void setup()
        throws SQLException
{
    this.connection = Mockito.mock(Connection.class, Mockito.RETURNS_DEEP_STUBS);
    this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
    Mockito.when(this.jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(this.connection);
    this.amazonS3 = Mockito.mock(AmazonS3.class);
    this.secretsManager = Mockito.mock(AWSSecretsManager.class);
    this.athena = Mockito.mock(AmazonAthena.class);
    this.queryStatusChecker = Mockito.mock(QueryStatusChecker.class);
    Mockito.when(this.secretsManager.getSecretValue(Mockito.eq(new GetSecretValueRequest().withSecretId("testSecret")))).thenReturn(new GetSecretValueResult().withSecretString("{\"username\": \"testUser\", \"password\": \"testPassword\"}"));
    this.preparedStatement = Mockito.mock(PreparedStatement.class);
    Mockito.when(this.connection.prepareStatement("someSql")).thenReturn(this.preparedStatement);
    DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", JdbcConnectionFactory.DatabaseEngine.MYSQL,
            "mysql://jdbc:mysql://hostname/${testSecret}", "testSecret");
    this.jdbcRecordHandler = new JdbcRecordHandler(this.amazonS3, this.secretsManager, this.athena, databaseConnectionConfig, this.jdbcConnectionFactory)
    {
        @Override
        public PreparedStatement buildSplitSql(Connection jdbcConnection, String catalogName, TableName tableName, Schema schema, Constraints constraints, Split split)
                throws SQLException
        {
            return jdbcConnection.prepareStatement("someSql");
        }
    };
    this.federatedIdentity = Mockito.mock(FederatedIdentity.class);
}
 

@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;
}
 

@Before
public void setUp()
{
    logger.info("setUpBefore - enter");
    allocator = new BlockAllocatorImpl();
    metadataHandler = new ExampleMetadataHandler(new LocalKeyFactory(),
            mock(AWSSecretsManager.class),
            mock(AmazonAthena.class),
            "spill-bucket",
            "spill-prefix");
    logger.info("setUpBefore - exit");
}
 

/**
 * Full DI constructor used mostly for testing
 *
 * @param awsGlue The glue client to use.
 * @param encryptionKeyFactory The EncryptionKeyFactory to use for spill encryption.
 * @param secretsManager The AWSSecretsManager client that can be used when attempting to resolve secrets.
 * @param athena The Athena client that can be used to fetch query termination status to fast-fail this handler.
 * @param spillBucket The S3 Bucket to use when spilling results.
 * @param spillPrefix The S3 prefix to use when spilling results.
 */
@VisibleForTesting
protected GlueMetadataHandler(AWSGlue awsGlue,
        EncryptionKeyFactory encryptionKeyFactory,
        AWSSecretsManager secretsManager,
        AmazonAthena athena,
        String sourceType,
        String spillBucket,
        String spillPrefix)
{
    super(encryptionKeyFactory, secretsManager, athena, sourceType, spillBucket, spillPrefix);
    this.awsGlue = awsGlue;
}
 

/**
 * @param sourceType Used to aid in logging diagnostic info when raising a support case.
 */
public RecordHandler(AmazonS3 amazonS3, AWSSecretsManager secretsManager, AmazonAthena athena, String sourceType)
{
    this.sourceType = sourceType;
    this.amazonS3 = amazonS3;
    this.secretsManager = new CachableSecretsManager(secretsManager);
    this.athena = athena;
}
 

/**
 * Full DI constructor used mostly for testing
 *
 * @param keyFactory The EncryptionKeyFactory to use for spill encryption.
 * @param awsSecretsManager The AWSSecretsManager client that can be used when attempting to resolve secrets.
 * @param athena The Athena client that can be used to fetch query termination status to fast-fail this handler.
 * @param spillBucket The S3 Bucket to use when spilling results.
 * @param spillPrefix The S3 prefix to use when spilling results.
 */
@VisibleForTesting
protected ExampleMetadataHandler(EncryptionKeyFactory keyFactory,
        AWSSecretsManager awsSecretsManager,
        AmazonAthena athena,
        String spillBucket,
        String spillPrefix)
{
    super(keyFactory, awsSecretsManager, athena, SOURCE_TYPE, spillBucket, spillPrefix);
    //Read the Lambda environment variable for controlling simulated throttles.
    this.simulateThrottle = (System.getenv(SIMULATE_THROTTLES) == null) ? 0 : Integer.parseInt(System.getenv(SIMULATE_THROTTLES));
}
 

@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);
}
 

@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;
}
 

@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 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 RedisRecordHandler(AmazonS3 amazonS3,
        AWSSecretsManager secretsManager,
        AmazonAthena athena,
        JedisPoolFactory jedisPoolFactory)
{
    super(amazonS3, secretsManager, athena, SOURCE_TYPE);
    this.amazonS3 = amazonS3;
    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;
}
 

@Override
public byte[] decrypt(EncryptedSecret encryptedSecret) {
  String secretName = encryptedSecret.getParams().get(SECRET_NAME);
  String secretRegion = encryptedSecret.getParams().get(SECRET_REGION);
  String secretKey = encryptedSecret.getParams().get(SECRET_KEY);

  AWSSecretsManager client =
      AWSSecretsManagerClientBuilder.standard().withRegion(secretRegion).build();

  byte[] binarySecret = null;
  GetSecretValueRequest getSecretValueRequest =
      new GetSecretValueRequest().withSecretId(secretName);
  GetSecretValueResult getSecretValueResult = null;

  try {
    getSecretValueResult = client.getSecretValue(getSecretValueRequest);
  } catch (Exception e) {
    log.error(
        "An error occurred when trying to use AWS Secrets Manager to fetch: [secretName: {}, secretRegion: {}, secretKey: {}]",
        secretName,
        secretRegion,
        secretKey,
        e);
    throw new RuntimeException("Failed to fetch secret from AWS Secrets Manager", e);
  }

  if (getSecretValueResult.getSecretString() != null) {
    String secret = getSecretValueResult.getSecretString();
    Gson gson = new Gson();
    Type type = new TypeToken<Map<String, String>>() {}.getType();
    Map<String, String> myMap = gson.fromJson(secret, type);
    binarySecret = myMap.get(secretKey).getBytes(StandardCharsets.UTF_8);
  } else {
    binarySecret = getSecretValueResult.getSecretBinary().array();
  }
  return binarySecret;
}
 

@Bean
@ConditionalOnMissingBean
AWSSecretsManager smClient(AwsSecretsManagerProperties awsSecretsManagerProperties) {
	return StringUtils.isNullOrEmpty(awsSecretsManagerProperties.getRegion())
			? AWSSecretsManagerClientBuilder.defaultClient()
			: AWSSecretsManagerClientBuilder.standard()
					.withRegion(awsSecretsManagerProperties.getRegion()).build();
}