Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added credentials option for iter_bucket. #372

Merged
merged 7 commits into from
Mar 11, 2020
Merged

Added credentials option for iter_bucket. #372

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
15edc08
Added credentials option for iter_bucket.
derpferd Oct 31, 2019
ad42535
Applied PR Feedback
derpferd Nov 3, 2019
9b0de4a
Formatting changes
derpferd Jan 15, 2020
731b150
explicitly construct session with session_kwargs
mpenkov Mar 8, 2020
50e4526
fix bug, update unit tests
mpenkov Mar 8, 2020
16e83e7
remove test with tricky mocks
mpenkov Mar 8, 2020
d0c3e43
isolate credentials test
mpenkov Mar 9, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 31 additions & 7 deletions smart_open/s3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -614,7 +614,9 @@ def _accept_all(key):


def iter_bucket(bucket_name, prefix='', accept_key=None,
key_limit=None, workers=16, retries=3):
key_limit=None, workers=16, retries=3,
aws_access_key_id=None, aws_secret_access_key=None,
aws_session_token=None):
"""
Iterate and download all S3 objects under `s3://bucket_name/prefix`.

Expand All @@ -634,6 +636,13 @@ def iter_bucket(bucket_name, prefix='', accept_key=None,
The number of subprocesses to use.
retries: int, optional
The number of time to retry a failed download.
aws_access_key_id: str, optional
The access key to use when creating the AWS client.
aws_secret_access_key: str, optional
The secret key to use when creating the AWS client.
aws_session_token: str, optional
The session token to use when creating the AWS client.


Yields
------
Expand Down Expand Up @@ -674,8 +683,14 @@ def iter_bucket(bucket_name, prefix='', accept_key=None,
pass

total_size, key_no = 0, -1
key_iterator = _list_bucket(bucket_name, prefix=prefix, accept_key=accept_key)
download_key = functools.partial(_download_key, bucket_name=bucket_name, retries=retries)
key_iterator = _list_bucket(bucket_name, prefix=prefix, accept_key=accept_key,
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token)
download_key = functools.partial(_download_key, bucket_name=bucket_name, retries=retries,
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token)

with _create_process_pool(processes=workers) as pool:
result_iterator = pool.imap_unordered(download_key, key_iterator)
Expand All @@ -694,8 +709,13 @@ def iter_bucket(bucket_name, prefix='', accept_key=None,
logger.info("processed %i keys, total size %i" % (key_no + 1, total_size))


def _list_bucket(bucket_name, prefix='', accept_key=lambda k: True):
client = boto3.client('s3')
def _list_bucket(bucket_name, prefix='', accept_key=lambda k: True,
aws_access_key_id=None, aws_secret_access_key=None,
aws_session_token=None):
client = boto3.client('s3',
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hanging indent here please.

aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token)
ctoken = None

while True:
Expand All @@ -720,14 +740,18 @@ def _list_bucket(bucket_name, prefix='', accept_key=lambda k: True):
break


def _download_key(key_name, bucket_name=None, retries=3):
def _download_key(key_name, bucket_name=None, retries=3,
aws_access_key_id=None, aws_secret_access_key=None,
aws_session_token=None):
if bucket_name is None:
raise ValueError('bucket_name may not be None')

#
# https://geekpete.com/blog/multithreading-boto3/
#
session = boto3.session.Session()
session = boto3.session.Session(aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token)
s3 = session.resource('s3')
bucket = s3.Bucket(bucket_name)

Expand Down
26 changes: 26 additions & 0 deletions smart_open/tests/test_s3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -460,6 +460,32 @@ def test_list_bucket_long(self):
expected = ['key_%d' % x for x in range(num_keys)]
self.assertEqual(sorted(keys), sorted(expected))

@mock.patch('boto3.client')
@mock.patch('boto3.session.Session')
def test_credentials(self, mock_session, mock_client):
mock_s3 = mock.Mock()
mock_client.return_value = mock_s3
mock_s3.list_objects_v2.return_value = {'Contents': [{'Key': 'test_key'}]}

list(smart_open.s3.iter_bucket(
BUCKET_NAME,
workers=None,
aws_access_key_id='access_id',
aws_secret_access_key='access_secret'
))

mock_client.assert_called_with(
's3',
aws_access_key_id='access_id',
aws_secret_access_key='access_secret',
aws_session_token=None
)
mock_session.assert_called_with(
aws_access_key_id='access_id',
aws_secret_access_key='access_secret',
aws_session_token=None
)

def test_old(self):
"""Does s3_iter_bucket work correctly?"""
#
Expand Down