def file_size(self):
s3 = AWS().s3_connection()
s3_object = s3.Object(self.bucket_name, self.key_name)
try:
return s3_object.content_length
except ClientError:
# file does not exist
return 0
def upload_local_file_to_s3_path(local_path, s3_path):
bucket_name, key_name, _ = parse_s3_path(s3_path)
s3 = AWS().s3_connection()
s3_file = s3.Object(bucket_name, key_name)
s3_file.upload_file(local_path)
if s3_file.content_length == 0:
raise NoDataFoundError(
'The file you\'ve uploaded to S3 has a size of 0 KB')
def _compose_unload_query(self, query, s3_path, options):
query_commands = ['MANIFEST'] if options.get(
'is_parallel_unload', False) else ['PARALLEL OFF']
query_commands.append('ALLOWOVERWRITE') if options.get(
'allow_overwrite', False) else None
query_commands.append('DELIMITER \'%s\'' % options.get('delimiter'))
query_commands.append('ADDQUOTES') if options.get('add_quotes',
False) else None
query_commands.append('ESCAPE') if options.get('escape',
False) else None
query_commands.append('HEADER') if options.get('header',
False) else None
query_commands.append(str(
options.get('compression_type'))) if options.get(
'compression_type', False) else None
query_commands.append('MAXFILESIZE %s' %
options.get('max_file_size')) if options.get(
'max_file_size', False) else None
unload_query_options = ' '.join(query_commands).strip()
unload_query = """UNLOAD ('{query}') TO '{s3_path}'
CREDENTIALS '{aws_connection_string}'
{unload_query_options};""".format(
s3_path=s3_path,
aws_connection_string=AWS().connection_string(),
query=query,
unload_query_options=unload_query_options)
unload_query = unload_query.replace("\n", "").replace(
" ", " ")
return unload_query
def test_initializes_and_returns_s3_connection(self, mock_boto_session,
mock_boto_resource):
mock_aws_credentials = Mock()
key_property = PropertyMock(return_value='aws_mock_key')
secret_property = PropertyMock(return_value='aws_mock_secret')
token_property = PropertyMock(return_value='aws_mock_token')
type(mock_aws_credentials).access_key = key_property
type(mock_aws_credentials).secret_key = secret_property
type(mock_aws_credentials).token = token_property
mock_boto_session.return_value.get_credentials.return_value = mock_aws_credentials
mock_s3_connection = Mock()
mock_boto_resource.return_value = mock_s3_connection
expected_calls = [
call('s3',
aws_secret_access_key='aws_mock_secret',
aws_access_key_id='aws_mock_key',
aws_session_token='aws_mock_token'),
call('s3',
aws_secret_access_key='aws_mock_secret',
aws_access_key_id='aws_mock_key',
aws_session_token='aws_mock_token')
]
s3_connection = AWS().s3_connection()
self.assertEqual(s3_connection, mock_s3_connection)
mock_boto_resource.assert_has_calls(expected_calls, any_order=True)
def test_initializes_and_returns_athena_connection(self, mock_boto_session,
mock_boto_client, _):
mock_aws_credentials = Mock()
key_property = PropertyMock(return_value='aws_mock_key')
secret_property = PropertyMock(return_value='aws_mock_secret')
token_property = PropertyMock(return_value='aws_mock_token')
region_property = PropertyMock(return_value='aws_mock_region_name')
type(mock_aws_credentials).access_key = key_property
type(mock_aws_credentials).secret_key = secret_property
type(mock_aws_credentials).token = token_property
mock_boto_session.return_value.get_credentials.return_value = mock_aws_credentials
mock_local_session = Mock()
type(mock_local_session).region_name = region_property
mock_boto_session.return_value = mock_local_session
mock_local_session.get_credentials.return_value = mock_aws_credentials
mock_athena_connection = Mock()
mock_boto_client.return_value = mock_athena_connection
expected_call = [
call('athena',
aws_secret_access_key='aws_mock_secret',
aws_access_key_id='aws_mock_key',
aws_session_token='aws_mock_token',
region_name='aws_mock_region_name')
]
athena_connection = AWS().athena_connection()
self.assertEqual(athena_connection, mock_athena_connection)
mock_boto_client.assert_has_calls(expected_call, any_order=True)
def test_can_connect_without_s3_base_path(self, *_):
config.S3_BASE_PATH = None
try:
aws_object = AWS().s3_connection()
except Exception as e:
self.fail(
"AWS() unexpectedly raised an exception related to S3_BASE_PATH: `{}`"
.format(e))
def test_from_local_authentication(self, mock_boto_session,
mock_boto_connection_request):
mock_aws_credentials = Mock()
key_property = PropertyMock(return_value='aws_mock_key')
secret_property = PropertyMock(return_value='aws_mock_secret')
type(mock_aws_credentials).access_key = key_property
type(mock_aws_credentials).secret_key = secret_property
mock_boto_session.return_value.get_credentials.return_value = mock_aws_credentials
connection_string = AWS().connection_string()
self.assertEqual(connection_string, self.LOCAL_CONNECTION_STRING)
def test_from_ec2_iam_authentication(self, mock_boto_resource,
mock_boto_session, mock_aws_request):
mock_resource = Mock()
client_mock = Mock()
meta_mock = Mock()
client_mock.head_bucket.side_effect = \
botocore.exceptions.ClientError(error_response={'Error': {'Code': '403', 'Message': 'NotFound'}}, operation_name='HeadBucket')
meta_property = PropertyMock(return_value=meta_mock)
type(mock_resource).meta = meta_property
client_property = PropertyMock(return_value=client_mock)
type(meta_mock).client = client_property
mock_boto_resource.return_value = mock_resource
mock_aws_request.return_value = {
'AccessKeyId': 'aws_mock_key',
'SecretAccessKey': 'aws_mock_secret',
'Token': 'aws_mock_token'
}
connection_string = AWS().connection_string()
self.assertEqual(connection_string, self.EC2_CONNECTION_STRING)
def download_from_s3_to_local_file(s3_path, local_path):
bucket_name, key_name, _ = parse_s3_path(s3_path)
s3 = AWS().s3_connection()
s3_file = s3.Object(bucket_name, key_name)
s3_file.download_file(local_path)
def connection_string(self):
return AWS().connection_string()