You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
99 lines
4.3 KiB
Python
99 lines
4.3 KiB
Python
#
|
|
# Copyright (c) Contributors to the Open 3D Engine Project.
|
|
# For complete copyright and license terms please see the LICENSE at the root of this distribution.
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0 OR MIT
|
|
#
|
|
#
|
|
|
|
'''
|
|
Usage:
|
|
Use EC2 role to download files to %WORKSPACE% folder from bucket bucket_name:
|
|
python download_from_s3.py --base_dir %WORKSPACE% --files_to_download "file1,file2" --bucket bucket_name
|
|
|
|
Use profile to download files to %WORKSPACE% folder from bucket bucket_name:
|
|
python download_from_s3.py --base_dir %WORKSPACE% --profile profile --files_to_download "file1,file2" --bucket bucket_name
|
|
'''
|
|
|
|
|
|
import os
|
|
import json
|
|
import boto3
|
|
from optparse import OptionParser
|
|
from util import error
|
|
|
|
|
|
def parse_args():
|
|
parser = OptionParser()
|
|
parser.add_option("--base_dir", dest="base_dir", default=os.getcwd(), help="Base directory to download files, If not given, then current directory is used.")
|
|
parser.add_option("--files_to_download", dest="files_to_download", default=None, help="Files to download, separated by comma.")
|
|
parser.add_option("--profile", dest="profile", default=None, help="The name of a profile to use. If not given, then the default profile is used.")
|
|
parser.add_option("--bucket", dest="bucket", default=None, help="S3 bucket the files are downloaded from.")
|
|
parser.add_option("--key_prefix", dest="key_prefix", default='', help="Object key prefix.")
|
|
'''
|
|
ExtraArgs used to call s3.download_file(), should be in json format. extra_args key must be one of: ACL, CacheControl, ContentDisposition, ContentEncoding, ContentLanguage, ContentType, Expires,
|
|
GrantFullControl, GrantRead, GrantReadACP, GrantWriteACP, Metadata, RequestPayer, ServerSideEncryption, StorageClass,
|
|
SSECustomerAlgorithm, SSECustomerKey, SSECustomerKeyMD5, SSEKMSKeyId, WebsiteRedirectLocation
|
|
'''
|
|
parser.add_option("--extra_args", dest="extra_args", default=None, help="Additional parameters used to download file.")
|
|
parser.add_option("--max_retry", dest="max_retry", default=1, help="Maximum retry times to download file.")
|
|
(options, args) = parser.parse_args()
|
|
if not os.path.isdir(options.base_dir):
|
|
error('{} is not a valid directory'.format(options.base_dir))
|
|
if not options.files_to_download:
|
|
error('Use --files_to_download to specify files to download, separated by comma.')
|
|
if not options.bucket:
|
|
error('Use --bucket to specify bucket that the files are downloaded from.')
|
|
return options
|
|
|
|
|
|
def get_client(service_name, profile_name=None):
|
|
session = boto3.session.Session(profile_name=profile_name)
|
|
client = session.client(service_name)
|
|
return client
|
|
|
|
|
|
def s3_download_file(client, base_dir, file, bucket, key_prefix=None, extra_args=None, max_retry=1):
|
|
print('Downloading file {} from bucket {}.'.format(file, bucket))
|
|
key = file if key_prefix is None else '{}/{}'.format(key_prefix, file)
|
|
for x in range(max_retry):
|
|
try:
|
|
client.download_file(
|
|
bucket, key, os.path.join(base_dir, file),
|
|
ExtraArgs=extra_args
|
|
)
|
|
print('Download succeeded')
|
|
return True
|
|
except:
|
|
print('Retrying download...')
|
|
print('Download failed')
|
|
return False
|
|
|
|
|
|
def download_files(base_dir, files_to_download, bucket, key_prefix=None, profile=None, extra_args=None, max_retry=1):
|
|
client = get_client('s3', profile)
|
|
files_to_download = files_to_download.split(',')
|
|
extra_args = json.loads(extra_args) if extra_args else None
|
|
|
|
print('Downloading {} files from bucket {}.'.format(len(files_to_download), bucket))
|
|
failure = []
|
|
success = []
|
|
for file in files_to_download:
|
|
if not s3_download_file(client, base_dir, file, bucket, key_prefix, extra_args, max_retry):
|
|
failure.append(file)
|
|
else:
|
|
success.append(file)
|
|
print('{} files are downloaded successfully:'.format(len(success)))
|
|
print('\n'.join(success))
|
|
print('{} files failed to download:'.format(len(failure)))
|
|
print('\n'.join(failure))
|
|
# Exit with error code 1 if any file is failed to download
|
|
if len(failure) > 0:
|
|
return False
|
|
return True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
options = parse_args()
|
|
download_files(options.base_dir, options.files_to_download, options.bucket, options.key_prefix, options.profile, options.extra_args, options.max_retry)
|