323 lines
13 KiB
Python
323 lines
13 KiB
Python
# 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.
|
|
import copy
|
|
import math
|
|
|
|
from s3transfer.tasks import Task
|
|
from s3transfer.tasks import SubmissionTask
|
|
from s3transfer.tasks import CreateMultipartUploadTask
|
|
from s3transfer.tasks import CompleteMultipartUploadTask
|
|
from s3transfer.utils import get_callbacks
|
|
from s3transfer.utils import calculate_range_parameter
|
|
from s3transfer.utils import get_filtered_dict
|
|
from s3transfer.utils import ChunksizeAdjuster
|
|
|
|
|
|
class CopySubmissionTask(SubmissionTask):
|
|
"""Task for submitting tasks to execute a copy"""
|
|
|
|
EXTRA_ARGS_TO_HEAD_ARGS_MAPPING = {
|
|
'CopySourceIfMatch': 'IfMatch',
|
|
'CopySourceIfModifiedSince': 'IfModifiedSince',
|
|
'CopySourceIfNoneMatch': 'IfNoneMatch',
|
|
'CopySourceIfUnmodifiedSince': 'IfUnmodifiedSince',
|
|
'CopySourceSSECustomerKey': 'SSECustomerKey',
|
|
'CopySourceSSECustomerAlgorithm': 'SSECustomerAlgorithm',
|
|
'CopySourceSSECustomerKeyMD5': 'SSECustomerKeyMD5',
|
|
'RequestPayer': 'RequestPayer'
|
|
}
|
|
|
|
UPLOAD_PART_COPY_ARGS = [
|
|
'CopySourceIfMatch',
|
|
'CopySourceIfModifiedSince',
|
|
'CopySourceIfNoneMatch',
|
|
'CopySourceIfUnmodifiedSince',
|
|
'CopySourceSSECustomerKey',
|
|
'CopySourceSSECustomerAlgorithm',
|
|
'CopySourceSSECustomerKeyMD5',
|
|
'SSECustomerKey',
|
|
'SSECustomerAlgorithm',
|
|
'SSECustomerKeyMD5',
|
|
'RequestPayer',
|
|
]
|
|
|
|
CREATE_MULTIPART_ARGS_BLACKLIST = [
|
|
'CopySourceIfMatch',
|
|
'CopySourceIfModifiedSince',
|
|
'CopySourceIfNoneMatch',
|
|
'CopySourceIfUnmodifiedSince',
|
|
'CopySourceSSECustomerKey',
|
|
'CopySourceSSECustomerAlgorithm',
|
|
'CopySourceSSECustomerKeyMD5',
|
|
'MetadataDirective'
|
|
]
|
|
|
|
COMPLETE_MULTIPART_ARGS = [
|
|
'RequestPayer'
|
|
]
|
|
|
|
def _submit(self, client, config, osutil, request_executor,
|
|
transfer_future):
|
|
"""
|
|
:param client: The client associated with the transfer manager
|
|
|
|
:type config: s3transfer.manager.TransferConfig
|
|
:param config: The transfer config associated with the transfer
|
|
manager
|
|
|
|
:type osutil: s3transfer.utils.OSUtil
|
|
:param osutil: The os utility associated to the transfer manager
|
|
|
|
:type request_executor: s3transfer.futures.BoundedExecutor
|
|
:param request_executor: The request executor associated with the
|
|
transfer manager
|
|
|
|
:type transfer_future: s3transfer.futures.TransferFuture
|
|
:param transfer_future: The transfer future associated with the
|
|
transfer request that tasks are being submitted for
|
|
"""
|
|
# Determine the size if it was not provided
|
|
if transfer_future.meta.size is None:
|
|
# If a size was not provided figure out the size for the
|
|
# user. Note that we will only use the client provided to
|
|
# the TransferManager. If the object is outside of the region
|
|
# of the client, they may have to provide the file size themselves
|
|
# with a completely new client.
|
|
call_args = transfer_future.meta.call_args
|
|
head_object_request = \
|
|
self._get_head_object_request_from_copy_source(
|
|
call_args.copy_source)
|
|
extra_args = call_args.extra_args
|
|
|
|
# Map any values that may be used in the head object that is
|
|
# used in the copy object
|
|
for param, value in extra_args.items():
|
|
if param in self.EXTRA_ARGS_TO_HEAD_ARGS_MAPPING:
|
|
head_object_request[
|
|
self.EXTRA_ARGS_TO_HEAD_ARGS_MAPPING[param]] = value
|
|
|
|
response = call_args.source_client.head_object(
|
|
**head_object_request)
|
|
transfer_future.meta.provide_transfer_size(
|
|
response['ContentLength'])
|
|
|
|
# If it is greater than threshold do a multipart copy, otherwise
|
|
# do a regular copy object.
|
|
if transfer_future.meta.size < config.multipart_threshold:
|
|
self._submit_copy_request(
|
|
client, config, osutil, request_executor, transfer_future)
|
|
else:
|
|
self._submit_multipart_request(
|
|
client, config, osutil, request_executor, transfer_future)
|
|
|
|
def _submit_copy_request(self, client, config, osutil, request_executor,
|
|
transfer_future):
|
|
call_args = transfer_future.meta.call_args
|
|
|
|
# Get the needed progress callbacks for the task
|
|
progress_callbacks = get_callbacks(transfer_future, 'progress')
|
|
|
|
# Submit the request of a single copy.
|
|
self._transfer_coordinator.submit(
|
|
request_executor,
|
|
CopyObjectTask(
|
|
transfer_coordinator=self._transfer_coordinator,
|
|
main_kwargs={
|
|
'client': client,
|
|
'copy_source': call_args.copy_source,
|
|
'bucket': call_args.bucket,
|
|
'key': call_args.key,
|
|
'extra_args': call_args.extra_args,
|
|
'callbacks': progress_callbacks,
|
|
'size': transfer_future.meta.size
|
|
},
|
|
is_final=True
|
|
)
|
|
)
|
|
|
|
def _submit_multipart_request(self, client, config, osutil,
|
|
request_executor, transfer_future):
|
|
call_args = transfer_future.meta.call_args
|
|
|
|
# Submit the request to create a multipart upload and make sure it
|
|
# does not include any of the arguments used for copy part.
|
|
create_multipart_extra_args = {}
|
|
for param, val in call_args.extra_args.items():
|
|
if param not in self.CREATE_MULTIPART_ARGS_BLACKLIST:
|
|
create_multipart_extra_args[param] = val
|
|
|
|
create_multipart_future = self._transfer_coordinator.submit(
|
|
request_executor,
|
|
CreateMultipartUploadTask(
|
|
transfer_coordinator=self._transfer_coordinator,
|
|
main_kwargs={
|
|
'client': client,
|
|
'bucket': call_args.bucket,
|
|
'key': call_args.key,
|
|
'extra_args': create_multipart_extra_args,
|
|
}
|
|
)
|
|
)
|
|
|
|
# Determine how many parts are needed based on filesize and
|
|
# desired chunksize.
|
|
part_size = config.multipart_chunksize
|
|
adjuster = ChunksizeAdjuster()
|
|
part_size = adjuster.adjust_chunksize(
|
|
part_size, transfer_future.meta.size)
|
|
num_parts = int(
|
|
math.ceil(transfer_future.meta.size / float(part_size)))
|
|
|
|
# Submit requests to upload the parts of the file.
|
|
part_futures = []
|
|
progress_callbacks = get_callbacks(transfer_future, 'progress')
|
|
|
|
for part_number in range(1, num_parts + 1):
|
|
extra_part_args = self._extra_upload_part_args(
|
|
call_args.extra_args)
|
|
# The part number for upload part starts at 1 while the
|
|
# range parameter starts at zero, so just subtract 1 off of
|
|
# the part number
|
|
extra_part_args['CopySourceRange'] = calculate_range_parameter(
|
|
part_size, part_number-1, num_parts, transfer_future.meta.size)
|
|
# Get the size of the part copy as well for the progress
|
|
# callbacks.
|
|
size = self._get_transfer_size(
|
|
part_size, part_number-1, num_parts, transfer_future.meta.size
|
|
)
|
|
part_futures.append(
|
|
self._transfer_coordinator.submit(
|
|
request_executor,
|
|
CopyPartTask(
|
|
transfer_coordinator=self._transfer_coordinator,
|
|
main_kwargs={
|
|
'client': client,
|
|
'copy_source': call_args.copy_source,
|
|
'bucket': call_args.bucket,
|
|
'key': call_args.key,
|
|
'part_number': part_number,
|
|
'extra_args': extra_part_args,
|
|
'callbacks': progress_callbacks,
|
|
'size': size
|
|
},
|
|
pending_main_kwargs={
|
|
'upload_id': create_multipart_future
|
|
}
|
|
)
|
|
)
|
|
)
|
|
|
|
complete_multipart_extra_args = self._extra_complete_multipart_args(
|
|
call_args.extra_args)
|
|
# Submit the request to complete the multipart upload.
|
|
self._transfer_coordinator.submit(
|
|
request_executor,
|
|
CompleteMultipartUploadTask(
|
|
transfer_coordinator=self._transfer_coordinator,
|
|
main_kwargs={
|
|
'client': client,
|
|
'bucket': call_args.bucket,
|
|
'key': call_args.key,
|
|
'extra_args': complete_multipart_extra_args,
|
|
},
|
|
pending_main_kwargs={
|
|
'upload_id': create_multipart_future,
|
|
'parts': part_futures
|
|
},
|
|
is_final=True
|
|
)
|
|
)
|
|
|
|
def _get_head_object_request_from_copy_source(self, copy_source):
|
|
if isinstance(copy_source, dict):
|
|
return copy.copy(copy_source)
|
|
else:
|
|
raise TypeError(
|
|
'Expecting dictionary formatted: '
|
|
'{"Bucket": bucket_name, "Key": key} '
|
|
'but got %s or type %s.'
|
|
% (copy_source, type(copy_source))
|
|
)
|
|
|
|
def _extra_upload_part_args(self, extra_args):
|
|
# Only the args in COPY_PART_ARGS actually need to be passed
|
|
# onto the upload_part_copy calls.
|
|
return get_filtered_dict(extra_args, self.UPLOAD_PART_COPY_ARGS)
|
|
|
|
def _extra_complete_multipart_args(self, extra_args):
|
|
return get_filtered_dict(extra_args, self.COMPLETE_MULTIPART_ARGS)
|
|
|
|
def _get_transfer_size(self, part_size, part_index, num_parts,
|
|
total_transfer_size):
|
|
if part_index == num_parts - 1:
|
|
# The last part may be different in size then the rest of the
|
|
# parts.
|
|
return total_transfer_size - (part_index * part_size)
|
|
return part_size
|
|
|
|
|
|
class CopyObjectTask(Task):
|
|
"""Task to do a nonmultipart copy"""
|
|
def _main(self, client, copy_source, bucket, key, extra_args, callbacks,
|
|
size):
|
|
"""
|
|
:param client: The client to use when calling PutObject
|
|
:param copy_source: The CopySource parameter to use
|
|
:param bucket: The name of the bucket to copy to
|
|
:param key: The name of the key to copy to
|
|
:param extra_args: A dictionary of any extra arguments that may be
|
|
used in the upload.
|
|
:param callbacks: List of callbacks to call after copy
|
|
:param size: The size of the transfer. This value is passed into
|
|
the callbacks
|
|
|
|
"""
|
|
client.copy_object(
|
|
CopySource=copy_source, Bucket=bucket, Key=key, **extra_args)
|
|
for callback in callbacks:
|
|
callback(bytes_transferred=size)
|
|
|
|
|
|
class CopyPartTask(Task):
|
|
"""Task to upload a part in a multipart copy"""
|
|
def _main(self, client, copy_source, bucket, key, upload_id, part_number,
|
|
extra_args, callbacks, size):
|
|
"""
|
|
:param client: The client to use when calling PutObject
|
|
:param copy_source: The CopySource parameter to use
|
|
:param bucket: The name of the bucket to upload to
|
|
:param key: The name of the key to upload to
|
|
:param upload_id: The id of the upload
|
|
:param part_number: The number representing the part of the multipart
|
|
upload
|
|
:param extra_args: A dictionary of any extra arguments that may be
|
|
used in the upload.
|
|
:param callbacks: List of callbacks to call after copy part
|
|
:param size: The size of the transfer. This value is passed into
|
|
the callbacks
|
|
|
|
:rtype: dict
|
|
:returns: A dictionary representing a part::
|
|
|
|
{'Etag': etag_value, 'PartNumber': part_number}
|
|
|
|
This value can be appended to a list to be used to complete
|
|
the multipart upload.
|
|
"""
|
|
response = client.upload_part_copy(
|
|
CopySource=copy_source, Bucket=bucket, Key=key,
|
|
UploadId=upload_id, PartNumber=part_number, **extra_args)
|
|
for callback in callbacks:
|
|
callback(bytes_transferred=size)
|
|
etag = response['CopyPartResult']['ETag']
|
|
return {'ETag': etag, 'PartNumber': part_number}
|