Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
from s3transfer.futures import BoundedExecutor
from s3transfer.subscribers import BaseSubscriber
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 CallArgs
from s3transfer.utils import FunctionContainer
class TaskFailureException(Exception):
pass
class SuccessTask(Task):
def _main(self, return_value='success', callbacks=None,
failure_cleanups=None):
if callbacks:
for callback in callbacks:
callback()
if failure_cleanups:
for failure_cleanup in failure_cleanups:
self._transfer_coordinator.add_failure_cleanup(failure_cleanup)
return return_value
class FailureTask(Task):
def _main(self, exception=TaskFailureException):
raise exception()
raise_exception(exception)
except:
return sys.exc_info()
class RecordingTransferCoordinator(TransferCoordinator):
def __init__(self):
self.all_transfer_futures_ever_associated = set()
super(RecordingTransferCoordinator, self).__init__()
def add_associated_future(self, future):
self.all_transfer_futures_ever_associated.add(future)
super(RecordingTransferCoordinator, self).add_associated_future(future)
class ReturnFooTask(Task):
def _main(self, **kwargs):
return 'foo'
class SleepTask(Task):
def _main(self, sleep_time, **kwargs):
time.sleep(sleep_time)
class TestTransferFuture(unittest.TestCase):
def setUp(self):
self.meta = TransferMeta()
self.coordinator = TransferCoordinator()
self.future = self._get_transfer_future()
def _get_transfer_future(self, **kwargs):
self._transfer_coordinator.submit(
request_executor,
DeleteObjectTask(
transfer_coordinator=self._transfer_coordinator,
main_kwargs={
'client': client,
'bucket': call_args.bucket,
'key': call_args.key,
'extra_args': call_args.extra_args,
},
is_final=True
)
)
class DeleteObjectTask(Task):
def _main(self, client, bucket, key, extra_args):
"""
:param client: The S3 client to use when calling DeleteObject
:type bucket: str
:param bucket: The name of the bucket.
:type key: str
:param key: The name of the object to delete.
:type extra_args: dict
:param extra_args: Extra arguments to pass to the DeleteObject call.
"""
client.delete_object(Bucket=bucket, Key=key, **extra_args)
# If the value is a list of futures, iterate though the list
# appending on the result from each future.
if isinstance(pending_value, list):
result = []
for future in pending_value:
result.append(future.result())
# Otherwise if the pending_value is a future, just wait for it.
else:
result = pending_value.result()
# Add the retrieved value to the kwargs to be sent to the
# main() call.
kwargs[key] = result
return kwargs
class SubmissionTask(Task):
"""A base class for any submission task
Submission tasks are the top-level task used to submit a series of tasks
to execute a particular transfer.
"""
def _main(self, transfer_future, **kwargs):
"""
:type transfer_future: s3transfer.futures.TransferFuture
:param transfer_future: The transfer future associated with the
transfer request that tasks are being submitted for
:param kwargs: Any additional kwargs that you may want to pass
to the _submit() method
"""
try:
self._transfer_coordinator.set_status_to_queued()
final_task = download_manager.get_final_io_task()
return FunctionContainer(
self._transfer_coordinator.submit, io_executor, final_task)
def _calculate_range_param(self, part_size, part_index, num_parts):
# Used to calculate the Range parameter
start_range = part_index * part_size
if part_index == num_parts - 1:
end_range = ''
else:
end_range = start_range + part_size - 1
range_param = 'bytes=%s-%s' % (start_range, end_range)
return range_param
class GetObjectTask(Task):
def _main(self, client, bucket, key, fileobj, extra_args, callbacks,
max_attempts, download_output_manager, io_chunksize,
start_index=0, bandwidth_limiter=None):
"""Downloads an object and places content into io queue
:param client: The client to use when calling GetObject
:param bucket: The bucket to download from
: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 max_attempts: The number of retries to do when downloading
:param download_output_manager: The download output manager associated
with the current download.
:param io_chunksize: The size of each io chunk to read from the
download stream and queue in the io queue.
class IORenameFileTask(Task):
"""A task to rename a temporary file to its final filename
:param fileobj: The file handle that content was written to.
:param final_filename: The final name of the file to rename to
upon completion of writing the contents.
:param osutil: OS utility
"""
def _main(self, fileobj, final_filename, osutil):
fileobj.close()
osutil.rename_file(fileobj.name, final_filename)
class IOCloseTask(Task):
"""A task to close out a file once the download is complete.
:param fileobj: The fileobj to close.
"""
def _main(self, fileobj):
fileobj.close()
class CompleteDownloadNOOPTask(Task):
"""A NOOP task to serve as an indicator that the download is complete
Note that the default for is_final is set to True because this should
always be the last task.
"""
def __init__(self, transfer_coordinator, main_kwargs=None,
pending_main_kwargs=None, done_callbacks=None,
:returns: The upload id of the multipart upload
"""
# Create the multipart upload.
response = client.create_multipart_upload(
Bucket=bucket, Key=key, **extra_args)
upload_id = response['UploadId']
# Add a cleanup if the multipart upload fails at any point.
self._transfer_coordinator.add_failure_cleanup(
client.abort_multipart_upload, Bucket=bucket, Key=key,
UploadId=upload_id
)
return upload_id
class CompleteMultipartUploadTask(Task):
"""Task to complete a multipart upload"""
def _main(self, client, bucket, key, upload_id, parts, extra_args):
"""
:param client: The client to use when calling CompleteMultipartUpload
: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 parts: A list of parts to use to complete the multipart upload::
[{'Etag': etag_value, 'PartNumber': part_number}, ...]
Each element in the list consists of a return value from
``UploadPartTask.main()``.
:param extra_args: A dictionary of any extra arguments that may be
used in completing the multipart transfer.
"""
# 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(
"""
def _main(self, fileobj, final_filename, osutil):
fileobj.close()
osutil.rename_file(fileobj.name, final_filename)
class IOCloseTask(Task):
"""A task to close out a file once the download is complete.
:param fileobj: The fileobj to close.
"""
def _main(self, fileobj):
fileobj.close()
class CompleteDownloadNOOPTask(Task):
"""A NOOP task to serve as an indicator that the download is complete
Note that the default for is_final is set to True because this should
always be the last task.
"""
def __init__(self, transfer_coordinator, main_kwargs=None,
pending_main_kwargs=None, done_callbacks=None,
is_final=True):
super(CompleteDownloadNOOPTask, self).__init__(
transfer_coordinator=transfer_coordinator,
main_kwargs=main_kwargs,
pending_main_kwargs=pending_main_kwargs,
done_callbacks=done_callbacks,
is_final=is_final
)
download_output_manager.queue_file_io_task(fileobj, chunk, index)
class ImmediatelyWriteIOGetObjectTask(GetObjectTask):
"""GetObjectTask that immediately writes to the provided file object
This is useful for downloads where it is known only one thread is
downloading the object so there is no reason to go through the
overhead of using an IO queue and executor.
"""
def _handle_io(self, download_output_manager, fileobj, chunk, index):
task = download_output_manager.get_io_write_task(fileobj, chunk, index)
task()
class IOWriteTask(Task):
def _main(self, fileobj, data, offset):
"""Pulls off an io queue to write contents to a file
:param fileobj: The file handle to write content to
:param data: The data to write
:param offset: The offset to write the data to.
"""
fileobj.seek(offset)
fileobj.write(data)
class IOStreamingWriteTask(Task):
"""Task for writing data to a non-seekable stream."""
def _main(self, fileobj, data):
"""Write data to a fileobj.