Function download_s3

SwissArmyTransformer/sat/resources/download.py:84–121

Download remote_dir into (under) local_dir

(local_dir, remote_uri)

SAT_BUCKET = 'sat'

def download_s3(local_dir, remote_uri):
    '''Download remote_dir into (under) local_dir
    '''
    import boto3
    s3_resource = boto3.resource('s3',
        endpoint_url = f'https://{SAT_ACCOUNT}.r2.cloudflarestorage.com',
        aws_access_key_id = f'{SAT_ACCESS_KEY}',
        aws_secret_access_key = f'{SAT_SECRET_KEY}'
        )
    client = boto3.client('s3',
        endpoint_url = f'https://{SAT_ACCOUNT}.r2.cloudflarestorage.com',
        aws_access_key_id = f'{SAT_ACCESS_KEY}',
        aws_secret_access_key = f'{SAT_SECRET_KEY}',
        verify=False
        )
    bucket = s3_resource.Bucket(SAT_BUCKET) 
    transfer_config = boto3.s3.transfer.TransferConfig(
        use_threads=True,
        multipart_threshold=8*1024*1024,
        max_concurrency=64,
        multipart_chunksize=8*1024*1024,
    )
    # remote_uri is file
    if '.' in os.path.basename(remote_uri):
        bucket.download_file(remote_uri, os.path.join(local_dir, os.path.basename(remote_uri)),Callback=ProgressPercentage(client, SAT_BUCKET, remote_uri), Config=transfer_config)
        return 
    # uri is path
    remote_dir = remote_uri
    key_prefix = remote_dir.split('/')[:-1]
    for obj in bucket.objects.filter(Prefix = remote_dir):
        key_suffix = obj.key[len(key_prefix):] # remote_dir/xxx/xxx.zip
        target_dir = os.path.join(local_dir, os.path.dirname(key_suffix))
        if not os.path.exists(target_dir):
            os.makedirs(target_dir)
        # skip if exists
        if os.path.exists(os.path.join(local_dir, key_suffix)) and os.path.getsize(os.path.join(local_dir, key_suffix)) == obj.size:
            continue
        bucket.download_file(obj.key, os.path.join(local_dir, key_suffix),Callback=ProgressPercentage(client, SAT_BUCKET, obj.key), Config=transfer_config) 


class ProgressPercentage(object):