"""S3 related helpers and Garcon tasks.""" from contextlib import contextmanager import csv import gzip from io import BytesIO import os from typing import Optional import zipfile import boto3 import botocore from botocore.exceptions import ClientError as BotoClientError from garcon import task from garcon_contrib.aws.utils import garcon_s3 import pandas as pd import s3fs import smart_open from feed_ingestion.conf import config from feed_ingestion.flows import helpers def _convert_zip_to_gzip( target_file_path, converted_file_path, local_temp_path, extract_original_filename=False): """Convert zip file to gzip file locally. Args: target_file_path (str): Full local path to target file. converted_file_path (str): Full local path to converted file. local_temp_path (str): Temp local file path. extract_original_filename (bool): If True, try to extract an original filename from an archive. """ with zipfile.ZipFile(target_file_path, 'r') as z: z.extractall(local_temp_path) if extract_original_filename: # try to obtain original filename from the archive archive_metadata = z.infolist() uncompressed_file = archive_metadata[0].orig_filename else: uncompressed_file = os.path.basename( target_file_path).replace('.zip', '') full_local_path_after_uncompressed = ( '{}{}').format(local_temp_path, uncompressed_file) with open(full_local_path_after_uncompressed, 'rb') as f_in: with gzip.open(converted_file_path, 'wb') as f_out: f_out.writelines(f_in) os.remove(full_local_path_after_uncompressed) # TODO: task should be moved to feed_ingestion.tasks.s3_tasks @task.decorate(timeout=2000) def convert_zip_to_gzip_on_s3( activity, zip_s3_path, gz_s3_path, local_temp_dir, extract_original_filename=False): """Convert zip file to gzip file on S3. Args: activity (ActivityWorker): The activity worker. zip_s3_path (str): full S3 path to zip file. gz_s3_path (str): full S3 path to gzip file. local_temp_dir (str): Optional path to the local temp directory. extract_original_filename (bool): If True, try to extract an original filename from an archive. Return: dict: Contains full S3 path to the .gz file. """ if local_temp_dir is None: local_temp_dir = './' # Download target file to local zip_file_name = os.path.basename(zip_s3_path) gz_file_name = os.path.basename(gz_s3_path) local_temp_path = '{}{}'.format(local_temp_dir, zip_file_name) helpers.download(zip_s3_path, local_temp_path) # Convert zip to gzip file _convert_zip_to_gzip( local_temp_path, '{}{}'.format( local_temp_dir, gz_file_name), local_temp_dir, extract_original_filename) # Upload converted file to S3 helpers.upload_raw_file_to_s3( '{}{}'.format(local_temp_dir, gz_file_name), gz_s3_path) # Remove local temp files os.remove(local_temp_path) os.remove('{}{}'.format(local_temp_dir, gz_file_name)) return {'converted_file_s3_full_path': gz_s3_path} # TODO: 0 usages. Remove unused @task.decorate(timeout=2000) def convert_s3_file_format( activity, converted_file_s3_dir, file_to_be_converted_dir, target_files, target_format, converted_format, local_temp_path): """Convert S3 zip file to gzip file. Args: activity (ActivityWorker): The activity worker. converted_file_s3_dir (str): Directory where converted S3 key object resides. file_to_be_converted_dir (str): Directory where file that will be converted resides. target_files (list): Target files to be converted. target_format (str): File format of the target file (eg. .zip, .gz). converted_format (str): File format of the converted file (eg. .gz). local_temp_path (str): Temp local file path. Returns: dict: Dictionary contains one key for the full path to the converted file. """ target_file = target_files.pop() converted_file = target_file.replace(target_format, converted_format) source_s3_path = '{}{}'.format(file_to_be_converted_dir, target_file) temp_path = '{}{}'.format(converted_file_s3_dir, converted_file) full_local_path_to_target_file = '{}{}'.format( local_temp_path, target_file) full_local_path_to_converted_file = ( '{}{}').format(local_temp_path, converted_file) # Download target file to local helpers.download(source_s3_path, full_local_path_to_target_file) # Convert the file if target_format == '.zip' and converted_format == '.gz': _convert_zip_to_gzip( full_local_path_to_target_file, full_local_path_to_converted_file, local_temp_path) else: raise AssertionError('Conversion is not supported.') # Upload converted file to S3 helpers.upload_raw_file_to_s3(full_local_path_to_converted_file, temp_path) # Remove local temp files os.remove(full_local_path_to_target_file) os.remove(full_local_path_to_converted_file) return {'converted_file_s3_full_path': temp_path} def copy_s3_key(old_s3_path, new_s3_path): """Copy S3 object to new path. Args: old_s3_path (str): Full s3 path - copying FROM. new_s3_path (str): Full S3 path - copying TO. """ s3 = boto3.client('s3') old_bucket, old_bucket_path = garcon_s3.extract_bucket_path(old_s3_path) new_bucket, new_bucket_path = garcon_s3.extract_bucket_path(new_s3_path) copy_source = { 'Bucket': old_bucket, 'Key': old_bucket_path } s3.copy( copy_source, new_bucket, new_bucket_path) def get_key_size(path, expected_bucket_owner=None): """Get size of file in key path. Args: path (str): Complete path in S3 to key/file. expected_bucket_owner (str): Expected bucket owner. Defaults to conf.config.EXPECTED_BUCKET_OWNER Returns: (int): File size of file in S3. """ s3_client = boto3.client('s3') bucket, key_path = garcon_s3.extract_bucket_path(path) if not expected_bucket_owner: expected_bucket_owner = config.EXPECTED_BUCKET_OWNER return s3_client.get_object( Bucket=bucket, Key=key_path, ExpectedBucketOwner=expected_bucket_owner)['ContentLength'] def expand_s3_csv_with_columns( csv_s3_key_path, columns_header, columns_values, path_to_unload, line_rstrip_char=None): """Expand passed file with passed columns. (And copy it to specified path in the same S3 bucket). Args: csv_s3_key_path (str): CSV file (S3 key for file to be transform). columns_header (tuple): Names of the new columns for a header row. columns_values (tuple): Values of the new columns for the data rows. path_to_unload (str): Path on S3 where to unload the expanded copy. line_rstrip_char (str): Optional right strip string for each line in CSV file. """ bucket_name, path = garcon_s3.extract_bucket_path(csv_s3_key_path) transformed_file_name = csv_s3_key_path.split('/')[-1] transformed_csv_path = garcon_s3.get_destination_s3key_path( bucket_name, '{path_to_unload}/{file_name}'.format( path_to_unload=path_to_unload, file_name=transformed_file_name)) with smart_open.smart_open(transformed_csv_path, 'wt') as temp_transformed: smart_open_obj = smart_open.smart_open(csv_s3_key_path) csv_reader_obj = csv.reader( (line.decode().rstrip(line_rstrip_char) for line in smart_open_obj) ) csv_writer_obj = csv.writer(temp_transformed) for row_number, row in enumerate(csv_reader_obj): if not row_number: row.extend(columns_header) else: row.extend(columns_values) csv_writer_obj.writerow(row) def delete_file(s3_path_to_file, expected_bucket_owner='437795906767'): """Delete a specified file on S3. Args: s3_path_to_file (str): Full S3 path to the file, including file_name. expected_bucket_owner (str): Expected bucket owner. """ s3_bucket, path = garcon_s3.extract_bucket_path(s3_path_to_file) s3_client = boto3.client('s3') s3_client.delete_object( Bucket=s3_bucket, Key=path, ExpectedBucketOwner=expected_bucket_owner ) def get_file_size(s3_path, expected_bucket_owner='437795906767'): """Get size of a file by a given S3 key. It is a duplicate for get_key_size. TODO: remove duplicate Args: s3_path (str): Full S3 path to the file. expected_bucket_owner (str): Expected bucket owner. Returns: int: Size of file in bytes. """ bucket, path = garcon_s3.extract_bucket_path(s3_path) s3_client = boto3.client('s3') return s3_client.head_object( Bucket=bucket, Key=path, ExpectedBucketOwner=expected_bucket_owner)['ContentLength'] def get_list_of_files_and_directories(s3_path): """Get the list of files and dirs in given S3 directory (prefix). Args: s3_path (str): S3 directory (prefix) to list the files and dirs. Returns: files_list (list): List of the files in directories in given S3 path. """ s3_client = boto3.client('s3') bucket, bucket_path = garcon_s3.extract_bucket_path(s3_path) paginator = s3_client.get_paginator('list_objects') operation_parameters = {'Bucket': bucket, 'Prefix': bucket_path} page_iterator = paginator.paginate(**operation_parameters) files_and_directories = [] for page in page_iterator: for key in page.get('Contents', []): files_and_directories.append(key['Key']) return files_and_directories def upload_processed_to_s3(csv_obj, s3_path, fileobj=None, expected_bucket_owner='437795906767'): """Compress and upload csv data to s3. Args: csv_obj (StringIO): Csv file object to upload. s3_path (str): S3 path for uploading. fileobj (file or None): File-like object storage for compressed data. expected_bucket_owner (str): Expected bucket owner. """ fileobj = fileobj or BytesIO() gz = gzip.GzipFile( filename=None, mode='wb', compresslevel=9, fileobj=fileobj) csv_obj.seek(0) for line in csv_obj: gz.write(line.encode('utf-8')) gz.close() bucket, bucket_path = garcon_s3.extract_bucket_path(s3_path) s3_client = boto3.client('s3') fileobj.seek(0) s3_client.upload_fileobj( fileobj, bucket, bucket_path, ExtraArgs={'ExpectedBucketOwner': expected_bucket_owner} ) return { 'source_files_dict': { 'files': [ {'file_name': s3_path.split('/')[-1], 'found': True, 'file_size': gz.size}]}} def get_source_files_content( s3_path, source_files_dict, lines=None, encoding='utf-8'): """Get source files content. Generator that yields s3 file content for each file in files list. Args: s3_path (str): S3 directory path to files. source_files_dict (dict): List of files to download. lines (int): Number of lines to read from source file. If None then read the whole file. encoding (str): Output string encoding. Return bytes if encoding is None. Yields: tuple: pair of file name and file content as a string. """ bucket, bucket_path = garcon_s3.extract_bucket_path(s3_path) for file in source_files_dict['files']: file_path = 's3://{}/{}{}'.format( bucket, bucket_path, file['file_name']) rows = [] with smart_open.smart_open(file_path, 'rb') as content: if lines: for i, line in enumerate(content, 1): if i > lines: break if encoding: rows.append(line.decode(encoding)) else: rows.append(line) else: for line in content: if encoding: rows.append(line.decode(encoding)) else: rows.append(line) if encoding: yield file['file_name'], ''.join(rows) else: yield file['file_name'], b''.join(rows) def upload_on_s3(bucket, archive_path, filename, fd, expected_bucket_owner='437795906767'): """Upload file on S3. Args: bucket (str): S3 bucket name. archive_path (str): Archive path on S3. filename (str): Filename. fd (File): File descriptor of file to upload. expected_bucket_owner (str): Expected bucket owner. """ s3_client = boto3.client('s3') key_name = '{bucket_path}{file_name}'.format( bucket_path=archive_path, file_name=filename) s3_client.upload_file( fd.name, bucket, key_name, ExtraArgs={'ExpectedBucketOwner': expected_bucket_owner} ) def get_both(s3_path): """Parse string of the s3_path to return the (bucket, key + filename).""" bucket_ = s3_path.replace('s3://', '').split('/')[0] key_ = s3_path.split(bucket_)[-1][1:] return (bucket_, key_) def read_csv( s3_path, sep=',', compression='infer', nrows=None, skiprows=None, header='infer', names=None, index_col=None, usecols=None, engine=None, parse_dates=False, lineterminator=None, escapechar=None, encoding=None, iterator=False, chunksize=None, dtype=None, low_memory=True, expected_bucket_owner='437795906767'): """Read a csv file from s3 into memory in a pandas dataframe.""" s3_client = boto3.client('s3') if 's3://' in s3_path: """ If a full s3 path is given BUCKET_NAME is parsed from the user param used rather than with the constant from constants.py. """ _bucket_, _key_ = get_both(s3_path) # todo: fix naming else: _key_ = s3_path _bucket_ = 'dev-rsaporta' # todo: remove dev-code try: buffer_in_binary = s3_client.get_object( Bucket=_bucket_, Key=_key_, ExpectedBucketOwner=expected_bucket_owner )['Body'] except botocore.exceptions.ClientError as e: return 'Unexpected error: %s' % e # note should just ask for kwargs. na_values = [ 'N/A', 'n/a', '#NA', 'NULL', 'null', 'NaN', '-NaN', 'nan', '-nan', ''] return pd.read_csv( filepath_or_buffer=buffer_in_binary, sep=sep, dtype=dtype, compression=compression, nrows=nrows, header=header, names=names, index_col=index_col, usecols=usecols, engine=engine, parse_dates=parse_dates, skiprows=skiprows, lineterminator=lineterminator, escapechar=escapechar, iterator=iterator, chunksize=chunksize, encoding=encoding, low_memory=low_memory, keep_default_na=False, na_values=na_values) def upload_to_s3( file_path: str, bucket_name: str, object_key: str ) -> Optional[int]: """Upload a local file into the archive folder. Args: file_path: dropped file local path bucket_name: destination S3 bucket name (i.e. 'dev-cucumbers') object_key: file object key (i.e. 'QQ/archives/2021-02-01/ORC_20210201_CONTENT_IDS.txt.gz') """ # Setup S3 bucket s3_resource = boto3.resource('s3') s3_bucket = s3_resource.Bucket(bucket_name) # Check local file size # Upload file to s3. # Check uploaded object length local_file_length = int(os.stat(file_path).st_size) try: s3_bucket.upload_file(file_path, object_key) remote_file_length = int(s3_bucket.Object(object_key).content_length) except BotoClientError as err: if err.response['Error']['Code'] == 'NoSuchKey': return None raise err # Verify uploaded file size if remote_file_length != local_file_length: raise ValueError( f'remote {object_key} size is not matching local {file_path} size') return remote_file_length def download_from_s3( bucket_name: str, object_key: str, file_path: str ) -> Optional[int]: """Download S3 object to a local path. Args: bucket_name: destination S3 bucket name (i.e. 'dev-cucumbers') object_key: file object key (i.e. 'QQ/archives/2021-02-01/ORC_20210201_CONTENT_IDS.txt.gz') file_path: dropped file local path """ # Setup S3 bucket s3_resource = boto3.resource('s3') s3_bucket = s3_resource.Bucket(bucket_name) # Upload file to s3. # Check uploaded object length # Check local file size try: s3_bucket.download_file(object_key, file_path) remote_file_length = int(s3_bucket.Object(object_key).content_length) except BotoClientError as err: if err.response['Error']['Code'] == 'NoSuchKey': return None raise err # Verify downloaded file size local_file_length = int(os.stat(file_path).st_size) if local_file_length != remote_file_length: raise ValueError( f'remote {object_key} size is not matching local {file_path} size') return remote_file_length def delete_s3_obj(bucket_name: str, object_key: str) -> None: """Delete S3 file object. Args: bucket_name: destination S3 bucket name (i.e. 'dev-cucumbers') object_key: file object key (i.e. 'QQ/archives/2021-02-01/ORC_20210201_CONTENT_IDS.txt.gz') """ # Setup s3 bucket s3_resource = boto3.resource('s3') s3_bucket = s3_resource.Bucket(bucket_name) # Check if file exists # Delete the file s3_bucket.delete_objects( Delete={'Objects': [ {'Key': object_key} ], 'Quiet': True}) def disk_2_s3(file, s3_path): """Send a file in local disk to s3 bucket. Duplicate for upload_on_s3. TODO: remove duplicate Please note that the s3_path needs to be this format: s3://bucket/key/filename.ext """ # make connection to s3 s3 = boto3.resource('s3') if 's3://' in s3_path: """ If a full s3 path is given bucket_name is parsed from the user param used rather than with the constant from constants.py. """ bucket_name, key_ = get_both(s3_path) bucket = s3.Bucket(bucket_name) else: bucket_name = 'dev-rsaporta' # mind-blowing! key_ = s3_path.replace(bucket_name, '') bucket = s3.Bucket(bucket_name) # upload file to s3. try: bucket.upload_file( Filename=file, Key=key_ ) except botocore.exceptions.ClientError as e: return ( "Unexpected error: {err} for pattern '{key}' in {bucket}".format( err=e, key=key_, bucket=bucket_name)) except Exception: return 'Write Permissions Denied' return "'{f}' loaded to '{path}'".format(f=file, path=s3_path) def to_csv(df, s3_path, index=False, compression=None, sep=',', quoting=None, chunksize=None, line_terminator='\n', escapechar=None, date_format=None, na_rep='', encoding=None): """Write a dataframe to local, then uploads the dataframe to s3."""'' _temp_file_ = s3_path.split('/')[-1] # note should just ask for kwargs. df.to_csv(_temp_file_, index=index, compression=compression, sep=sep, quoting=quoting, chunksize=chunksize, lineterminator=line_terminator, escapechar=escapechar, date_format=date_format, na_rep=na_rep, encoding=encoding) disk_2_s3(_temp_file_, s3_path) os.remove(_temp_file_) return "File uploaded to '%s'" % s3_path @contextmanager def open_s3_stream(path, mode, block_size=None): """Open S3 file stream. Args: path (str): S3 file path (e.g. 's3://bucket/Dir/file') mode (str): file access mode. The same as builtin open(). block_size (int): Buffer block size in bytes. Returns: File: File-like object to s3 file. """ fs = s3fs.S3FileSystem(default_fill_cache=False) with fs.open(path, mode, block_size=block_size) as f: yield f @contextmanager def open_s3_gzip_stream(path, mode, block_size=None): """Open Gzip S3 file stream. This context manager is used to stream gzipped to/from S3 transparently. Args: path (str): S3 file path (e.g. 's3://bucket/Dir/file') mode (str): file access mode. The same as builtin open(). block_size (int): Buffer block size in bytes. Returns: File: File-like object to gzipped file on S3. """ with open_s3_stream(path, mode, block_size) as f: with gzip.GzipFile(fileobj=f) as gz: yield gz