"""TIKTOK FRAUDULENT STREAMS REPORT Ingestion Workflow.""" from datetime import datetime from garcon import task from garcon_contrib.dynamo_feed_status import garcon_feed_status from snowflake_connector.etl_connector import SQLLoader 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.flows.tiktok_fraudulent_streams_report import config from feed_ingestion.tasks import bootstrap as reload, check_status from feed_ingestion.util import task_status from feed_ingestion.util.aws import s3 # Load SQL templates sql_loader = SQLLoader(__file__) @task.decorate(timeout=1000) @reload.reset_dynamodb_status_on_reload(config.feed_name) def bootstrap(activity, date, reload): """Bootstrap workflow by getting the correct configurations. Args: activity (ActivityWorker): The activity worker. date (str): Reporting date (YYYY-MM-DD). Returns: dict: Context. """ date_obj = ( datetime.strptime(date, '%Y-%m-%d') if date else datetime.today() ) # we will use the full date for the 'processed' timestamp (but the files \ # are shipped monthly with YYYYMM format dates) processing_date = date_obj.strftime('%Y-%m-%d') activity.logger.info('Bootstrap flow: {}'.format(date_obj)) # this is the ftp share bucket where the raw files are dropped drop_path = config.s3['drop'] # files are dropped as .zip if date_obj < datetime(2023, 6, 1): filename = config.old_filename else: filename = config.new_filename source_file_name = filename.format(date=date_obj, format='zip') s3_drop_file_domain = f's3://{config.drop_bucket}/{drop_path}' s3_drop_file_path = f'{s3_drop_file_domain}/{source_file_name}' # and converted to .gzip before being reuploaded to 'archive' new_file_name = filename.format(date=date_obj, format='gz') archive_path = config.s3['archive'].format( feed_name=config.feed_name, date=date_obj ) s3_archive_path = f's3://{config.data_bucket}/{archive_path}/' # archive is where the file is uploaded to snowflake from s3_archive_file_path = f'{s3_archive_path}{new_file_name}' return dict( feed_name=config.feed_name, date=processing_date, drop_path=drop_path, archive_path=archive_path, source_file_name=source_file_name, new_file_name=new_file_name, s3_archive_path=s3_archive_path, key_dir=s3_archive_file_path, secrets_path=config.secrets_path, staging_raw_table=config.staging_raw_table, temp_staging_raw_table=config.temp_staging_raw_table, s3_archive_file_path=s3_archive_file_path, s3_drop_file_path=s3_drop_file_path, reload=reload, ) @task.decorate(timeout=3600) @check_status(feed_name=config.feed_name) def grab_drop_files(activity, date, s3_drop_file_path, s3_archive_file_path): """Copy feed file (in zip format) from the drop location into the archive \ folder (as a gz archive that can be loaded into snowflake). Args: activity (ActivityWorker): The activity worker date (str): Reporting date (YYYY-MM-DD) s3_drop_path (str): s3 path where raw (zip) files dropped s3_rchive_path (str): s3 path to write the converted (gzip) files """ activity.logger.info( f'Convert ZIP {s3_drop_file_path} ' f'to {s3_archive_file_path}' ) activity.logger.info( f'Checking if there are any files in {s3_drop_file_path} to process ' f'- {s3.get_list_of_files_and_directories(s3_drop_file_path)}') if not s3.get_list_of_files_and_directories(s3_drop_file_path): activity.logger.info( f'No files found in {s3_drop_file_path}') return {'stop': True, 'missing_file': s3_drop_file_path} return s3.convert_zip_to_gzip_on_s3( activity, zip_s3_path=s3_drop_file_path, gz_s3_path=s3_archive_file_path, local_temp_dir='./', extract_original_filename=True, ) @task.decorate(timeout=300) def create_temp_staging_raw_table( activity, date, feed_name, temp_staging_raw_table, sfdb_params, secrets_path=None, kwargs=None): """Create temporary table in Snowflake. Args: activity (ActivityWorker): The activity worker. date (str): Reporting date (YYYY-MM-DD). feed_name (str): Name of the feed to get executor class. temp_staging_raw_table (str): Name of the temp staging_raw table. 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. kwargs (dict): Custom activity params. """ if task_status.is_completed_task(feed_name, date, 'staging_raw_table_tasks'): activity.logger.info( 'create_temp_staging_raw_table for {date} already complete, and ' '"reload" flag was not passed, skipping...'.format(date=date)) return activity.logger.info(f'Creating temp raw table {temp_staging_raw_table}') 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.drop_table(temp_staging_raw_table) activity.logger.info('{table} was dropped, if existed'.format( table=temp_staging_raw_table)) sf_executor.create_temp_staging_raw_table( temp_staging_raw_table, date) activity.logger.info( '{table} was created'.format(table=temp_staging_raw_table)) @task.decorate(timeout=600) def check_latest_report_downloaded(activity, date, feed_name, filename): """Check if this month's report is available. Args: activity (ActivityWorker): The activity worker date (str): Today's date (YYYY-MM-DD) filename (str): Expected filename feed_name (str): Feed name Returns: filename (str | None): filename if available """ if task_status.is_completed_task(feed_name, date, 'grab_drop_files'): garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_DOWNLOADED ) activity.logger.info( f'File {filename} of {feed_name} on {date} was downloaded and \ is available to ingest' ) return {'downloaded_file': filename} else: # This doesn't work for a single filename, it splits the string. # TODO: Write logic to set a single filename as missing. # garcon_feed_status.set_missing_files(feed_name, date, filename) garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_AVAILABLE ) activity.logger.warn( f'File {filename} of {feed_name} on {date} is not available to \ ingest, stopping' ) return {'stop': True, 'missing_file': filename} @task.decorate(timeout=2000) def set_status_to_ingested(activity, date, feed_name, filename): """Copy the temp_staging_raw data to the feed's staging_raw tables. Args: activity (ActivityWorker): The activity worker. feed_name (str): The name of the feed. date (str): Reporting date of the data file. downloaded_files (dict): Dictionary with files which were available. """ # This doesn't work for a single file because it splits the string # TODO: Write logic to set a single filename as ingested (if necessary) # task_status.set_values(feed_name, date, 'ingested_files', filename) if filename: activity.logger.info( f'Setting status for feed {feed_name} status for {date} to \ {garcon_feed_status.STATUS_INGESTED} ' ) garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_INGESTED ) else: activity.logger.warn( f'Setting status for feed {feed_name} status for {date} to \ {garcon_feed_status.STATUS_NOT_AVAILABLE}' ) garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_AVAILABLE )