You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

323 lines
13 KiB

4 years ago
  1. # Copyright 2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License"). You
  4. # may not use this file except in compliance with the License. A copy of
  5. # the License is located at
  6. #
  7. # http://aws.amazon.com/apache2.0/
  8. #
  9. # or in the "license" file accompanying this file. This file is
  10. # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
  11. # ANY KIND, either express or implied. See the License for the specific
  12. # language governing permissions and limitations under the License.
  13. import copy
  14. import math
  15. from s3transfer.tasks import Task
  16. from s3transfer.tasks import SubmissionTask
  17. from s3transfer.tasks import CreateMultipartUploadTask
  18. from s3transfer.tasks import CompleteMultipartUploadTask
  19. from s3transfer.utils import get_callbacks
  20. from s3transfer.utils import calculate_range_parameter
  21. from s3transfer.utils import get_filtered_dict
  22. from s3transfer.utils import ChunksizeAdjuster
  23. class CopySubmissionTask(SubmissionTask):
  24. """Task for submitting tasks to execute a copy"""
  25. EXTRA_ARGS_TO_HEAD_ARGS_MAPPING = {
  26. 'CopySourceIfMatch': 'IfMatch',
  27. 'CopySourceIfModifiedSince': 'IfModifiedSince',
  28. 'CopySourceIfNoneMatch': 'IfNoneMatch',
  29. 'CopySourceIfUnmodifiedSince': 'IfUnmodifiedSince',
  30. 'CopySourceSSECustomerKey': 'SSECustomerKey',
  31. 'CopySourceSSECustomerAlgorithm': 'SSECustomerAlgorithm',
  32. 'CopySourceSSECustomerKeyMD5': 'SSECustomerKeyMD5',
  33. 'RequestPayer': 'RequestPayer'
  34. }
  35. UPLOAD_PART_COPY_ARGS = [
  36. 'CopySourceIfMatch',
  37. 'CopySourceIfModifiedSince',
  38. 'CopySourceIfNoneMatch',
  39. 'CopySourceIfUnmodifiedSince',
  40. 'CopySourceSSECustomerKey',
  41. 'CopySourceSSECustomerAlgorithm',
  42. 'CopySourceSSECustomerKeyMD5',
  43. 'SSECustomerKey',
  44. 'SSECustomerAlgorithm',
  45. 'SSECustomerKeyMD5',
  46. 'RequestPayer',
  47. ]
  48. CREATE_MULTIPART_ARGS_BLACKLIST = [
  49. 'CopySourceIfMatch',
  50. 'CopySourceIfModifiedSince',
  51. 'CopySourceIfNoneMatch',
  52. 'CopySourceIfUnmodifiedSince',
  53. 'CopySourceSSECustomerKey',
  54. 'CopySourceSSECustomerAlgorithm',
  55. 'CopySourceSSECustomerKeyMD5',
  56. 'MetadataDirective'
  57. ]
  58. COMPLETE_MULTIPART_ARGS = [
  59. 'RequestPayer'
  60. ]
  61. def _submit(self, client, config, osutil, request_executor,
  62. transfer_future):
  63. """
  64. :param client: The client associated with the transfer manager
  65. :type config: s3transfer.manager.TransferConfig
  66. :param config: The transfer config associated with the transfer
  67. manager
  68. :type osutil: s3transfer.utils.OSUtil
  69. :param osutil: The os utility associated to the transfer manager
  70. :type request_executor: s3transfer.futures.BoundedExecutor
  71. :param request_executor: The request executor associated with the
  72. transfer manager
  73. :type transfer_future: s3transfer.futures.TransferFuture
  74. :param transfer_future: The transfer future associated with the
  75. transfer request that tasks are being submitted for
  76. """
  77. # Determine the size if it was not provided
  78. if transfer_future.meta.size is None:
  79. # If a size was not provided figure out the size for the
  80. # user. Note that we will only use the client provided to
  81. # the TransferManager. If the object is outside of the region
  82. # of the client, they may have to provide the file size themselves
  83. # with a completely new client.
  84. call_args = transfer_future.meta.call_args
  85. head_object_request = \
  86. self._get_head_object_request_from_copy_source(
  87. call_args.copy_source)
  88. extra_args = call_args.extra_args
  89. # Map any values that may be used in the head object that is
  90. # used in the copy object
  91. for param, value in extra_args.items():
  92. if param in self.EXTRA_ARGS_TO_HEAD_ARGS_MAPPING:
  93. head_object_request[
  94. self.EXTRA_ARGS_TO_HEAD_ARGS_MAPPING[param]] = value
  95. response = call_args.source_client.head_object(
  96. **head_object_request)
  97. transfer_future.meta.provide_transfer_size(
  98. response['ContentLength'])
  99. # If it is greater than threshold do a multipart copy, otherwise
  100. # do a regular copy object.
  101. if transfer_future.meta.size < config.multipart_threshold:
  102. self._submit_copy_request(
  103. client, config, osutil, request_executor, transfer_future)
  104. else:
  105. self._submit_multipart_request(
  106. client, config, osutil, request_executor, transfer_future)
  107. def _submit_copy_request(self, client, config, osutil, request_executor,
  108. transfer_future):
  109. call_args = transfer_future.meta.call_args
  110. # Get the needed progress callbacks for the task
  111. progress_callbacks = get_callbacks(transfer_future, 'progress')
  112. # Submit the request of a single copy.
  113. self._transfer_coordinator.submit(
  114. request_executor,
  115. CopyObjectTask(
  116. transfer_coordinator=self._transfer_coordinator,
  117. main_kwargs={
  118. 'client': client,
  119. 'copy_source': call_args.copy_source,
  120. 'bucket': call_args.bucket,
  121. 'key': call_args.key,
  122. 'extra_args': call_args.extra_args,
  123. 'callbacks': progress_callbacks,
  124. 'size': transfer_future.meta.size
  125. },
  126. is_final=True
  127. )
  128. )
  129. def _submit_multipart_request(self, client, config, osutil,
  130. request_executor, transfer_future):
  131. call_args = transfer_future.meta.call_args
  132. # Submit the request to create a multipart upload and make sure it
  133. # does not include any of the arguments used for copy part.
  134. create_multipart_extra_args = {}
  135. for param, val in call_args.extra_args.items():
  136. if param not in self.CREATE_MULTIPART_ARGS_BLACKLIST:
  137. create_multipart_extra_args[param] = val
  138. create_multipart_future = self._transfer_coordinator.submit(
  139. request_executor,
  140. CreateMultipartUploadTask(
  141. transfer_coordinator=self._transfer_coordinator,
  142. main_kwargs={
  143. 'client': client,
  144. 'bucket': call_args.bucket,
  145. 'key': call_args.key,
  146. 'extra_args': create_multipart_extra_args,
  147. }
  148. )
  149. )
  150. # Determine how many parts are needed based on filesize and
  151. # desired chunksize.
  152. part_size = config.multipart_chunksize
  153. adjuster = ChunksizeAdjuster()
  154. part_size = adjuster.adjust_chunksize(
  155. part_size, transfer_future.meta.size)
  156. num_parts = int(
  157. math.ceil(transfer_future.meta.size / float(part_size)))
  158. # Submit requests to upload the parts of the file.
  159. part_futures = []
  160. progress_callbacks = get_callbacks(transfer_future, 'progress')
  161. for part_number in range(1, num_parts + 1):
  162. extra_part_args = self._extra_upload_part_args(
  163. call_args.extra_args)
  164. # The part number for upload part starts at 1 while the
  165. # range parameter starts at zero, so just subtract 1 off of
  166. # the part number
  167. extra_part_args['CopySourceRange'] = calculate_range_parameter(
  168. part_size, part_number-1, num_parts, transfer_future.meta.size)
  169. # Get the size of the part copy as well for the progress
  170. # callbacks.
  171. size = self._get_transfer_size(
  172. part_size, part_number-1, num_parts, transfer_future.meta.size
  173. )
  174. part_futures.append(
  175. self._transfer_coordinator.submit(
  176. request_executor,
  177. CopyPartTask(
  178. transfer_coordinator=self._transfer_coordinator,
  179. main_kwargs={
  180. 'client': client,
  181. 'copy_source': call_args.copy_source,
  182. 'bucket': call_args.bucket,
  183. 'key': call_args.key,
  184. 'part_number': part_number,
  185. 'extra_args': extra_part_args,
  186. 'callbacks': progress_callbacks,
  187. 'size': size
  188. },
  189. pending_main_kwargs={
  190. 'upload_id': create_multipart_future
  191. }
  192. )
  193. )
  194. )
  195. complete_multipart_extra_args = self._extra_complete_multipart_args(
  196. call_args.extra_args)
  197. # Submit the request to complete the multipart upload.
  198. self._transfer_coordinator.submit(
  199. request_executor,
  200. CompleteMultipartUploadTask(
  201. transfer_coordinator=self._transfer_coordinator,
  202. main_kwargs={
  203. 'client': client,
  204. 'bucket': call_args.bucket,
  205. 'key': call_args.key,
  206. 'extra_args': complete_multipart_extra_args,
  207. },
  208. pending_main_kwargs={
  209. 'upload_id': create_multipart_future,
  210. 'parts': part_futures
  211. },
  212. is_final=True
  213. )
  214. )
  215. def _get_head_object_request_from_copy_source(self, copy_source):
  216. if isinstance(copy_source, dict):
  217. return copy.copy(copy_source)
  218. else:
  219. raise TypeError(
  220. 'Expecting dictionary formatted: '
  221. '{"Bucket": bucket_name, "Key": key} '
  222. 'but got %s or type %s.'
  223. % (copy_source, type(copy_source))
  224. )
  225. def _extra_upload_part_args(self, extra_args):
  226. # Only the args in COPY_PART_ARGS actually need to be passed
  227. # onto the upload_part_copy calls.
  228. return get_filtered_dict(extra_args, self.UPLOAD_PART_COPY_ARGS)
  229. def _extra_complete_multipart_args(self, extra_args):
  230. return get_filtered_dict(extra_args, self.COMPLETE_MULTIPART_ARGS)
  231. def _get_transfer_size(self, part_size, part_index, num_parts,
  232. total_transfer_size):
  233. if part_index == num_parts - 1:
  234. # The last part may be different in size then the rest of the
  235. # parts.
  236. return total_transfer_size - (part_index * part_size)
  237. return part_size
  238. class CopyObjectTask(Task):
  239. """Task to do a nonmultipart copy"""
  240. def _main(self, client, copy_source, bucket, key, extra_args, callbacks,
  241. size):
  242. """
  243. :param client: The client to use when calling PutObject
  244. :param copy_source: The CopySource parameter to use
  245. :param bucket: The name of the bucket to copy to
  246. :param key: The name of the key to copy to
  247. :param extra_args: A dictionary of any extra arguments that may be
  248. used in the upload.
  249. :param callbacks: List of callbacks to call after copy
  250. :param size: The size of the transfer. This value is passed into
  251. the callbacks
  252. """
  253. client.copy_object(
  254. CopySource=copy_source, Bucket=bucket, Key=key, **extra_args)
  255. for callback in callbacks:
  256. callback(bytes_transferred=size)
  257. class CopyPartTask(Task):
  258. """Task to upload a part in a multipart copy"""
  259. def _main(self, client, copy_source, bucket, key, upload_id, part_number,
  260. extra_args, callbacks, size):
  261. """
  262. :param client: The client to use when calling PutObject
  263. :param copy_source: The CopySource parameter to use
  264. :param bucket: The name of the bucket to upload to
  265. :param key: The name of the key to upload to
  266. :param upload_id: The id of the upload
  267. :param part_number: The number representing the part of the multipart
  268. upload
  269. :param extra_args: A dictionary of any extra arguments that may be
  270. used in the upload.
  271. :param callbacks: List of callbacks to call after copy part
  272. :param size: The size of the transfer. This value is passed into
  273. the callbacks
  274. :rtype: dict
  275. :returns: A dictionary representing a part::
  276. {'Etag': etag_value, 'PartNumber': part_number}
  277. This value can be appended to a list to be used to complete
  278. the multipart upload.
  279. """
  280. response = client.upload_part_copy(
  281. CopySource=copy_source, Bucket=bucket, Key=key,
  282. UploadId=upload_id, PartNumber=part_number, **extra_args)
  283. for callback in callbacks:
  284. callback(bytes_transferred=size)
  285. etag = response['CopyPartResult']['ETag']
  286. return {'ETag': etag, 'PartNumber': part_number}