""" Membran Snowflake-only Data Ingestion Workflow tasks. Tasks to ingest data from Membran feed into staging_raw_membran """ from datetime import date as date_module from datetime import timedelta from gzip import open as gzip_open import os import shutil from garcon import task from garcon_contrib.aws.utils import garcon_s3 from garcon_contrib.dynamo_feed_status import garcon_feed_status from feed_ingestion.flows.physical_warehouse_reports import config from feed_ingestion.tasks import check_status, s3_tasks from feed_ingestion.util import os_tools from feed_ingestion.util.aws import s3 as s3utils from feed_ingestion.util.context_util import strtobool @task.decorate(timeout=600) def bootstrap(activity, date, reload=None, use_s3='False', date_format='DD-MM-YYYY'): """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' then clear all feed statuses. use_s3 (str): If flow should backfill data from S3 only. date_format (str): Format of date in files. Returns: dict: Initial context of the workflow. """ activity.logger.info( 'Bootstrapping {feed_name}...'.format(feed_name=config.feed_name)) # date is the date passed in or yesterday's date if not date: date = (date_module.today() - timedelta(days=1)).strftime('%Y-%m-%d') feed_name = config.feed_name 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 } date_format = date_format if date_format is not None else 'DD-MM-YYYY' use_s3 = use_s3 if use_s3 is not None else 'False' use_s3 = bool(strtobool(use_s3)) source_path = config.drop_path drop_file_name = config.drop_file_name file_name_s3 = f"{drop_file_name.replace(' ', '')}.gz" staging_raw_table = config.snowflake_table_names['staging_raw'] s3_archive_bucket, s3_temp_staging_raw_bucket = \ config.s3.get('archive_bucket').format( datestamp=date, s3_bucket=config.s3_bucket), \ config.s3.get( 'temp_staging_raw_bucket').format( datestamp=date, s3_bucket=config.s3_bucket) return dict( date=date, source_files_dict={'files': [file_name_s3]}, date_as_in_uuid=date.replace('-', ''), feed_name=feed_name, secrets_path=config.secrets_path, source_path=source_path, s3_drop_bucket=config.drop_bucket, s3_archive_bucket=s3_archive_bucket, s3_temp_staging_raw_bucket=s3_temp_staging_raw_bucket, drop_file_name=drop_file_name, staging_raw_date_col='download_date', staging_raw_table=staging_raw_table, fact_analytics_table='fact_analytics', fact_analytics_error_table='fact_analytics_error', use_s3=use_s3, date_format=date_format ) @task.decorate(timeout=3000) @check_status() def grab_drop_files_s3( activity, feed_name, date, drop_file_name, source_bucket_name, source_path, destination_arch_full_path): """Copy drop file from S3 bucket to cucumbers s3.""" source_key_name = '{dir}{file}'.format( dir=source_path, file=drop_file_name) archive_file = '{dir}{file}'.format( dir=destination_arch_full_path, file=drop_file_name) destination_key_name = garcon_s3.extract_bucket_path(archive_file)[1] destination_bucket_name = garcon_s3.extract_bucket_path(archive_file)[0] result = s3_tasks.copy_file( activity, source_bucket_name, source_key_name, destination_bucket_name, destination_key_name, replace=True) file_name = destination_key_name.split('/')[-1] file_name_exists = result.get(file_name) if not file_name_exists: garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_AVAILABLE) garcon_feed_status.set_missing_files(feed_name, date, [file_name]) return {'stop': True} @task.decorate(timeout=10800) def gzip_and_put( activity, date, feed_name, source_s3_path, target_s3_path, drop_file_name): """Download a feed file from the drop location into the temp folder. Args: activity (ActivityWorker): The activity worker. date (str): Reporting date (YYYY-MM-DD). feed_name (str): Name of the feed. source_s3_path (str): Bucket from which to fetch file. target_s3_path (str): Destination bucket. drop_file_name (str): The file name. """ # Prepare downloads dir _result = [] local_dir = os_tools.create_temp_dir(feed_name=feed_name, date=date) local_file_path = os.path.join(local_dir, drop_file_name) src_s3_bucket_name, src_s3_dir_key = garcon_s3.extract_bucket_path( source_s3_path ) src_s3_file_key = os.path.join(src_s3_dir_key, drop_file_name) # Download file from s3 to local file_size = s3utils.download_from_s3( src_s3_bucket_name, src_s3_file_key, local_file_path) if file_size is None: garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_INGESTED) details = "'{file}' download from '{bucket}' bucket has failed".format( file=drop_file_name, bucket=source_s3_path) activity.logger.error(details) return {'stop': True, 'message': str(details)} # put file in gzip archive gzip_file = '{}.gz'.format(local_file_path.replace(' ', '')) _gzip_local_file(gzip_file, local_file_path) # upload to temp path tgt_s3_bucket_name, tgt_s3_dir_key = garcon_s3.extract_bucket_path( target_s3_path ) tgt_s3_file_key = os.path.join( tgt_s3_dir_key, '{}.gz'.format(drop_file_name.replace(' ', '')) ) message = f'feed_name {feed_name}, date {date}, \ gzip_file {gzip_file}, s3_path {target_s3_path}' activity.logger.info(message) s3utils.upload_to_s3(gzip_file, tgt_s3_bucket_name, tgt_s3_file_key) # cleanup local folder shutil.rmtree(local_dir) activity.logger.info( '{file} deleted local unzipped folder and original .zip'.format( file=drop_file_name)) return dict(source_files_dict={'files': _result}) def _gzip_local_file(gzip_file, local_file_path): with open(local_file_path, 'rb') as f_in: with gzip_open(gzip_file, 'wb') as f_out: shutil.copyfileobj(f_in, f_out)