def test_get_credentials_from_extra_with_s3_config_and_profile(
self, mock_get_connection, mock_parse_s3_config):
mock_connection = Connection(extra='{"s3_config_format": "aws", '
'"profile": "test", '
'"s3_config_file": "aws-credentials", '
'"region_name": "us-east-1"}')
mock_get_connection.return_value = mock_connection
hook = AwsHook()
hook._get_credentials(region_name=None)
mock_parse_s3_config.assert_called_once_with('aws-credentials', 'aws',
'test')
def _inject_aws_credentials(self):
if TRANSFER_SPEC in self.body and AWS_S3_DATA_SOURCE in self.body[
TRANSFER_SPEC]:
aws_hook = AwsHook(self.aws_conn_id)
aws_credentials = aws_hook.get_credentials()
aws_access_key_id = aws_credentials.access_key
aws_secret_access_key = aws_credentials.secret_key
self.body[TRANSFER_SPEC][AWS_S3_DATA_SOURCE][AWS_ACCESS_KEY] = {
ACCESS_KEY_ID: aws_access_key_id,
SECRET_ACCESS_KEY: aws_secret_access_key,
}
def test_get_credentials_from_extra_without_token(self,
mock_get_connection):
mock_connection = Connection(
extra='{"aws_access_key_id": "aws_access_key_id",'
'"aws_secret_access_key": "aws_secret_access_key"}')
mock_get_connection.return_value = mock_connection
hook = AwsHook()
credentials_from_hook = hook.get_credentials()
self.assertEqual(credentials_from_hook.access_key, 'aws_access_key_id')
self.assertEqual(credentials_from_hook.secret_key,
'aws_secret_access_key')
self.assertIsNone(credentials_from_hook.token)
def test_get_credentials_from_login_with_token(self, mock_get_connection):
mock_connection = Connection(
login='******',
password='******',
extra='{"aws_session_token": "test_token"}')
mock_get_connection.return_value = mock_connection
hook = AwsHook()
credentials_from_hook = hook.get_credentials()
self.assertEqual(credentials_from_hook.access_key, 'aws_access_key_id')
self.assertEqual(credentials_from_hook.secret_key,
'aws_secret_access_key')
self.assertEqual(credentials_from_hook.token, 'test_token')
def test_get_credentials_from_role_arn(self, mock_get_connection):
mock_connection = Connection(
extra='{"role_arn":"arn:aws:iam::123456:role/role_arn"}')
mock_get_connection.return_value = mock_connection
hook = AwsHook()
credentials_from_hook = hook.get_credentials()
self.assertIn("ASIA", credentials_from_hook.access_key)
# We assert the length instead of actual values as the values are random:
# Details: https://github.com/spulec/moto/commit/ab0d23a0ba2506e6338ae20b3fde70da049f7b03
self.assertEqual(20, len(credentials_from_hook.access_key))
self.assertEqual(40, len(credentials_from_hook.secret_key))
self.assertEqual(356, len(credentials_from_hook.token))
def test_get_credentials_from_login_without_token(self,
mock_get_connection):
mock_connection = Connection(
login='******',
password='******',
)
mock_get_connection.return_value = mock_connection
hook = AwsHook()
credentials_from_hook = hook.get_credentials()
self.assertEqual(credentials_from_hook.access_key, 'aws_access_key_id')
self.assertEqual(credentials_from_hook.secret_key,
'aws_secret_access_key')
self.assertIsNone(credentials_from_hook.token)
def get_iam_token(self, conn):
"""
Uses AWSHook to retrieve a temporary password to connect to MySQL
Port is required. If none is provided, default 3306 is used
"""
from airflow.providers.amazon.aws.hooks.aws_hook import AwsHook
aws_conn_id = conn.extra_dejson.get('aws_conn_id', 'aws_default')
aws_hook = AwsHook(aws_conn_id)
if conn.port is None:
port = 3306
else:
port = conn.port
client = aws_hook.get_client_type('rds')
token = client.generate_db_auth_token(conn.host, port, conn.login)
return token, port
def hook(self) -> AwsHook:
"""
An AWS API connection manager (wraps boto3)
:return: the connected hook to AWS
:rtype: AwsHook
"""
if self._hook is None:
self._hook = AwsHook(aws_conn_id=self.aws_conn_id)
return self._hook
def expand_role(self):
if 'RoleArn' in self.config:
hook = AwsHook(self.aws_conn_id)
self.config['RoleArn'] = hook.expand_role(self.config['RoleArn'])
def get_hook(self):
return AwsHook(aws_conn_id=self.aws_conn_id)
def expand_role(self):
if 'TrainingJobDefinition' in self.config:
config = self.config['TrainingJobDefinition']
if 'RoleArn' in config:
hook = AwsHook(self.aws_conn_id)
config['RoleArn'] = hook.expand_role(config['RoleArn'])
def get_hook(self):
"""Create and return an AwsHook."""
return AwsHook(
aws_conn_id=self.aws_conn_id
)