""" YouTube Monthly Data Ingestion Workflow. Tasks to ingest data from YouTube Monthly feed into staging_raw_table_youtube_monthly (SnowFlake). """ from datetime import datetime from functools import partial import gzip import re from tempfile import TemporaryDirectory from garcon import task from garcon_contrib.dynamo_feed_status import garcon_feed_status from feed_ingestion.conf.config import merge_configs from feed_ingestion.flows.helpers import get_sf_config from feed_ingestion.flows.youtube_monthly import config from feed_ingestion.flows.youtube_monthly import util from feed_ingestion.flows.youtube_monthly.snowflake_executor \ import YoutubeMonthlySF from feed_ingestion.util import task_status, youtube_util from feed_ingestion.util.aws import s3 @task.decorate(timeout=1000) def bootstrap(activity, date, report_name, reload, accounts=''): """Bootstrap workflow for multiple monthly reports. Args: activity (ActivityWorker): The activity worker. date (str): Reporting date (YYYY-MM-DD). report_name (str): one of reports from config.reports reload (str): if == 'True' then force reload accounts (str): comma-separated MCN accounts, default to all Returns: dict: Context. """ assert report_name in config.reports report_config = config.reports[report_name] if accounts: report_accounts = {} for account in accounts.split(','): if account not in report_config['accounts']: raise ValueError(f'No account {account} for {report_name}') report_accounts[account] = report_config['accounts'][account] else: report_accounts = report_config['accounts'] if not report_accounts: raise ValueError('No accounts to process') date_obj = datetime.strptime(date, '%Y-%m-%d') file_template = report_config['file_pattern'] is_monthly = (report_config.get('period') == 'monthly' or '_M_' in file_template) is_weekly = (report_config.get('period') == 'weekly' or '_W_' in file_template) if is_weekly: # todo: move function to date_util.py _, first_day, last_day = youtube_util.get_days(date) elif is_monthly: # todo: move function to date_util.py first_day, last_day = util.get_first_last_day(date_obj) else: raise ValueError(f'Unknown report period for {file_template}') report_date = first_day.strftime('%Y-%m-%d') feed_name = f'{config.feed_name}_{report_name}' activity.logger.info(f'Bootstrap report {report_name} for {report_date}') if reload == 'True': activity.logger.info('Delete status for feed: {} {} '.format( feed_name, date)) garcon_feed_status.delete_status(feed_name, date) else: status = garcon_feed_status.get_overall_status(feed_name, report_date) if status == garcon_feed_status.STATUS_INGESTED: activity.logger.info( f'Feed {feed_name} already ingested for {report_date}' ) return { 'stop': True, 'message': f'{feed_name} is already ingested for {report_date}' } accounts_str = \ '({})'.format('|'.join(report_accounts)) file_pattern_regexp = \ '.+' + re.sub('{compression}', '', file_template).format( start_date=first_day, end_date=last_day, account=accounts_str) files = [] drop_location = report_config.get('drop_location', config.drop_location) s3_drop_path = config.s3_sftp_drop_path.format( drop_location=drop_location) for account, account_kwargs in report_accounts.items(): kwargs = { 'start_date': first_day, 'end_date': last_day, 'account': account, **account_kwargs } filename = file_template.format(**kwargs) csv_header_lines = report_config.get('csv_header_lines', 1) file_entry = { 'file_name': filename, 'mcn_account': account, 'file_size': None, 'csv_header_lines': csv_header_lines, } subreports_config = report_config.get('subreports') if subreports_config: subreports_file_entry = {} for subreport_name, info in subreports_config.items(): subreports_file_entry[subreport_name] = { 'report_title': info['report_title'], 'file_name': info['file_pattern'].format(**kwargs), 'report_name': f'{report_name}__{subreport_name}', 'staging_raw_table': info['staging_raw_table'], } file_entry['subreports'] = subreports_file_entry files.append(file_entry) source_files_dict = {'files': files} s3_archive_path = config.s3_report_archive_path_template.format( report_name=report_name, date=first_day, ) return dict( feed_name=feed_name, s3_archive_path=s3_archive_path, s3_drop_path=s3_drop_path, staging_raw_table=report_config['staging_raw_table'], source_files_dict=source_files_dict, date=report_date, report_name=report_name, file_pattern_regexp=file_pattern_regexp, ) @task.decorate(timeout=5000) def move_and_extract_files( activity, source_files_dict, old_s3_path, new_s3_path): """Move raw files from archives to temp folder and convert to files to gz. Args: activity (ActivityWorker): The activity worker. source_files_dict (dict): dict containing metadata of files being processed: { 'files": [ { 'file_name': 'my_cute_file' 'subreports': { 'sub_report1': { 'report_title': 'Sub Report name inside File', 'staging_raw_table': 'staging_raw_table_report_name__subreport_name', # NOQA: E501 'file_name': 'my_cute_file_sub_report1' } }, 'file_size': None, }, { 'file_name': 'my_nasty_file' 'file_size': None }, ] }. old_s3_path (str): Original drop location. new_s3_path (str): Location of extracted files for each source. """ for file_info in source_files_dict['files']: with TemporaryDirectory() as temp_path: local_temp_path = f'{temp_path}/' file_name = file_info['file_name'] old_s3_path_full = f'{old_s3_path}{file_name}' new_s3_path_full = f'{new_s3_path}{file_name}' if file_name.endswith('.zip'): file_name = file_name.replace('.zip', '.gz') new_temp_s3_path = f'{new_s3_path}{file_name}' # update file_info in order to return updated file_name to flow file_info['file_name'] = file_name activity.logger.info( f'Convert ZIP {old_s3_path_full} ' f'to {new_temp_s3_path}') s3.convert_zip_to_gzip_on_s3( activity, zip_s3_path=old_s3_path_full, gz_s3_path=new_temp_s3_path, local_temp_dir=local_temp_path) else: activity.logger.info( f'Copy {old_s3_path_full} to {new_s3_path_full}') s3.copy_s3_key( old_s3_path=old_s3_path_full, new_s3_path=new_s3_path_full ) if 'subreports' in file_info: # process multy report file source_bucket, source_key = s3.get_both(old_s3_path_full) downloaded_source_file_name = f'{local_temp_path}{file_name}' s3.download_from_s3( source_bucket, source_key, file_path=downloaded_source_file_name) report_filename_map = {} for subreport_name, info in file_info['subreports'].items(): subreport_filename = info['file_name'] report_filename_map[info['report_title']] = \ subreport_filename activity.logger.info( f'Going to extracted subreports: {report_filename_map}') if 'report_v1-0' in file_name: extracted_files = \ extract_financial_reports_from_multireport_file( activity, source_filename=downloaded_source_file_name, target_directory=local_temp_path, report_filename_map=report_filename_map ) else: extracted_files = extract_reports_from_multireport_file( activity, source_filename=downloaded_source_file_name, target_directory=local_temp_path, report_filename_map=report_filename_map ) activity.logger.info(f'Extracted subreports {extracted_files}') for subreport_file_name in extracted_files['files']: target_s3_path = f'{new_s3_path}{subreport_file_name}' target_bucket, target_key = s3.get_both(target_s3_path) file_path = f'{local_temp_path}{subreport_file_name}' activity.logger.info( f'Upload subreport file to {target_s3_path}') s3.upload_to_s3( file_path=file_path, bucket_name=target_bucket, object_key=target_key ) return dict(source_files_dict=source_files_dict) def extract_reports_from_multireport_file( activity, source_filename, target_directory, report_filename_map ): """Extract multiple reports stored in single source "CSV" file. The reports in youtube's pseudo-CSV file are separated by empty line. Next line is report name like 'Music Streams' Then it has normal CSV content with header line. Args: activity (ActivityWorker): The activity worker. source_filename (str): report file (csv or gz). target_directory (str): directory for extracted reports, report_filename_map (dict): key: report name, value: cvs filename, """ if source_filename.endswith('.gz'): open_func = partial(gzip.open, source_filename, mode='rt') else: open_func = partial(open, source_filename) processed_reports = [] processed_files = [] with open_func() as source_file: try: while True: line = next(source_file) report_name = line.strip() if report_name not in report_filename_map: raise ValueError(f'Unexpected report {report_name}') result_filename = report_filename_map[report_name] processed_reports.append(report_name) processed_files.append(result_filename) if not result_filename.endswith('.csv') \ and not result_filename.endswith('.csv.gz'): raise ValueError( f'Non-CSV/GZ report file name {result_filename} ' f'for report {report_name}') target_file_path = f'{target_directory}/{result_filename}' if result_filename.endswith('.gz'): open_write_func = partial( gzip.open, target_file_path, mode='wt') else: open_write_func = partial( open, target_file_path, mode='w') with open_write_func() as result_file: while True: line = next(source_file) if line == '\n': # reached empty line, going to next report break result_file.write(line) except StopIteration: # all good. The source file has ended pass missed_reports = set(report_filename_map) - set(processed_reports) if missed_reports: raise ValueError( f'Missed report(s) {missed_reports} in file {source_filename}') return {'files': processed_files} def extract_financial_reports_from_multireport_file( activity, source_filename, target_directory, report_filename_map ): """ Extract multiple financial reports stored in single file. The reports are separated by empty line except the Summary report which requires to skip an empty line. Args: activity (ActivityWorker): The activity worker. source_filename (str): report file (csv or gz). target_directory (str): directory for extracted reports, report_filename_map (dict): key: report name, value: cvs filename, """ if source_filename.endswith('.gz'): open_func = partial(gzip.open, source_filename, mode='rt') else: open_func = partial(open, source_filename) processed_reports = [] processed_files = [] with open_func() as source_file: # get metadata meta = {} for i in range(8): line = next(source_file) if not line == '\n': line_p = line.split(',') meta[line_p[0]] = line_p[1].replace('\n', '') try: while True: line = next(source_file) report_name = line.strip() if report_name not in report_filename_map: raise ValueError(f'Unexpected report {report_name}') result_filename = report_filename_map[report_name] processed_reports.append(report_name) processed_files.append(result_filename) if not result_filename.endswith('.csv') \ and not result_filename.endswith('.csv.gz'): raise ValueError( f'Non-CSV/GZ report file name {result_filename} ' f'for report {report_name}') target_file_path = f'{target_directory}/{result_filename}' if result_filename.endswith('.gz'): open_write_func = partial( gzip.open, target_file_path, mode='wt') else: open_write_func = partial( open, target_file_path, mode='w') # metadata meta_header = ','.join(list(meta.keys())) meta_values = ','.join(list(meta.values())) if report_name == 'Summary': with open_write_func() as result_file: while True: # line 0 header_0 = next(source_file).replace('\n', '') data_0 = next(source_file).replace('\n', '') # line break to ignore next(source_file) # line 1 header_1 = next(source_file) data_1 = next(source_file) # Put together header = '{},{},{}'.format( meta_header, header_0, header_1 ) data = '{},{},{}'.format( meta_values, data_0, data_1 ) # Write them up result_file.write(header) result_file.write(data) line = next(source_file) if line == '\n': # reached empty line, going to next report break else: with open_write_func() as result_file: line = next(source_file) header = '{},{}'.format( meta_header, line ) result_file.write(header) while True: line = next(source_file) if line == '\n': # reached empty line, going to next report break data = '{},{}'.format( meta_values, line ) result_file.write(data) except StopIteration: # all good. The source file has ended pass missed_reports = set(report_filename_map) - set(processed_reports) if missed_reports: raise ValueError( f'Missed report(s) {missed_reports} in file {source_filename}') return {'files': processed_files} @task.decorate(timeout=3600 * 5) def load_staging_raw_table_reports( activity, feed_name, date, report_name, s3_dir_path, staging_raw_table_name, source_files_dict, ): """Save data to staging raw table(s) for given report. Args: activity (ActivityWorker): The Garcon activity worker. feed_name (str): Feed name of workflow execution for status updates. date (str): Reporting date (YYYY-MM-DD). report_name (str): Name of the report to ingest. staging_raw_table_name (str): The staging table name for report. """ def load_staging_raw_for_filenames( staging_raw_table_name, report_name, filenames, mcn_account, lines_to_skip): with YoutubeMonthlySF(sf_config) as executor: temp_table_name = 'temp_{staging_raw_table_name}_{date}'.format( staging_raw_table_name=staging_raw_table_name, date=date.replace('-', '_')) executor.create_temp_staging_raw_table( report_name=report_name, table_name=temp_table_name ) activity.logger.info(f'{temp_table_name} was created') executor.load_temp_staging_raw_table( temp_staging_raw_table=temp_table_name, filenames=filenames, s3_dir_path=s3_dir_path, lines_to_skip=lines_to_skip, ) executor.clean_staging_raw_table( staging_raw_table=staging_raw_table_name, download_date=date, mcn_account=mcn_account ) executor.load_staging_raw_table( staging_raw_table=staging_raw_table_name, temp_staging_raw_table=temp_table_name, report_name=report_name, mcn_account=mcn_account, date=date, filenames=filenames, ) executor.drop_table(table=temp_table_name) if task_status.is_completed_task( feed_name, date, 'load_staging_raw_table_reports'): activity.logger.info( 'load_staging_raw_table_reports for {date} for {report_name} ' 'already complete'.format(date=date, report_name=report_name)) return sf_config = merge_configs(get_sf_config(config.secrets_path), {}) for file_info in source_files_dict['files']: mcn_account = file_info['mcn_account'] if staging_raw_table_name: lines_to_skip = file_info.get('csv_header_lines', 1) load_staging_raw_for_filenames( staging_raw_table_name=staging_raw_table_name, report_name=report_name, filenames=[file_info['file_name']], mcn_account=mcn_account, lines_to_skip=lines_to_skip, ) else: for subreport_name, info in file_info['subreports'].items(): filename = info['file_name'] load_staging_raw_for_filenames( staging_raw_table_name=info['staging_raw_table'], report_name=info['report_name'], filenames=[filename], mcn_account=mcn_account, lines_to_skip=1 )