# Copyright 2016 Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this file except in compliance with the License. A copy of # the License is located at # # http://aws.amazon.com/apache2.0/ # # or in the "license" file accompanying this file. This file is # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """Abstractions over S3's upload/download operations. This module provides high level abstractions for efficient uploads/downloads. It handles several things for the user: * Automatically switching to multipart transfers when a file is over a specific size threshold * Uploading/downloading a file in parallel * Throttling based on max bandwidth * Progress callbacks to monitor transfers * Retries. While botocore handles retries for streaming uploads, it is not possible for it to handle retries for streaming downloads. This module handles retries for both cases so you don't need to implement any retry logic yourself. This module has a reasonable set of defaults. It also allows you to configure many aspects of the transfer process including: * Multipart threshold size * Max parallel downloads * Max bandwidth * Socket timeouts * Retry amounts There is no support for s3->s3 multipart copies at this time. .. _ref_s3transfer_usage: Usage ===== The simplest way to use this module is: .. code-block:: python client = boto3.client('s3', 'us-west-2') transfer = S3Transfer(client) # Upload /tmp/myfile to s3://bucket/key transfer.upload_file('/tmp/myfile', 'bucket', 'key') # Download s3://bucket/key to /tmp/myfile transfer.download_file('bucket', 'key', '/tmp/myfile') The ``upload_file`` and ``download_file`` methods also accept ``**kwargs``, which will be forwarded through to the corresponding client operation. Here are a few examples using ``upload_file``:: # Making the object public transfer.upload_file('/tmp/myfile', 'bucket', 'key', extra_args={'ACL': 'public-read'}) # Setting metadata transfer.upload_file('/tmp/myfile', 'bucket', 'key', extra_args={'Metadata': {'a': 'b', 'c': 'd'}}) # Setting content type transfer.upload_file('/tmp/myfile.json', 'bucket', 'key', extra_args={'ContentType': "application/json"}) The ``S3Transfer`` class also supports progress callbacks so you can provide transfer progress to users. Both the ``upload_file`` and ``download_file`` methods take an optional ``callback`` parameter. Here's an example of how to print a simple progress percentage to the user: .. code-block:: python class ProgressPercentage(object): def __init__(self, filename): self._filename = filename self._size = float(os.path.getsize(filename)) self._seen_so_far = 0 self._lock = threading.Lock() def __call__(self, bytes_amount): # To simplify we'll assume this is hooked up # to a single filename. with self._lock: self._seen_so_far += bytes_amount percentage = (self._seen_so_far / self._size) * 100 sys.stdout.write( "\r%s %s / %s (%.2f%%)" % (self._filename, self._seen_so_far, self._size, percentage)) sys.stdout.flush() transfer = S3Transfer(boto3.client('s3', 'us-west-2')) # Upload /tmp/myfile to s3://bucket/key and print upload progress. transfer.upload_file('/tmp/myfile', 'bucket', 'key', callback=ProgressPercentage('/tmp/myfile')) You can also provide a TransferConfig object to the S3Transfer object that gives you more fine grained control over the transfer. For example: .. code-block:: python client = boto3.client('s3', 'us-west-2') config = TransferConfig( multipart_threshold=8 * 1024 * 1024, max_concurrency=10, num_download_attempts=10, ) transfer = S3Transfer(client, config) transfer.upload_file('/tmp/foo', 'bucket', 'key') """ import concurrent.futures import functools import logging import math import os import queue import random import socket import string import threading from botocore.compat import six # noqa: F401 from botocore.exceptions import IncompleteReadError from botocore.vendored.requests.packages.urllib3.exceptions import ( ReadTimeoutError, ) import s3transfer.compat from s3transfer.exceptions import RetriesExceededError, S3UploadFailedError __author__ = 'Amazon Web Services' __version__ = '0.6.0' class NullHandler(logging.Handler): def emit(self, record): pass logger = logging.getLogger(__name__) logger.addHandler(NullHandler()) MB = 1024 * 1024 SHUTDOWN_SENTINEL = object() def random_file_extension(num_digits=8): return ''.join(random.choice(string.hexdigits) for _ in range(num_digits)) def disable_upload_callbacks(request, operation_name, **kwargs): if operation_name in ['PutObject', 'UploadPart'] and hasattr( request.body, 'disable_callback' ): request.body.disable_callback() def enable_upload_callbacks(request, operation_name, **kwargs): if operation_name in ['PutObject', 'UploadPart'] and hasattr( request.body, 'enable_callback' ): request.body.enable_callback() class QueueShutdownError(Exception): pass class ReadFileChunk: def __init__( self, fileobj, start_byte, chunk_size, full_file_size, callback=None, enable_callback=True, ): """ Given a file object shown below: |___________________________________________________| 0 | | full_file_size |----chunk_size---| start_byte :type fileobj: file :param fileobj: File like object :type start_byte: int :param start_byte: The first byte from which to start reading. :type chunk_size: int :param chunk_size: The max chunk size to read. Trying to read pass the end of the chunk size will behave like you've reached the end of the file. :type full_file_size: int :param full_file_size: The entire content length associated with ``fileobj``. :type callback: function(amount_read) :param callback: Called whenever data is read from this object. """ self._fileobj = fileobj self._start_byte = start_byte self._size = self._calculate_file_size( self._fileobj, requested_size=chunk_size, start_byte=start_byte, actual_file_size=full_file_size, ) self._fileobj.seek(self._start_byte) self._amount_read = 0 self._callback = callback self._callback_enabled = enable_callback @classmethod def from_filename( cls, filename, start_byte, chunk_size, callback=None, enable_callback=True, ): """Convenience factory function to create from a filename. :type start_byte: int :param start_byte: The first byte from which to start reading. :type chunk_size: int :param chunk_size: The max chunk size to read. Trying to read pass the end of the chunk size will behave like you've reached the end of the file. :type full_file_size: int :param full_file_size: The entire content length associated with ``fileobj``. :type callback: function(amount_read) :param callback: Called whenever data is read from this object. :type enable_callback: bool :param enable_callback: Indicate whether to invoke callback during read() calls. :rtype: ``ReadFileChunk`` :return: A new instance of ``ReadFileChunk`` """ f = open(filename, 'rb') file_size = os.fstat(f.fileno()).st_size return cls( f, start_byte, chunk_size, file_size, callback, enable_callback ) def _calculate_file_size( self, fileobj, requested_size, start_byte, actual_file_size ): max_chunk_size = actual_file_size - start_byte return min(max_chunk_size, requested_size) def read(self, amount=None): if amount is None: amount_to_read = self._size - self._amount_read else: amount_to_read = min(self._size - self._amount_read, amount) data = self._fileobj.read(amount_to_read) self._amount_read += len(data) if self._callback is not None and self._callback_enabled: self._callback(len(data)) return data def enable_callback(self): self._callback_enabled = True def disable_callback(self): self._callback_enabled = False def seek(self, where): self._fileobj.seek(self._start_byte + where) if self._callback is not None and self._callback_enabled: # To also rewind the callback() for an accurate progress report self._callback(where - self._amount_read) self._amount_read = where def close(self): self._fileobj.close() def tell(self): return self._amount_read def __len__(self): # __len__ is defined because requests will try to determine the length # of the stream to set a content length. In the normal case # of the file it will just stat the file, but we need to change that # behavior. By providing a __len__, requests will use that instead # of stat'ing the file. return self._size def __enter__(self): return self def __exit__(self, *args, **kwargs): self.close() def __iter__(self): # This is a workaround for http://bugs.python.org/issue17575 # Basically httplib will try to iterate over the contents, even # if its a file like object. This wasn't noticed because we've # already exhausted the stream so iterating over the file immediately # stops, which is what we're simulating here. return iter([]) class StreamReaderProgress: """Wrapper for a read only stream that adds progress callbacks.""" def __init__(self, stream, callback=None): self._stream = stream self._callback = callback def read(self, *args, **kwargs): value = self._stream.read(*args, **kwargs) if self._callback is not None: self._callback(len(value)) return value class OSUtils: def get_file_size(self, filename): return os.path.getsize(filename) def open_file_chunk_reader(self, filename, start_byte, size, callback): return ReadFileChunk.from_filename( filename, start_byte, size, callback, enable_callback=False ) def open(self, filename, mode): return open(filename, mode) def remove_file(self, filename): """Remove a file, noop if file does not exist.""" # Unlike os.remove, if the file does not exist, # then this method does nothing. try: os.remove(filename) except OSError: pass def rename_file(self, current_filename, new_filename): s3transfer.compat.rename_file(current_filename, new_filename) class MultipartUploader: # These are the extra_args that need to be forwarded onto # subsequent upload_parts. UPLOAD_PART_ARGS = [ 'SSECustomerKey', 'SSECustomerAlgorithm', 'SSECustomerKeyMD5', 'RequestPayer', ] def __init__( self, client, config, osutil, executor_cls=concurrent.futures.ThreadPoolExecutor, ): self._client = client self._config = config self._os = osutil self._executor_cls = executor_cls def _extra_upload_part_args(self, extra_args): # Only the args in UPLOAD_PART_ARGS actually need to be passed # onto the upload_part calls. upload_parts_args = {} for key, value in extra_args.items(): if key in self.UPLOAD_PART_ARGS: upload_parts_args[key] = value return upload_parts_args def upload_file(self, filename, bucket, key, callback, extra_args): response = self._client.create_multipart_upload( Bucket=bucket, Key=key, **extra_args ) upload_id = response['UploadId'] try: parts = self._upload_parts( upload_id, filename, bucket, key, callback, extra_args ) except Exception as e: logger.debug( "Exception raised while uploading parts, " "aborting multipart upload.", exc_info=True, ) self._client.abort_multipart_upload( Bucket=bucket, Key=key, UploadId=upload_id ) raise S3UploadFailedError( "Failed to upload {} to {}: {}".format( filename, '/'.join([bucket, key]), e ) ) self._client.complete_multipart_upload( Bucket=bucket, Key=key, UploadId=upload_id, MultipartUpload={'Parts': parts}, ) def _upload_parts( self, upload_id, filename, bucket, key, callback, extra_args ): upload_parts_extra_args = self._extra_upload_part_args(extra_args) parts = [] part_size = self._config.multipart_chunksize num_parts = int( math.ceil(self._os.get_file_size(filename) / float(part_size)) ) max_workers = self._config.max_concurrency with self._executor_cls(max_workers=max_workers) as executor: upload_partial = functools.partial( self._upload_one_part, filename, bucket, key, upload_id, part_size, upload_parts_extra_args, callback, ) for part in executor.map(upload_partial, range(1, num_parts + 1)): parts.append(part) return parts def _upload_one_part( self, filename, bucket, key, upload_id, part_size, extra_args, callback, part_number, ): open_chunk_reader = self._os.open_file_chunk_reader with open_chunk_reader( filename, part_size * (part_number - 1), part_size, callback ) as body: response = self._client.upload_part( Bucket=bucket, Key=key, UploadId=upload_id, PartNumber=part_number, Body=body, **extra_args, ) etag = response['ETag'] return {'ETag': etag, 'PartNumber': part_number} class ShutdownQueue(queue.Queue): """A queue implementation that can be shutdown. Shutting down a queue means that this class adds a trigger_shutdown method that will trigger all subsequent calls to put() to fail with a ``QueueShutdownError``. It purposefully deviates from queue.Queue, and is *not* meant to be a drop in replacement for ``queue.Queue``. """ def _init(self, maxsize): self._shutdown = False self._shutdown_lock = threading.Lock() # queue.Queue is an old style class so we don't use super(). return queue.Queue._init(self, maxsize) def trigger_shutdown(self): with self._shutdown_lock: self._shutdown = True logger.debug("The IO queue is now shutdown.") def put(self, item): # Note: this is not sufficient, it's still possible to deadlock! # Need to hook into the condition vars used by this class. with self._shutdown_lock: if self._shutdown: raise QueueShutdownError( "Cannot put item to queue when " "queue has been shutdown." ) return queue.Queue.put(self, item) class MultipartDownloader: def __init__( self, client, config, osutil, executor_cls=concurrent.futures.ThreadPoolExecutor, ): self._client = client self._config = config self._os = osutil self._executor_cls = executor_cls self._ioqueue = ShutdownQueue(self._config.max_io_queue) def download_file( self, bucket, key, filename, object_size, extra_args, callback=None ): with self._executor_cls(max_workers=2) as controller: # 1 thread for the future that manages the uploading of files # 1 thread for the future that manages IO writes. download_parts_handler = functools.partial( self._download_file_as_future, bucket, key, filename, object_size, callback, ) parts_future = controller.submit(download_parts_handler) io_writes_handler = functools.partial( self._perform_io_writes, filename ) io_future = controller.submit(io_writes_handler) results = concurrent.futures.wait( [parts_future, io_future], return_when=concurrent.futures.FIRST_EXCEPTION, ) self._process_future_results(results) def _process_future_results(self, futures): finished, unfinished = futures for future in finished: future.result() def _download_file_as_future( self, bucket, key, filename, object_size, callback ): part_size = self._config.multipart_chunksize num_parts = int(math.ceil(object_size / float(part_size))) max_workers = self._config.max_concurrency download_partial = functools.partial( self._download_range, bucket, key, filename, part_size, num_parts, callback, ) try: with self._executor_cls(max_workers=max_workers) as executor: list(executor.map(download_partial, range(num_parts))) finally: self._ioqueue.put(SHUTDOWN_SENTINEL) def _calculate_range_param(self, part_size, part_index, num_parts): start_range = part_index * part_size if part_index == num_parts - 1: end_range = '' else: end_range = start_range + part_size - 1 range_param = f'bytes={start_range}-{end_range}' return range_param def _download_range( self, bucket, key, filename, part_size, num_parts, callback, part_index ): try: range_param = self._calculate_range_param( part_size, part_index, num_parts ) max_attempts = self._config.num_download_attempts last_exception = None for i in range(max_attempts): try: logger.debug("Making get_object call.") response = self._client.get_object( Bucket=bucket, Key=key, Range=range_param ) streaming_body = StreamReaderProgress( response['Body'], callback ) buffer_size = 1024 * 16 current_index = part_size * part_index for chunk in iter( lambda: streaming_body.read(buffer_size), b'' ): self._ioqueue.put((current_index, chunk)) current_index += len(chunk) return except ( socket.timeout, OSError, ReadTimeoutError, IncompleteReadError, ) as e: logger.debug( "Retrying exception caught (%s), " "retrying request, (attempt %s / %s)", e, i, max_attempts, exc_info=True, ) last_exception = e continue raise RetriesExceededError(last_exception) finally: logger.debug("EXITING _download_range for part: %s", part_index) def _perform_io_writes(self, filename): with self._os.open(filename, 'wb') as f: while True: task = self._ioqueue.get() if task is SHUTDOWN_SENTINEL: logger.debug( "Shutdown sentinel received in IO handler, " "shutting down IO handler." ) return else: try: offset, data = task f.seek(offset) f.write(data) except Exception as e: logger.debug( "Caught exception in IO thread: %s", e, exc_info=True, ) self._ioqueue.trigger_shutdown() raise class TransferConfig: def __init__( self, multipart_threshold=8 * MB, max_concurrency=10, multipart_chunksize=8 * MB, num_download_attempts=5, max_io_queue=100, ): self.multipart_threshold = multipart_threshold self.max_concurrency = max_concurrency self.multipart_chunksize = multipart_chunksize self.num_download_attempts = num_download_attempts self.max_io_queue = max_io_queue class S3Transfer: ALLOWED_DOWNLOAD_ARGS = [ 'VersionId', 'SSECustomerAlgorithm', 'SSECustomerKey', 'SSECustomerKeyMD5', 'RequestPayer', ] ALLOWED_UPLOAD_ARGS = [ 'ACL', 'CacheControl', 'ContentDisposition', 'ContentEncoding', 'ContentLanguage', 'ContentType', 'Expires', 'GrantFullControl', 'GrantRead', 'GrantReadACP', 'GrantWriteACL', 'Metadata', 'RequestPayer', 'ServerSideEncryption', 'StorageClass', 'SSECustomerAlgorithm', 'SSECustomerKey', 'SSECustomerKeyMD5', 'SSEKMSKeyId', 'SSEKMSEncryptionContext', 'Tagging', ] def __init__(self, client, config=None, osutil=None): self._client = client if config is None: config = TransferConfig() self._config = config if osutil is None: osutil = OSUtils() self._osutil = osutil def upload_file( self, filename, bucket, key, callback=None, extra_args=None ): """Upload a file to an S3 object. Variants have also been injected into S3 client, Bucket and Object. You don't have to use S3Transfer.upload_file() directly. """ if extra_args is None: extra_args = {} self._validate_all_known_args(extra_args, self.ALLOWED_UPLOAD_ARGS) events = self._client.meta.events events.register_first( 'request-created.s3', disable_upload_callbacks, unique_id='s3upload-callback-disable', ) events.register_last( 'request-created.s3', enable_upload_callbacks, unique_id='s3upload-callback-enable', ) if ( self._osutil.get_file_size(filename) >= self._config.multipart_threshold ): self._multipart_upload(filename, bucket, key, callback, extra_args) else: self._put_object(filename, bucket, key, callback, extra_args) def _put_object(self, filename, bucket, key, callback, extra_args): # We're using open_file_chunk_reader so we can take advantage of the # progress callback functionality. open_chunk_reader = self._osutil.open_file_chunk_reader with open_chunk_reader( filename, 0, self._osutil.get_file_size(filename), callback=callback, ) as body: self._client.put_object( Bucket=bucket, Key=key, Body=body, **extra_args ) def download_file( self, bucket, key, filename, extra_args=None, callback=None ): """Download an S3 object to a file. Variants have also been injected into S3 client, Bucket and Object. You don't have to use S3Transfer.download_file() directly. """ # This method will issue a ``head_object`` request to determine # the size of the S3 object. This is used to determine if the # object is downloaded in parallel. if extra_args is None: extra_args = {} self._validate_all_known_args(extra_args, self.ALLOWED_DOWNLOAD_ARGS) object_size = self._object_size(bucket, key, extra_args) temp_filename = filename + os.extsep + random_file_extension() try: self._download_file( bucket, key, temp_filename, object_size, extra_args, callback ) except Exception: logger.debug( "Exception caught in download_file, removing partial " "file: %s", temp_filename, exc_info=True, ) self._osutil.remove_file(temp_filename) raise else: self._osutil.rename_file(temp_filename, filename) def _download_file( self, bucket, key, filename, object_size, extra_args, callback ): if object_size >= self._config.multipart_threshold: self._ranged_download( bucket, key, filename, object_size, extra_args, callback ) else: self._get_object(bucket, key, filename, extra_args, callback) def _validate_all_known_args(self, actual, allowed): for kwarg in actual: if kwarg not in allowed: raise ValueError( "Invalid extra_args key '%s', " "must be one of: %s" % (kwarg, ', '.join(allowed)) ) def _ranged_download( self, bucket, key, filename, object_size, extra_args, callback ): downloader = MultipartDownloader( self._client, self._config, self._osutil ) downloader.download_file( bucket, key, filename, object_size, extra_args, callback ) def _get_object(self, bucket, key, filename, extra_args, callback): # precondition: num_download_attempts > 0 max_attempts = self._config.num_download_attempts last_exception = None for i in range(max_attempts): try: return self._do_get_object( bucket, key, filename, extra_args, callback ) except ( socket.timeout, OSError, ReadTimeoutError, IncompleteReadError, ) as e: # TODO: we need a way to reset the callback if the # download failed. logger.debug( "Retrying exception caught (%s), " "retrying request, (attempt %s / %s)", e, i, max_attempts, exc_info=True, ) last_exception = e continue raise RetriesExceededError(last_exception) def _do_get_object(self, bucket, key, filename, extra_args, callback): response = self._client.get_object( Bucket=bucket, Key=key, **extra_args ) streaming_body = StreamReaderProgress(response['Body'], callback) with self._osutil.open(filename, 'wb') as f: for chunk in iter(lambda: streaming_body.read(8192), b''): f.write(chunk) def _object_size(self, bucket, key, extra_args): return self._client.head_object(Bucket=bucket, Key=key, **extra_args)[ 'ContentLength' ] def _multipart_upload(self, filename, bucket, key, callback, extra_args): uploader = MultipartUploader(self._client, self._config, self._osutil) uploader.upload_file(filename, bucket, key, callback, extra_args)