diff --git a/requirements-test.txt b/requirements-test.txt index a8373019..cf37e641 100644 --- a/requirements-test.txt +++ b/requirements-test.txt @@ -1,7 +1,7 @@ -e git://github.com/boto/botocore.git@develop#egg=botocore nose==1.3.3 mock==1.3.0 -coverage==4.0.1 +coverage==4.4 wheel==0.24.0 # Note you need at least pip --version of 6.0 or # higher to be able to pick on these version specifiers. diff --git a/s3transfer/download.py b/s3transfer/download.py index d7ca2b4b..a15fd930 100644 --- a/s3transfer/download.py +++ b/s3transfer/download.py @@ -382,6 +382,7 @@ def _submit_download_request(self, client, config, osutil, transfer_future) # Get the needed callbacks for the task + response_callbacks = get_callbacks(transfer_future, 'response') progress_callbacks = get_callbacks(transfer_future, 'progress') # Get any associated tags for the get object task. @@ -401,7 +402,8 @@ def _submit_download_request(self, client, config, osutil, 'key': call_args.key, 'fileobj': fileobj, 'extra_args': call_args.extra_args, - 'callbacks': progress_callbacks, + 'response_callbacks': response_callbacks, + 'progress_callbacks': progress_callbacks, 'max_attempts': config.num_download_attempts, 'download_output_manager': download_output_manager, 'io_chunksize': config.io_chunksize, @@ -420,6 +422,7 @@ def _submit_ranged_download_request(self, client, config, osutil, call_args = transfer_future.meta.call_args # Get the needed progress callbacks for the task + response_callbacks = get_callbacks(transfer_future, 'response') progress_callbacks = get_callbacks(transfer_future, 'progress') # Get a handle to the file that will be used for writing downloaded @@ -463,7 +466,8 @@ def _submit_ranged_download_request(self, client, config, osutil, 'key': call_args.key, 'fileobj': fileobj, 'extra_args': extra_args, - 'callbacks': progress_callbacks, + 'response_callbacks': response_callbacks, + 'progress_callbacks': progress_callbacks, 'max_attempts': config.num_download_attempts, 'start_index': i * part_size, 'download_output_manager': download_output_manager, @@ -494,7 +498,8 @@ def _calculate_range_param(self, part_size, part_index, num_parts): class GetObjectTask(Task): - def _main(self, client, bucket, key, fileobj, extra_args, callbacks, + def _main(self, client, bucket, key, fileobj, extra_args, + response_callbacks, progress_callbacks, max_attempts, download_output_manager, io_chunksize, start_index=0, bandwidth_limiter=None): """Downloads an object and places content into io queue @@ -504,7 +509,8 @@ def _main(self, client, bucket, key, fileobj, extra_args, callbacks, :param key: The key to download from :param fileobj: The file handle to write content to :param exta_args: Any extra arguements to include in GetObject request - :param callbacks: List of progress callbacks to invoke on download + :param response_callbacks: List of response callbacks to invoke on download + :param progress_callbacks: List of progress callbacks to invoke on download :param max_attempts: The number of retries to do when downloading :param download_output_manager: The download output manager associated with the current download. @@ -520,8 +526,10 @@ def _main(self, client, bucket, key, fileobj, extra_args, callbacks, try: response = client.get_object( Bucket=bucket, Key=key, **extra_args) + for callback in response_callbacks: + callback(response=response) streaming_body = StreamReaderProgress( - response['Body'], callbacks) + response['Body'], progress_callbacks) if bandwidth_limiter: streaming_body = \ bandwidth_limiter.get_bandwith_limited_stream( @@ -551,7 +559,7 @@ def _main(self, client, bucket, key, fileobj, extra_args, callbacks, # are trying to download the stream again and all progress # for this GetObject has been lost. invoke_progress_callbacks( - callbacks, start_index - current_index) + progress_callbacks, start_index - current_index) continue raise RetriesExceededError(last_exception) diff --git a/s3transfer/subscribers.py b/s3transfer/subscribers.py index 329b01ec..e5e522af 100644 --- a/s3transfer/subscribers.py +++ b/s3transfer/subscribers.py @@ -61,6 +61,9 @@ def on_queued(self, future, **kwargs): """ pass + def on_response(self, future, response, **kwargs): + pass + def on_progress(self, future, bytes_transferred, **kwargs): """Callback to be invoked when progress is made on transfer diff --git a/tests/__init__.py b/tests/__init__.py index 3cef6e76..5a965b1b 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -207,12 +207,16 @@ def rename_file(self, current_filename, new_filename): class RecordingSubscriber(BaseSubscriber): def __init__(self): self.on_queued_calls = [] + self.on_response_calls = [] self.on_progress_calls = [] self.on_done_calls = [] def on_queued(self, **kwargs): self.on_queued_calls.append(kwargs) + def on_response(self, **kwargs): + self.on_response_calls.append(kwargs) + def on_progress(self, **kwargs): self.on_progress_calls.append(kwargs) diff --git a/tests/unit/test_download.py b/tests/unit/test_download.py index b2d9fe55..f68804e5 100644 --- a/tests/unit/test_download.py +++ b/tests/unit/test_download.py @@ -530,7 +530,8 @@ def setUp(self): self.bucket = 'mybucket' self.key = 'mykey' self.extra_args = {} - self.callbacks = [] + self.response_callbacks = [] + self.progress_callbacks = [] self.max_attempts = 5 self.io_executor = BoundedExecutor(1000, 1) self.content = b'my content' @@ -546,7 +547,8 @@ def get_download_task(self, **kwargs): default_kwargs = { 'client': self.client, 'bucket': self.bucket, 'key': self.key, 'fileobj': self.fileobj, 'extra_args': self.extra_args, - 'callbacks': self.callbacks, + 'response_callbacks': self.response_callbacks, + 'progress_callbacks': self.progress_callbacks, 'max_attempts': self.max_attempts, 'download_output_manager': self.download_output_manager, 'io_chunksize': self.io_chunksize,