"""Tasks of the Amazon Digital Services Ingestion Workflow.""" import bz2 import datetime import gzip import os import subprocess import uuid import zipfile import boto3 from botocore.exceptions import ClientError from garcon import task from garcon_contrib.aws.utils import garcon_s3 from garcon_contrib.dynamo_feed_status import garcon_feed_status from garcon_contrib.ftp import garcon_ftp from feed_ingestion.conf.config import BOTO3_CONFIG from feed_ingestion.flows.amazon_digital_services import config from feed_ingestion.tasks import s3_tasks from feed_ingestion.util.aws import s3 NOTIFICATION_EMAIL_SUBJECT = ( 'Alert: Amazon Digital Services files are missing') NOTIFICATION_EMAIL_BODY = ( 'Amazon Digital Services missing files:\n{}') def _unzip_file(full_local_zip_path): """Unzip a file. Args: full_local_zip_path (str): Full path to a local .zip archive. Returns: str: Full path to an extracted file. """ local_temp_dir = os.path.dirname(full_local_zip_path) with zipfile.ZipFile(full_local_zip_path) as z: z.extractall(local_temp_dir) archive_metadata = z.infolist() uncompressed_file = archive_metadata[0].orig_filename # type: ignore return os.path.join(local_temp_dir, uncompressed_file) def _unbz2_file(full_local_bz2_path): """Decompress bz2 file. Args: full_local_bz2_path (str): Full path to a local .bz2 archive. Returns: str: Full path to an extracted file. """ uncompressed_file_path = full_local_bz2_path.replace('.bz2', '') with bz2.open(full_local_bz2_path) as bz2_file: file_content = bz2_file.read() with open(uncompressed_file_path, 'wb') as f: f.write(file_content) return uncompressed_file_path def _remove_footer(full_local_path_uncompressed, lines_to_remove=0): """Remove footer. (We have to remove one footer line from summary_statement files). Args: full_local_path_uncompressed (str): Full path to an uncompressed file. lines_to_remove (int): Number of lines to remove. """ if lines_to_remove: tmp_filename = uuid.uuid4() sed_cmd_to_execute = ( 'sed -e "$(($(wc -l < {path}) - {lines_to_remove}))' r',\$d" {path} > {tmp_filename} ; mv {tmp_filename} ' # noqa '{path}').format( path=full_local_path_uncompressed, lines_to_remove=lines_to_remove - 1, tmp_filename=tmp_filename) subprocess.check_call(sed_cmd_to_execute, shell=True) @task.decorate(timeout=1000) def bootstrap(activity, date, reload, licensor): """Bootstrap workflow by getting the correct configurations. Args: activity (ActivityWorker): The activity worker. date (str): Reporting date (YYYY-MM-DD). reload (str or None): If 'True' delete all feed statuses in DynamoDB. licensor (str): one of config.licensors Returns: dict: Context. """ # date is the date passed in or yesterday's date date_obj = datetime.datetime.strptime( date, '%Y-%m-%d').date() if date else datetime.date.today() if not licensor: licensor = 'theorchard' assert licensor in config.licensors, f'unsupported licensor "{licensor}"' feed_name = '_'.join([config.feed_name, licensor]) activity.logger.info('Bootstrap {} for date {}'.format(feed_name, date_obj)) if reload == 'True': activity.logger.info('Delete status for feed: {} {} '.format( feed_name, date)) garcon_feed_status.delete_status(feed_name, date) else: overall_status = garcon_feed_status.get_overall_status( feed_name, date) if overall_status == garcon_feed_status.STATUS_INGESTED: message = 'already ingested for {}'.format(date) activity.logger.info(message) return { 'message': message, 'stop': True } s3_archive_path = config.s3_archive_path_template.format( licensor=licensor, date=date_obj) s3_tmp_path = config.s3_tmp_path_template.format( licensor=licensor, date=date_obj) file_pattern_template = \ config.filename_templates[licensor]['search_pattern'] file_pattern = file_pattern_template.format( date=date_obj ) staging_raw_table = config.snowflake_table_names['staging_raw'][licensor] return dict( feed_name=feed_name, secrets_path=config.secrets_path, licensor=licensor, staging_raw_table=staging_raw_table, s3_archive_path=s3_archive_path, s3_tmp_path=s3_tmp_path, file_pattern=file_pattern, date=date_obj.strftime('%Y-%m-%d'), ) def _is_required_vendor_country_for_date(vendor, country, date_obj, licensor): # let's find most recent requirements for given date vendor_countries = _get_vendor_countries_for_date(date_obj, licensor) return vendor in vendor_countries and \ country in vendor_countries[vendor]['required'] def _get_vendor_countries_for_date(date_obj, licensor): vendor_countries = () for date, value in config.vendors_countries[licensor]: if date_obj >= date: vendor_countries = value if not vendor_countries: raise ValueError(f'cannot find requirements for date {date_obj}') return vendor_countries @task.decorate(timeout=7200) def fetch_from_drop_location( activity, date, feed_name, s3_tmp_path, licensor, use_s3): """Copy source files from source S3/FTP to target S3. Args: activity (Activity): Activity instance. date (str): Date being processed. feed_name (str): Name of the feed. s3_tmp_path (str): S3 directory to write files to. licensor (str): one of config.licensors use_s3 (str): 'True' if it is need to load file from s3. Returns: dict: A dict with the the metadata about the downloaded files. """ def _get_filenames(vendor, country, date_obj, licensor): """Get source and target file and folder names.""" filename_templates = config.filename_templates[licensor] format_values = dict(vendor_code=vendor, date=date_obj, country=country, country_lower=country.lower()) if licensor in ('theorchard'): if country == 'JP': filename_template = filename_templates['daily_zip'] target_filename_template = filename_template else: # Amazon appends lowercase country code for statement files # since 2021-05-05 if date_obj >= datetime.date(year=2021, month=5, day=5): filename_template = filename_templates[ 'summary_country_lower_zip'] else: filename_template = filename_templates['summary_zip'] # we need to add country suffix to filenamtarget e # in order to keep them in one directory target_filename_template = \ filename_templates['summary_country_zip'] else: assert licensor == 'sme' if date_obj < datetime.date(year=2017, month=7, day=12): filename_template = filename_templates['d'] else: filename_template = filename_templates['d_country'] target_filename_template = filename_templates['d_country'] source_filename = filename_template.format(**format_values) source_dir = filename_templates['remote_dir'].format( **format_values) target_filename = target_filename_template.format(**format_values) return source_dir, source_filename, target_filename def _copy_file(source_dir, source_filename, target_filename): """Copy file from source to target S3. Returns: dict: - file: filename - status: true/false successful - file_size: file size in bytes - exception: exception (on exception) """ if licensor in ('theorchard'): try: return garcon_ftp.copy_file_from_ftp_to_s3( activity=activity, ftp_creds=config.sftp[licensor], ftp_path=source_dir, ftp_file_name=source_filename, s3_path=s3_tmp_path, s3_file_name=target_filename) except Exception as e: activity.logger.warning( f'File processing failed: {source_dir}{source_filename}: ' f'{repr(e)}') return { 'file': source_filename, 'status': False, 'file_size': -1, 'exception': e, } else: assert licensor == 'sme' exception = None local_filename = None file_size = -1 status = False try: source_key_name = '{}{}'.format(source_dir, source_filename) path = garcon_s3.extract_bucket_path(s3_tmp_path)[1] destination_key_name = '{}{}'.format(path, target_filename) copy_result = s3_tasks.copy_file_from_sme_s3_to_theocrhard( activity=activity, secrets_path=config.sme_secrets_path, source_bucket_name=config.sme_drop_bucket, source_key_name=source_key_name, destination_bucket_name=config.s3_bucket, destination_key_name=destination_key_name, replace=True) if copy_result[target_filename] is True: status = True file_size = copy_result['file_size'] except Exception as e: activity.logger.warning( f'File processing filed: {source_dir}{source_filename}: ' f'{repr(e)}') exception = e finally: _remove_file(local_filename) return { 'file': target_filename, 'status': status, 'file_size': file_size, 'exception': exception, } def _validate_response(copy_response): """Check copy_response for requirements.""" found = copy_response['status'] required_not_found = False message = 'Validate response for file {}. Details: {}'.format( source_filename, copy_response) activity.logger.info(message) if found: activity.logger.info('File downloaded: %s', copy_response) elif not _is_required_vendor_country_for_date( vendor_code, country_code, date_obj, licensor): # if this vendor/country is not required, continue activity.logger.info('File not found: %s. Skip it.', copy_response) else: message = 'Expected file {} was not downloaded'.format( source_filename) activity.logger.info(message) required_not_found = source_filename return { 'file_name': copy_response['file'], 'file_size': copy_response.get('file_size', -1), 'found': found, 'required_not_found': required_not_found, 'exception': str(copy_response.get('exception', None)) } activity.logger.info('Fetching source files: %s', date) date_obj = datetime.datetime.strptime(date, '%Y-%m-%d').date() result = [] # copy files and validate if they are required vendor_countries = _get_vendor_countries_for_date(date_obj, licensor) for vendor_code, countries in vendor_countries.items(): all_countries = countries['required'] + countries['optional'] for country_code in all_countries: source_dir, source_filename, target_filename = \ _get_filenames(vendor_code, country_code, date_obj, licensor) copy_response = _copy_file(source_dir, source_filename, target_filename) validated_response = _validate_response(copy_response) result.append(validated_response) required_not_found_files = [ s3_tmp_path + item['file_name'] for item in filter(lambda item: item['required_not_found'], result) ] file_statuses = list(map(lambda f: f['found'], result)) if any(required_not_found_files): garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_AVAILABLE) garcon_feed_status.set_missing_files( feed_name, date, required_not_found_files) return dict( stop=True, message='Not all the required files are present yet', missing_files=','.join(required_not_found_files) ) if not any(file_statuses): return dict(stop=True, message='No files present yet') return dict(source_files_dict={'files': result}) def _remove_file(local_filename): if local_filename: try: os.remove(local_filename) except OSError: pass @task.decorate(timeout=5000) def prepare_files_for_ingestion( activity, feed_name, date, source_files_dict, s3_tmp_path, s3_archive_path, licensor): """Сonvert source files to .gz, remove footer if necessary, upload on S3. Args: activity (ActivityWorker): The activity worker. date (str): Date being processed. feed_name (str): Name of the feed. source_files_dict (dict): The dict with the file names and metadata. s3_tmp_path (str): Temp S3 dir where source .zip files reside. s3_archive_path (str): Archive S3 dir where converted to .gz files will reside. licensor (str): one of config.licensors. """ local_temp_dir = os.path.realpath(os.getcwd()) + '/' for file_dict in filter(lambda f: f['found'], source_files_dict['files']): source_s3_path = '{}{}'.format( s3_tmp_path, file_dict['file_name']) full_local_source_path = os.path.join( local_temp_dir, file_dict['file_name']) # file type compressions and functions to decompress file_types = { 'zip': _unzip_file, # theorchard files from amazon sftp } file_type = file_dict['file_name'].split('.')[-1] if file_type in file_types: full_local_gzip_path = os.path.join( local_temp_dir, file_dict['file_name'].replace(f'.{file_type}', '.gz')) gz_s3_path = '{}{}'.format( s3_archive_path, file_dict['file_name'].replace(f'.{file_type}', '.gz')) else: full_local_gzip_path = os.path.join( local_temp_dir, f"{file_dict['file_name']}.gz") gz_s3_path = '{}{}.gz'.format( s3_archive_path, file_dict['file_name']) # download source file from S3 try: s3.helpers.download( source_s3_path, full_local_source_path, config.expected_bucket_owner ) except ClientError as err: activity.logger.error(f'Download failed {source_s3_path}') garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_INGESTED) raise err if file_type in file_types: # unzip file try: full_local_path_uncompressed = file_types[file_type]( full_local_source_path) except Exception as e: garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_INGESTED) activity.logger.error( '{file} was not a valid zipfile '.format( file=full_local_source_path)) raise e else: full_local_path_uncompressed = full_local_source_path if licensor in ('theorchard'): # remove footer lines_to_remove = 1 if ('JP' in full_local_source_path or 'AmazonMP3' in full_local_source_path): lines_to_remove = 0 _remove_footer(full_local_path_uncompressed, lines_to_remove) # gzip file with open(full_local_path_uncompressed, 'rb') as f_in: with gzip.open(full_local_gzip_path, 'wb') as f_out: f_out.writelines(f_in) # upload converted file to S3 try: s3.helpers.upload_raw_file_to_s3( full_local_gzip_path, gz_s3_path, config.expected_bucket_owner ) except ClientError as err: garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_INGESTED) raise err # Remove local temp files _remove_file(full_local_path_uncompressed) _remove_file(full_local_source_path) _remove_file(full_local_gzip_path) activity.logger.info( 'Source files were converted to .gz, footers were removed') @task.decorate(timeout=5000) def notify_missing_files(activity, source_files_dict): """Notify about missing files. Args: activity (ActivityWorker): The activity worker. source_files_dict (dict): The dict with the file names and metadata. """ missing_files = [ f['file_name'] for f in source_files_dict['files'] if f['required_not_found']] client = boto3.client('ses', config=BOTO3_CONFIG) client.send_email( Source=config.EMAIL_SOURCE, Destination={'ToAddresses': config.EMAIL_LIST}, Message={ 'Subject': { 'Data': NOTIFICATION_EMAIL_SUBJECT, 'Charset': 'UTF-8'}, 'Body': { 'Text': { 'Data': NOTIFICATION_EMAIL_BODY.format( '\n'.join(missing_files)), 'Charset': 'UTF-8' }}})