"""TikTok Weekly Data Ingestion Workflow.""" from collections import defaultdict from datetime import datetime from datetime import timedelta from itertools import product 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 import TikTokWeekly from feed_ingestion.flows.helpers import get_sf_config from feed_ingestion.flows.tiktok_weekly import config from feed_ingestion.tasks import bootstrap as reload from feed_ingestion.tasks import s3_tasks from feed_ingestion.util import task_status from feed_ingestion.util.context_util import get_context_values @task.decorate(timeout=1000) @reload.reset_dynamodb_status_on_reload(config.feed_name) def bootstrap(activity, date, dw_config=None): """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()) # if date is not Sunday, get last sunday as a date if date_obj.weekday() != 6: date_obj = date_obj - timedelta(days=date_obj.weekday() + 1) date = date_obj.strftime('%Y-%m-%d') # check overall status for this date overall_status = garcon_feed_status.get_overall_status( config.feed_name, date) if overall_status == garcon_feed_status.STATUS_INGESTED: return {'stop': True} activity.logger.info('Bootstrap flow: {}'.format(date_obj)) drop_path = config.s3['drop'].format(date=date_obj) archive_path = config.s3['archive'].format(date=date_obj) filename_template = config.filename.format(date=date_obj) return dict( feed_name=config.feed_name, date=date, drop_path=drop_path, archive_path=archive_path, filename_template=filename_template) @task.decorate(timeout=3600) def grab_drop_files( activity, feed_name, date, drop_path, archive_path, platforms, filename_template): """Download a feed file from the drop location into the archive folder. Args: activity (ActivityWorker): The activity worker. feed_name (str): The name of the feed. date (str): Reporting date (YYYY-MM-DD). drop_path (str): Source S3 path to the archive location. archive_path (str): Destination S3 path to the archive location. platforms (str): The list of platforms (e.g. 'TikTok, Douyin'). filename_template (str): The template of filename formatted with date. """ platforms_list = get_context_values(platforms, config.platforms) missing_files = [] ingested_files = task_status.get_values( feed_name, date, 'ingested_files_status') downloaded_files = defaultdict(dict) for platform, report in product(platforms_list, config.reports): filename = filename_template.format(platform=platform, report=report) if filename in ingested_files: continue source_key_name = f'{drop_path}{filename}' destination_key_name = f'{archive_path}{filename}' result = s3_tasks.copy_file_from_sme_s3_to_theocrhard( activity, secrets_path=config.sme_secrets_path, source_bucket_name=config.drop_bucket, source_key_name=source_key_name, destination_bucket_name=config.data_bucket, destination_key_name=destination_key_name, replace=True) # file is not downloaded if not result.get(filename): missing_files.append(filename) else: downloaded_files[platform][report] = filename garcon_feed_status.set_missing_files(feed_name, date, missing_files) if missing_files: garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_AVAILABLE) return {'stop': True, 'missing_files': missing_files} return { 'downloaded_files': downloaded_files, 'platforms': platforms_list} @task.decorate(timeout=2000) def load_staging_raw_table(activity, date, sfdb_params, downloaded_files): """Copy the temp_staging_raw data to the feed's staging_raw tables. Args: activity (ActivityWorker): The activity worker. date (str): Reporting date of the data file. sfdb_params (dict): Dictionary stores optional Snowflake db and schema. downloaded_files (dict): Dictionary with files which were available. """ sf_config = merge_configs(get_sf_config(config.secrets_path), sfdb_params) for platform, reports in downloaded_files.items(): for report, filename in reports.items(): temp_table_name = config.temp_staging_raw_table.format( platform=platform, report=report, date=date.replace('-', '')) staging_raw_table = config.reports[report].format( platform=platform) with TikTokWeekly(sf_config) as executor: executor.clean_staging_raw_table(staging_raw_table, date) activity.logger.info( f'{staging_raw_table} was cleaned for date {date}.') executor.load_staging_raw_table( temp_table_name, staging_raw_table, date, report=report) activity.logger.info( f'{staging_raw_table} was loaded for date {date}.') executor.drop_table(temp_table_name) activity.logger.info(f'{temp_table_name} was dropped.') @task.decorate(timeout=2000) def set_status_to_ingested( activity, feed_name, date, downloaded_files, platforms): """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. platforms (list): The list of platforms. """ ingested_files = task_status.get_values( feed_name, date, 'ingested_files_status') files = [file for platform, reports in downloaded_files.items() for report, file in reports.items()] files.extend(ingested_files) task_status.set_values(feed_name, date, 'ingested_files', files) if len(files) == len(config.reports) * len(platforms): garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_INGESTED) else: garcon_feed_status.set_overall_status( feed_name, date, garcon_feed_status.STATUS_NOT_AVAILABLE)