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.
132 lines
5.1 KiB
Python
132 lines
5.1 KiB
Python
"""
|
|
All or portions of this file Copyright (c) Amazon.com, Inc. or its affiliates or
|
|
its licensors.
|
|
|
|
For complete copyright and license terms please see the LICENSE at the root of this
|
|
distribution (the "License"). All use of this software is governed by the License,
|
|
or, if provided, by the license below or the license accompanying this file. Do not
|
|
remove or modify any license notices. This file is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
"""
|
|
import pytest
|
|
pytest.importorskip("boto3")
|
|
import boto3
|
|
import botocore.exceptions
|
|
import logging
|
|
import os
|
|
|
|
import ly_test_tools.environment.file_system as file_system
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class KeyExistsError(Exception):
|
|
pass
|
|
|
|
|
|
class KeyDoesNotExistError(Exception):
|
|
pass
|
|
|
|
|
|
class S3Utils(object):
|
|
"""
|
|
Stores a boto3 S3 client to use for AWS S3 functionalities.
|
|
"""
|
|
DEFAULT_REGION = 'us-west-2'
|
|
|
|
def __init__(self, boto3_session=None):
|
|
# type: (boto3.Session) -> None
|
|
"""
|
|
The boto3 session can be set during init, or a default one will be created.
|
|
:param boto3_session: A boto3 session
|
|
"""
|
|
if boto3_session:
|
|
self._session = boto3_session
|
|
else:
|
|
logger.info("No session provided, using default profile for s3 resource")
|
|
self._session = boto3.session.Session()
|
|
self._s3_resource = self._session.resource('s3')
|
|
|
|
def upload_to_bucket(self, bucket_name, file_path, overwrite=False):
|
|
"""
|
|
Uploads a given file to the given S3 bucket.
|
|
:param bucket_name: Name of the S3 bucket where the file should be uploaded.
|
|
:param file_path: Path to the target file.
|
|
:param overwrite: Overwrite the key if it exists.
|
|
"""
|
|
if not self.bucket_exists_in_s3(bucket_name):
|
|
self._s3_resource.create_bucket(Bucket=bucket_name)
|
|
|
|
s3_bucket = self._s3_resource.Bucket(bucket_name)
|
|
|
|
file_key = os.path.basename(file_path)
|
|
if not overwrite and self.key_exists_in_bucket(bucket_name, file_key):
|
|
raise KeyExistsError("Key '{}' already exists in S3 bucket {}".format(file_key, bucket_name))
|
|
|
|
s3_bucket.upload_file(file_path, file_key)
|
|
logger.info("Uploading {} to S3 bucket {}".format(file_key, bucket_name))
|
|
|
|
def download_from_bucket(self, bucket_name, file_key, destination_dir, file_name=None):
|
|
"""
|
|
Download the given key from the given S3 bucket to the given destination. Logs an error if there is not enough \
|
|
space available for the download.
|
|
:param bucket_name: Name of the S3 bucket containing the desired file.
|
|
:param file_key: Name of the file stored in S3.
|
|
:param destination_dir: Directory where the file should be downloaded to.
|
|
:param file_name: The name of the file you want to save it as. Defaults to the file_key.
|
|
"""
|
|
self.bucket_exists_in_s3(bucket_name)
|
|
|
|
if not self.key_exists_in_bucket(bucket_name, file_key):
|
|
raise KeyDoesNotExistError("Key '{}' does not exist in S3 bucket {}".format(file_key, bucket_name))
|
|
|
|
obj_summary = self._s3_resource.ObjectSummary(bucket_name, file_key)
|
|
required_space = obj_summary.size
|
|
|
|
file_system.check_free_space(destination_dir, required_space, "Insufficient space available for download:")
|
|
|
|
if not os.path.exists(destination_dir):
|
|
os.makedirs(destination_dir)
|
|
|
|
if file_name is None:
|
|
file_name = file_key
|
|
destination_path = os.path.join(destination_dir, file_name)
|
|
self._s3_resource.Object(bucket_name, file_key).download_file(destination_path)
|
|
logger.info("Downloading {} to {}".format(file_key, destination_path))
|
|
|
|
def bucket_exists_in_s3(self, bucket_name):
|
|
"""
|
|
Verifies that the S3 bucket exists.
|
|
:param bucket_name: Name of the S3 bucket that may or may not exist.
|
|
:return: True if the bucket exists. False otherwise.
|
|
"""
|
|
bucket_exists = True
|
|
|
|
try:
|
|
self._s3_resource.meta.client.head_bucket(Bucket=bucket_name)
|
|
except botocore.exceptions.ClientError as err:
|
|
if err.response['Error']['Code'] == '404':
|
|
bucket_exists = False
|
|
|
|
return bucket_exists
|
|
|
|
def key_exists_in_bucket(self, bucket_name, file_key):
|
|
"""
|
|
Verifies that the given key does not already exist in the given S3 bucket.
|
|
:param bucket_name: Name of the S3 bucket that may or may not contain the file key.
|
|
:param file_key: Name of the file key in question.
|
|
:return: True if the key exists. False otherwise.
|
|
"""
|
|
key_exists = True
|
|
obj_summary = self._s3_resource.ObjectSummary(bucket_name, file_key)
|
|
|
|
# Attempting to access any member of ObjectSummary for a nonexistent key will throw an exception
|
|
# There is no built-in way to check key existence otherwise
|
|
try:
|
|
obj_summary.size
|
|
except botocore.exceptions.ClientError as err:
|
|
if err.response['Error']['Code'] == '404':
|
|
key_exists = False
|
|
|
|
return key_exists
|