"""Tasks for handling YouTube specific operations.""" import json import os from pathlib import Path from tempfile import gettempdir, TemporaryDirectory import boto3 from garcon import task from garcon_contrib.aws.utils import garcon_s3 from garcon_contrib.dynamo_feed_status import garcon_feed_status from snowflake_connector.etl_connector import SQLLoader from feed_ingestion import conf from feed_ingestion.conf.config import merge_configs from feed_ingestion.flows import registered_executors from feed_ingestion.flows.helpers import get_sf_config from feed_ingestion.tasks import s3_tasks from feed_ingestion.util import file_operations from feed_ingestion.util import sentry_util from feed_ingestion.util import task_status from feed_ingestion.util import youtube_util from feed_ingestion.util.aws import athena from feed_ingestion.util.aws import s3 as s3utils GENERATED_PARQUET_FILENAME_PATTERN = \ r'^\d{8}_\d{6}_\d{5}_\w+_\w+\-\w+\-\w+\-\w+\-\w+$' ATHENA_QUERY_TIMEOUT_SECONDS = 60 * 60 # 60 minutes @task.decorate(timeout=12600) def update_channel_names_table( activity, date, feed_name, sfdb_params, secrets_path=None): """Update YouTube channel names mapping table. Args: activity (ActivityWorker): The activity worker. date (str): Reporting date (YYYY-MM-DD). feed_name (str): Name of the feed to get executor class. sfdb_params (dict): Dict with params to optionally override default ones (Snowflake db and schema name). secrets_path (str): Secrets manager path of the flow. """ task_id = 'update_channel_names_table' if task_status.is_completed_task(feed_name, date, task_id): activity.logger.info( 'update_channel_names_table for {date} already complete, and ' '"reload" flag was not passed, skipping...'.format(date=date)) return sf_config = get_sf_config(secrets_path) sf_config_custom = merge_configs(sf_config, sfdb_params) ExecutorSR = registered_executors.get(feed_name) with ExecutorSR(sf_config_custom) as sf_executor: sf_executor.update_channel_names_table(date) sf_executor.cleanup_duplicated_channel_names() activity.logger.info('Updated channel names mapping table.') @task.decorate(timeout=3600) def sme_copy_from_athena_to_s3( activity, date, report_status_name, sme_athena_database, sme_athena_temp_database, sme_athena_source_table, destination_s3_bucket, destination_s3_path, athena_workgroup, aws_region='us-east-1'): """Extract PARQUET fields from SME YouTube report. Args: activity (ActivityWorker): The activity worker. date (str): Reporting date (YYYY-MM-DD). report_status_name (str): Status name for this report in DynamoDB sme_athena_database (str): name of SME database in Athena sme_athena_temp_database (str): name of TEMP database in Athena sme_athena_source_table (str): table name in Athena database destination_s3_bucket (str): bucket where to put the result PARQUET destination_s3_path (str): path in destination_s3_bucket (should starts and ends with "/") athena_workgroup (str): Athena Workgroup """ if not sme_athena_source_table: raise ValueError('athena_source_table required') if not destination_s3_path.endswith('/'): raise ValueError('destination_s3_path should end with /') def _load_athena_query(sql_template_name, sql_params): sql_loader = SQLLoader(__file__) activity.logger.info(f'Loading query file {sql_template_name}') sql_template = sql_loader.load_query(sql_template_name) sql_query = sql_template.format( **sql_params ) return sql_query task_id = 'copy_from_athena_to_s3' activity.logger.info(f'Start copy_from_athena_to_s3 date {date} ' f'report_status_name {report_status_name}') if task_status.is_completed_task(report_status_name, date, task_id): activity.logger.info( 'copy_from_athena_to_s3 for {date} already complete, and ' '"reload" flag was not passed, skipping...'.format(date=date)) return sql_params = { 'athena_sme_db': sme_athena_database, 'athena_sme_table': sme_athena_source_table, 'date': date, } sql_query = _load_athena_query( 'sme_youtube_athena_load_for_date', sql_params) s3_url = f's3://{destination_s3_bucket}/{destination_s3_path}' s3_tasks.remove_files_from_path( activity=activity, path=s3_url, return_deleted_files=False) athena.run_query( athena_query=sql_query, athena_temp_database=sme_athena_temp_database, athena_workgroup=athena_workgroup, destination_s3_bucket=destination_s3_bucket, destination_s3_path=destination_s3_path, timeout=ATHENA_QUERY_TIMEOUT_SECONDS ) # check if we have any data files and update overall status try: result = s3_tasks.source_files( activity=activity, s3_bucket=destination_s3_bucket, s3_path=destination_s3_path, file_pattern=GENERATED_PARQUET_FILENAME_PATTERN ) overall_status = garcon_feed_status.STATUS_DOWNLOADED activity.logger.info(f'Done copy_from_athena_to_s3 date {date} ' f'report_status_name {report_status_name}') except ValueError: activity.logger.error(f'Athena data files not found in {s3_url}') overall_status = garcon_feed_status.STATUS_NOT_AVAILABLE result = {'stop': True, 'message': 'Athena data files not found'} garcon_feed_status.set_overall_status(report_status_name, date, overall_status) return result def save_youtube_access_token_from_secrets( path: Path, secrets_path='swf-shared-youtube-access-token'): """Save YouTube access token from secrets manager to file.""" if path.exists(): return environment = conf.config.ENV secrets_client = boto3.client('secretsmanager', region_name=conf.config.AWS_REGION) secret_name = (f'{environment}' f'/{secrets_path}' f'/YOUTUBE_ACCESS_TOKEN') with open(path, 'w') as f: secret_value = secrets_client.get_secret_value( SecretId=secret_name) f.write(secret_value['SecretString']) @task.decorate(timeout=16000) def grab_reports_files( activity, report_name, report_status_name, date, archive_path, credentials_path, api_service_name, api_version, jobs_meta_path, cms_dict, gz=True, split_file=False, split_path=None, selected_owner=None): """Download report files into archive location. Args: activity (ActivityWorker): The activity worker. report_name (str): Name of the report to ingest. report_status_name (str): DynamoDB status name. date (str): Reporting date (YYYY-MM). archive_path (str): Archive location on S3. credentials_path (str): Path to credentials file. api_service_name (str): YouTube API service name. api_version (str): YouTube API version. jobs_meta_path (str): Path to jobs metadata file. cms_dict (dict): Content owner name <-> id mapping. gz (bool): Gzip output if True. split_file (bool): Split the report file if True. split_path (str): s3 path to store splitted files. selected_owner (str): Single content owner to work with provided via context. """ feed_name = report_status_name task_name = 'grab_report_files' if task_status.is_completed_task(report_status_name, date, task_name): return save_youtube_access_token_from_secrets( path=Path(credentials_path) ) youtube_reports = youtube_util.get_authenticated_services( credentials_path, api_service_name, api_version) with open(jobs_meta_path) as jobs_file: jobs = json.loads(jobs_file.read()) # check if there was single owner selected content_owners = selected_owner if selected_owner else cms_dict for content_owner_id, content_owner in content_owners.items(): job = youtube_util.get_job(content_owner_id, jobs, report_name) activity.logger.info('Downloading {} report for {}.'.format( report_name, content_owner)) try: local_dir, local_file_path = ( youtube_util.grab_reports_files_for_content_owner( content_owner, content_owner_id, job, date, youtube_reports, report_name, feed_name, gz)) activity.logger.info( f'content_owner: {content_owner}, ' f'local_dir: {local_dir}, ' f'local_file_path: {local_file_path}', ) except Exception as e: msg = ('Cannot download file from YouTube API: ' '{exception_body}'.format(exception_body=str(e))) sentry_util.send_error_or_warning(Exception(msg)) garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_AVAILABLE) activity.logger.error(msg) return {'stop': True, 'message': msg} activity.logger.info("Downloaded '{}' report to '{}'.".format( report_name, local_file_path)) report_file = '{}.{}.csv.gz'.format( report_name, content_owner) tgt_s3_bucket_name, tgt_s3_dir_key = garcon_s3.extract_bucket_path( archive_path) tgt_s3_file_key = os.path.join(tgt_s3_dir_key, report_file) # If exists delete '.tsv.gz' file from s3 s3utils.delete_s3_obj(tgt_s3_bucket_name, tgt_s3_file_key) activity.logger.info( "'{}' was succesfully deleted from '{}'".format( tgt_s3_file_key, tgt_s3_bucket_name)) # Send gzipped file to s3 temp bucket file_size = s3utils.upload_to_s3( local_file_path, tgt_s3_bucket_name, tgt_s3_file_key) if file_size is None: garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_INGESTED) return {'stop': True, 'message': "Upload to '{}' has failed".format( tgt_s3_file_key)} activity.logger.info( "'{file}' ({size}) was uploaded to '{key}'".format( file=report_file, size=file_size, key=tgt_s3_file_key)) try: if split_file: _split_file_and_upload(local_file_path, split_path) except Exception as e: garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_INGESTED) raise e finally: os.remove(local_file_path) activity.logger.info(f'temp dir content: {os.listdir(gettempdir())}') activity.logger.info( 'Deleted local temp folder and {file}'.format( file=report_file)) # mark task as completed of all content owners are processed task_status.mark_completed_task(report_status_name, date, task_name) def _split_file_and_upload(local_file_path, s3_destination_path): with TemporaryDirectory() as temp_dir: run_result = file_operations.split_file( source_file=local_file_path, destination_path=f'{temp_dir}/', chunk_size='1000m', compression_level=5, cut_header_lines=1, archiver='pigz', ) if not run_result['result']: result_stderr = str(run_result['stderr']) raise RuntimeError( f'Splitting file failed. Return_code ' f'{run_result["return_code"]}: {result_stderr[:200]}') run_result = file_operations.upload_path_to_s3( path=temp_dir, destination=s3_destination_path, num_threads=20) if not run_result['result']: result_stderr = str(run_result['stderr']) raise RuntimeError( f'Upload to s3 failed. Return_code ' f'{run_result["return_code"]}: {result_stderr[:200]}')