def test_create_transfer_manager(self):
client = object()
config = TransferConfig()
osutil = OSUtils()
with mock.patch('boto3.s3.transfer.TransferManager') as manager:
create_transfer_manager(client, config, osutil)
assert manager.call_args == mock.call(client, config, osutil, None)
def test_create_transfer_manager_with_no_threads(self):
client = object()
config = TransferConfig()
config.use_threads = False
with mock.patch('boto3.s3.transfer.TransferManager') as manager:
create_transfer_manager(client, config)
assert manager.call_args == mock.call(client, config, None,
NonThreadedExecutor)
def __init__(self, s3_client, bucket_name, prefix=None, kms_key_id=None, force_upload=False):
self.s3 = s3_client
self.bucket_name = bucket_name
self.prefix = prefix
self.kms_key_id = kms_key_id or None
self.force_upload = force_upload
self.transfer_manager = transfer.create_transfer_manager(self.s3, transfer.TransferConfig())
self._artifact_metadata = None
def download_fileobj(self, Bucket, Key, Fileobj, ExtraArgs=None,
Callback=None, Config=None):
"""Download an object from S3 to a file-like object.
The file-like object must be in binary mode.
This is a managed transfer which will perform a multipart download in
multiple threads if necessary.
Usage::
import boto3
s3 = boto3.client('s3')
with open('filename', 'wb') as data:
s3.download_fileobj('mybucket', 'mykey', data)
:type Fileobj: a file-like object
:param Fileobj: A file-like object to download into. At a minimum, it must
implement the `write` method and must accept bytes.
:type Bucket: str
:param Bucket: The name of the bucket to download from.
:type Key: str
:param Key: The name of the key to download from.
:type ExtraArgs: dict
:param ExtraArgs: Extra arguments that may be passed to the
client operation.
:type Callback: method
:param Callback: A method which takes a number of bytes transferred to
be periodically called during the download.
:type Config: boto3.s3.transfer.TransferConfig
:param Config: The transfer configuration to be used when performing the
download.
"""
if not hasattr(Fileobj, 'write'):
raise ValueError('Fileobj must implement write')
subscribers = None
if Callback is not None:
subscribers = [ProgressCallbackInvoker(Callback)]
config = Config
if config is None:
config = TransferConfig()
with create_transfer_manager(self, config) as manager:
future = manager.download(
bucket=Bucket, key=Key, fileobj=Fileobj,
extra_args=ExtraArgs, subscribers=subscribers)
return future.result()
def __init__(
self,
s3_client: Any,
bucket_name: str,
prefix: Optional[str] = None,
kms_key_id: Optional[str] = None,
force_upload: bool = False,
no_progressbar: bool = False,
):
self.s3 = s3_client
self.bucket_name = bucket_name
self.prefix = prefix
self.kms_key_id = kms_key_id or None
self.force_upload = force_upload
self.no_progressbar = no_progressbar
self.transfer_manager = transfer.create_transfer_manager(
self.s3, transfer.TransferConfig())
self._artifact_metadata = None
def _upload_file(source_path: str, bucket: str, target_path: str,
storage_class: StorageClass) -> bool:
logger.debug(f'Uploading {source_path} to {target_path}...')
botocore_config = botocore.config.Config(max_pool_connections=CONNECTIONS)
s3 = boto3.client('s3', config=botocore_config)
transfer_config = s3transfer.TransferConfig(
use_threads=True,
max_concurrency=CONNECTIONS,
)
s3t = s3transfer.create_transfer_manager(s3, transfer_config)
total_size = os.path.getsize(source_path)
progress = tqdm.tqdm(
desc='upload',
total=total_size,
unit='B',
unit_scale=1,
position=0,
bar_format='{desc:<10}{percentage:3.0f}%|{bar:10}{r_bar}')
future = s3t.upload(
source_path,
bucket,
target_path,
extra_args={
'ServerSideEncryption': 'AES256',
'StorageClass': storage_class.name
},
subscribers=[s3transfer.ProgressCallbackInvoker(progress.update)])
try:
future.result()
except ClientError as e:
logger.error(f'Upload failed for {source_path} to {target_path}', e)
return False
finally:
s3t.shutdown()
progress.close()
return True
def copy(self, CopySource, Bucket, Key, ExtraArgs=None, Callback=None,
SourceClient=None, Config=None):
"""Copy an object from one S3 location to another.
This is a managed transfer which will perform a multipart copy in
multiple threads if necessary.
Usage::
import boto3
s3 = boto3.resource('s3')
copy_source = {
'Bucket': 'mybucket',
'Key': 'mykey'
}
s3.meta.client.copy(copy_source, 'otherbucket', 'otherkey')
:type CopySource: dict
:param CopySource: The name of the source bucket, key name of the
source object, and optional version ID of the source object. The
dictionary format is:
``{'Bucket': 'bucket', 'Key': 'key', 'VersionId': 'id'}``. Note
that the ``VersionId`` key is optional and may be omitted.
:type Bucket: str
:param Bucket: The name of the bucket to copy to
:type Key: str
:param Key: The name of the key to copy to
:type ExtraArgs: dict
:param ExtraArgs: Extra arguments that may be passed to the
client operation
:type Callback: method
:param Callback: A method which takes a number of bytes transferred to
be periodically called during the copy.
:type SourceClient: botocore or boto3 Client
:param SourceClient: The client to be used for operation that
may happen at the source object. For example, this client is
used for the head_object that determines the size of the copy.
If no client is provided, the current client is used as the client
for the source object.
:type Config: boto3.s3.transfer.TransferConfig
:param Config: The transfer configuration to be used when performing the
copy.
"""
subscribers = None
if Callback is not None:
subscribers = [ProgressCallbackInvoker(Callback)]
config = Config
if config is None:
config = TransferConfig()
with create_transfer_manager(self, config) as manager:
future = manager.copy(
copy_source=CopySource, bucket=Bucket, key=Key,
extra_args=ExtraArgs, subscribers=subscribers,
source_client=SourceClient)
return future.result()
def create_s3t():
s3_client = create_s3_client()
transfer_config = s3transfer.TransferConfig(
use_threads=True, max_concurrency=s3_client.meta.config.max_pool_connections
)
return s3transfer.create_transfer_manager(s3_client, config=transfer_config)
import json
import os
import pathlib
from threading import Lock, get_ident
import boto3
from boto3.s3.transfer import S3Transfer, TransferConfig, create_transfer_manager
from s3transfer.subscribers import BaseSubscriber
from tqdm.autonotebook import tqdm
from .util import HeliumException, split_path
HELIUM_METADATA = 'helium'
s3_client = boto3.client('s3')
s3_manager = create_transfer_manager(s3_client, TransferConfig())
class SizeCallback(BaseSubscriber):
def __init__(self, size):
self.size = size
def on_queued(self, future, **kwargs):
future.meta.provide_transfer_size(self.size)
class ProgressCallback(BaseSubscriber):
def __init__(self, progress):
self._progress = progress
self._lock = Lock()
def upload_fileobj(self,
Fileobj,
Bucket,
Key,
ExtraArgs=None,
Callback=None,
Config=None):
"""Upload a file-like object to S3.
The file-like object must be in binary mode.
This is a managed transfer which will perform a multipart upload in
multiple threads if necessary.
Usage::
import boto3
s3 = boto3.client('s3')
with open('filename', 'rb') as data:
s3.upload_fileobj(data, 'mybucket', 'mykey')
:type Fileobj: a file-like object
:param Fileobj: A file-like object to upload. At a minimum, it must
implement the `read` method, and must return bytes.
:type Bucket: str
:param Bucket: The name of the bucket to upload to.
:type Key: str
:param Key: The name of the key to upload to.
:type ExtraArgs: dict
:param ExtraArgs: Extra arguments that may be passed to the
client operation. For allowed upload arguments see
boto3.s3.transfer.S3Transfer.ALLOWED_UPLOAD_ARGS.
:type Callback: function
:param Callback: A method which takes a number of bytes transferred to
be periodically called during the upload.
:type Config: boto3.s3.transfer.TransferConfig
:param Config: The transfer configuration to be used when performing the
upload.
"""
if not hasattr(Fileobj, 'read'):
raise ValueError('Fileobj must implement read')
subscribers = None
if Callback is not None:
subscribers = [ProgressCallbackInvoker(Callback)]
config = Config
if config is None:
config = TransferConfig()
with create_transfer_manager(self, config) as manager:
future = manager.upload(
fileobj=Fileobj,
bucket=Bucket,
key=Key,
extra_args=ExtraArgs,
subscribers=subscribers,
)
return future.result()
