"""Tasks of Source of Stream pre-aggregation workflow.""" import datetime from datetime import timedelta from garcon import task from garcon.contrib.dynamo_feed_status import \ feed_status_ingestion as feed_status from analytics_aggregation import base_config from analytics_aggregation.flows.apple_music_sos import utils from analytics_aggregation.flows.apple_music_sos.snowflake_executor import \ AppleExecutor from analytics_aggregation.util import common as common_utils @task.decorate(timeout=300) def extract_date( activity, feed_name, context_date_range=None, reload=False, labelids=None): """Extract date either from passed context, either from ETLs' statuses. Args: activity (ActivityWorker): The activity worker. feed_name (str): Name of the feed. context_date_range (str): Date range from context (context_date). reload (bool): Flag to force reload and ignore all the statuses. labelids (list | str | int): List of label ids. Optional. If is not None then the whole flow will be run only for these label ids. Can be list of integers or comma-separated ids or single int value. """ reload = common_utils.get_bool_from_flag(reload) if not context_date_range and reload: return common_utils.exit_message( "'reload' flag passed but date range is not specified.") if context_date_range and not reload: return common_utils.exit_message( "'explicit_date_range' specified but 'reload' flag omitted.") if context_date_range: start_date, end_date = context_date_range.split('_') else: range_result = _find_range_to_process(feed_name) if not range_result['found_days']: return common_utils.exit_message( 'There is no data to aggregate for {} range'.format( range_result['week_range'])) start_date = range_result['start_date'] end_date = range_result['end_date'] activity.logger.info( 'Start the workflow for {start_date} - {end_date}'.format( start_date=start_date, end_date=end_date)) context = { 'reload': reload, 'date_range_as_str': start_date + '_' + end_date, 'date_range': { 'start_date': start_date, 'end_date': end_date}} if labelids: try: labelids_validated = common_utils.format_id_list(labelids) except common_utils.UnprocessableFlowParamsException: raise ValueError('Incorrect label ids. Must be list of integers.') context['labelids'] = labelids_validated return context def _find_range_to_process(feed_name): """Determine range of dates that need processing within past week. Args: feed_name (str): Name of the feed. Returns: dict: { 'found_days': Boolean, 'week_range': String date range for past week, 'start_date': First date that needs processing, 'end_date': Last date that needs processing } """ result = {'found_days': False} today = datetime.date.today().strftime('%Y-%m-%d') today_minus_seven_days = ( datetime.date.today() - timedelta(days=7)).strftime('%Y-%m-%d') result['week_range'] = today_minus_seven_days + '_' + today day_range = [] for day in common_utils.list_of_dates(today_minus_seven_days, today): is_am_ingested = utils.is_apple_music_ingested(day) is_day_processed = feed_status.get_overall_status( feed_name, day) == feed_status.STATUS_INGESTED if is_am_ingested and not is_day_processed: result['found_days'] = True day_range.append(day) if result['found_days']: result['start_date'] = day_range[0] result['end_date'] = day_range[-1] return result @task.decorate(timeout=1200) def check_apple_music_ingested(activity, date_range): """Check if Apple Music report is already ingested. Args: activity (ActivityWorker): The activity worker. date_range (dict): Dict with 2 keys start_date and end_date. Returns: dict: If data is missing, return {'stop': True}. """ dates_list = common_utils.list_of_dates(**date_range) for day in dates_list: if utils.is_apple_music_ingested(day): continue else: return common_utils.exit_message( 'There is no Apple Music Streams raw data for {}'.format(day)) @task.decorate(timeout=2400) def cleanup_staging_sos(activity, date_range, labelids=None): """Delete data from staging_sos for the given period. Args: activity (ActivityWorker): The activity worker. date_range (dict): Dict with 2 keys start_date and end_date. labelids (list[int]): List of label ids. Optional. """ activity.logger.info( ('Delete Apple Music data ' 'from staging_sos for {start_date} - {end_date}').format( start_date=date_range['start_date'], end_date=date_range['end_date'])) labelids_clause, labelids_values = common_utils.sos_labelid_filter( labelids) with AppleExecutor(sf_config=base_config.SNOWFLAKE_CONFIG) as executor: executor.cleanup_staging_sos( params={ 'db': base_config.SNOWFLAKE_CONFIG['db'], 'schema': base_config.SNOWFLAKE_CONFIG['schema'], 'start_date': date_range['start_date'], 'end_date': date_range['end_date'], 'labelids': labelids_values}, labelids_clause=labelids_clause) activity.logger.info('Cleanup of staging_sos complete') @task.decorate(timeout=7200) def populate_staging_with_apple_music_data( activity, date_range, labelids=None): """Populate staging_sos table with Apple Music data. Args: activity (ActivityWorker): The activity worker. date_range (dict): Dict with 2 keys start_date and end_date. labelids (list[int]): List of label ids. Optional. """ activity.logger.info( 'Starting to populate staging_sos with Apple Music data...') labelids_clause, labelids_values = common_utils.sos_labelid_filter( labelids, 'dt') with AppleExecutor( sf_config=base_config.SNOWFLAKE_CONFIG) as executor: executor.populate_staging_sos( params={ 'db': base_config.SNOWFLAKE_CONFIG['db'], 'schema': base_config.SNOWFLAKE_CONFIG['schema'], 'start_date': date_range['start_date'], 'end_date': date_range['end_date'], 'labelids': labelids_values}, labelids_clause=labelids_clause) activity.logger.info('Apple Music data ingested to staging_sos.')