"""Tasks for handling the generic S3 operations.""" import os import re from tempfile import NamedTemporaryFile import boto3 from boto3.s3.transfer import TransferConfig from botocore.exceptions import ClientError from garcon import task from garcon_contrib.aws.utils import garcon_s3 from feed_ingestion.flows import helpers from feed_ingestion.util.aws.s3 import copy_s3_key from feed_ingestion.util.aws.s3 import get_key_size from feed_ingestion.util.aws.s3 import get_list_of_files_and_directories from feed_ingestion.util.aws.s3 import upload_on_s3 @task.decorate(timeout=30) def rename_files_on_s3( activity, files_to_rename, expected_bucket_owner='437795906767'): """Copy S3 objects to new path and delete from old. Args: activity (ActivityWorker): The activity worker. files_to_rename (list): Pairs of old and new files names. expected_bucket_owner (str): The expected bucket owner. Example of a pair of the names in the list: files_to_rename = [ ["s3://bucket/path_to_the_file/GoodsIn_ESSN_20160827.csv.done", "s3://bucket/path_to_the_file/GoodsIn_ESSN_20160827.csv"], ] Returns: dict: {"renamed_s3_files": [list_of_renamed_files]}. """ s3_client = boto3.client('s3') renamed_s3_files = [] for old_file, new_file in files_to_rename: new_bucket, new_bucket_path = garcon_s3.extract_bucket_path(new_file) old_bucket, old_bucket_path = garcon_s3.extract_bucket_path(old_file) try: s3_client.copy_object( CopySource={'Bucket': old_bucket, 'Key': old_bucket_path}, Bucket=new_bucket, Key=new_bucket_path, ExpectedBucketOwner=expected_bucket_owner) except ClientError as err: # when source file not exists if err.response['Error']['Code'] == 'NoSuchKey': activity.logger.info( 'File {old_file} does not exist'.format(old_file=old_file)) # If old file is missing (already renamed on previous flow run) if helpers.check_s3_key_exist( new_file, expected_bucket_owner): renamed_s3_files.append( garcon_s3.get_destination_s3key_path( new_bucket, new_bucket_path)) continue activity.logger.error( 'Old file {old_file} and new file {new_file} both ' 'not exists'.format(old_file=old_file, new_file=new_file)) raise else: s3_client.delete_object( Bucket=old_bucket, Key=old_bucket_path, ExpectedBucketOwner=expected_bucket_owner ) activity.logger.info('{old_file} moved to {new_file}'.format( old_file=old_file, new_file=new_file)) renamed_s3_files.append(garcon_s3.get_destination_s3key_path( new_bucket, new_bucket_path)) return dict(renamed_s3_files=renamed_s3_files) @task.decorate(timeout=600) def remove_files_from_path(activity, path, return_deleted_files): """Remove all objects from a given s3 path. This is a boto3 version of garcon_s3.remove_files_from_path task. Args: activity (ActivityWorker): the swf activity worker. path (str): S3 URL formatted like 's3://bucket_name/folder1/folder2/' return_deleted_files (bool): whether return deleted file names or not Returns: dict: One element dict with a list of deleted files. """ bucket_name, key_prefix = garcon_s3.extract_bucket_path(path) s3_resource = boto3.resource('s3') bucket = s3_resource.Bucket(bucket_name) if not bucket.creation_date: activity.logger.info( 'S3 bucket s3://{} does not exist'.format(bucket_name)) return None result = bucket.objects.filter(Prefix=key_prefix).delete() if result: file_keys = [entry['Key'] for entry in result[0]['Deleted']] else: file_keys = [] if return_deleted_files is True: return { 's3.files_removed': file_keys } return {} @task.decorate(timeout=3000) def copy_file( activity, source_bucket_name, source_key_name, destination_bucket_name, destination_key_name, replace=True): """Copy file from one S3 bucket to another. Args: activity (ActivityWorker): The activity worker. source_bucket_name (str): Name of bucket to copy from. source_key_name (str): Name of file to copy. destination_bucket_name (str): Name of bucket to copy to. destination_key_name (str): Name of copied file. replace (bool|str): Optional. Replace the file if it already exists. Returns: dict: Dictionary indicating whether copy was successful. """ if type(replace) is str: replace = False if replace.lower() == 'false' else True file_name = destination_key_name.split('/')[-1] s3_resource = boto3.resource('s3') destination_obj = _get_key( activity, destination_bucket_name, destination_key_name, s3_resource) if not destination_obj: return {file_name: False} if not replace and destination_obj: return {file_name: True} # copy file try: destination_obj.copy( {'Bucket': source_bucket_name, 'Key': source_key_name}) activity.logger.info( 's3://{}/{} copied to s3://{}/{}'.format( source_bucket_name, source_key_name, destination_bucket_name, destination_key_name)) except ClientError: activity.logger.info( 's3://{}/{} does not exist'.format( source_bucket_name, source_key_name)) return {file_name: False} return {file_name: True} @task.decorate(timeout=3000) def copy_files( activity, s3_archive_path, s3_download_path, source_files_dict, need_all_files=True): """Archive certain files from the drop location to archive location. Args: activity (ActivityWorker): The activity worker. source_files_dict (dict): dict containing metadata of files. s3_archive_path (str): Archive location on S3. s3_download_path (str): Drop location on S3. need_all_files (bool): If True, all the files need to be present. Returns: dict: Context patch. """ def copy_one_file(download_path, archive_path, file_name): """Copy one file from download_path s3 path to archive_path.""" from_path = '{s3_path}{file_name}'.format( s3_path=download_path, file_name=file_name) to_path = '{s3_path}{file_name}'.format( s3_path=archive_path, file_name=file_name) try: copy_s3_key(from_path, to_path) activity.logger.info( 'File was copied from {from_path} to {to_path}'.format( from_path=from_path, to_path=to_path)) except ClientError as err: activity.logger.error( ('Cannot copy files from {drop_location} ' 'to {archive_location}. {exception_body}').format( drop_location=from_path, archive_location=to_path, exception_body=err)) found = False else: found = True return { 'file_name': file_name, 'file_size': get_key_size(to_path) if found else 0, 'file_path': from_path, 'found': found } result = [ copy_one_file(s3_download_path, s3_archive_path, fd['file_name']) for fd in source_files_dict['files']] file_statuses = list(map(lambda f: f['found'], result)) if not any(file_statuses) or (need_all_files and not all(file_statuses)): return dict(stop=True, msg=f'Need all files to proceed. ' f'Download result: "{result}"') return dict(source_files_dict={'files': result}) def _get_key(activity, bucket_name, key_name, s3_resource): """Get key for an S3 file. Args: activity (ActivityWorker): The activity worker. bucket_name (str): Name of bucket where file exists. key_name (str): Name of file. s3_resource (Resource): A Boto S3 Resource. Returns: Object: S3 object. """ if not s3_resource.Bucket(bucket_name).creation_date: activity.logger.info( 'S3 bucket s3://{} does not exist'.format(bucket_name)) return None return s3_resource.Object(bucket_name, key_name) def _get_sme_s3_client(secrets_path: str = None): aws_access_key_id = ( os.environ.get('SME_AWS_ACCESS_KEY_ID') or helpers.get_secret(secrets_path, 'SME_AWS_ACCESS_KEY_ID')) aws_secret_access_key = ( os.environ.get('SME_AWS_SECRET_ACCESS_KEY') or helpers.get_secret(secrets_path, 'SME_AWS_SECRET_ACCESS_KEY')) return boto3.client( 's3', aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key ) @task.decorate(timeout=3000) def copy_file_from_sme_s3_to_theocrhard( activity, secrets_path, source_bucket_name, source_key_name, destination_bucket_name, destination_key_name, replace=True): """Copy file from sme S3 bucket to theorchard s3. Args: activity (ActivityWorker): The activity worker. secrets_path (str): [DEPRECATED] Secrets manager path. source_bucket_name (str): Name of bucket to copy from. source_key_name (str): Name of file to copy. destination_bucket_name (str): Name of bucket to copy to. destination_key_name (str): Name of copied file. replace (bool|str): Optional. Replace the file if it already exists. Returns: dict: Dictionary indicating whether copy was successful. """ if type(replace) is str: replace = False if replace.lower() == 'false' else True file_name = destination_key_name.split('/')[-1] s3_resource = boto3.resource('s3') destination_obj = _get_key( activity, destination_bucket_name, destination_key_name, s3_resource) if not destination_obj: return {file_name: False} # check that file already on theorchard s3 if not replace and destination_obj: return { file_name: True, 'file_size': destination_obj.content_length } sme_s3_client = _get_sme_s3_client(secrets_path) with NamedTemporaryFile('wb') as file: try: sme_s3_client.download_fileobj( source_bucket_name, source_key_name, file, Config=TransferConfig()) file.flush() except ClientError as err: activity.logger.info( 'Cannot find file {}: error {}'.format( 's3://{}/{}'.format( source_bucket_name, source_key_name), err)) return {file_name: False} activity.logger.info('Uploading file {} to s3 {}'.format( file_name, destination_key_name)) upload_on_s3( destination_bucket_name, destination_key_name, '', file) # compare temp file size and file which was uploaded on theorchard s3 files_size = os.path.getsize(file.name) uploaded_size = get_key_size( 's3://{}/{}'.format(destination_bucket_name, destination_key_name)) if files_size != uploaded_size: return {file_name: 'Uploading failed'} else: return { file_name: True, 'file_size': files_size} def remove_prefix(text, prefix): """Remove prefix from string if string starts with it.""" return text[text.startswith(prefix) and len(prefix):] @task.decorate(timeout=1000) def source_files(activity, s3_bucket, s3_path, file_pattern): """Get list of files to by s3 location and file pattern. If no files found it raises ValueError. Args: activity (ActivityWorker): The activity worker. s3_bucket (str): bucket where to find search files. s3_path (str): prefix for the files in the bucket. file_pattern (str): regexp to match soruce files in the s3_path. it matches against the each s3 key in the prefix. Returns: dict: dict containing 'source_files_dict' with list of source files metadata. Metadata includes file_name and file_size """ logger = activity.logger s3_full_path = 's3://{}/{}'.format(s3_bucket, s3_path) source_files_dict = {'files': []} for file_path in get_list_of_files_and_directories(s3_full_path): file_name = remove_prefix(file_path, s3_path) if not re.match(file_pattern, file_name): logger.warning( f'Skipping non-matching source file {file_name}') continue file_size = get_key_size( 's3://{}/{}'.format(s3_bucket, file_path)) source_files_dict['files'].append({ 'file_name': file_name, 'file_size': file_size}) if not source_files_dict['files']: raise ValueError(f'No source files found in {s3_full_path}') return { 'source_files_dict': source_files_dict, }