#!/usr/bin/env python3 """Backload Scheduler exec script. It manages tasks in backload_tasks snowflake table. """ import argparse from datetime import datetime, timedelta import json import logging import os import boto3 from feed_ingestion.bin.snowflake_executor \ import BackloadTasksExecutor from feed_ingestion.conf.config import SF_CONFIG from feed_ingestion.flows import get_flow from feed_ingestion.flows.apple_music_streams.utils import \ get_fact_analytics_report as get_apple_fact_analytics_report from feed_ingestion.flows.spotify import config as spotify_config DATE_FORMAT = '%Y-%m-%d' logging.basicConfig(level=logging.INFO) FEED_INGESTION_TABLE_NAME = os.environ.get('FEED_INGESTION_TABLE') def get_fact_analytics_report(feed_name, date, licensor): """Get fact_analytics_report for spotify and apple music etls.""" if feed_name.startswith('spotify'): return spotify_config.fact_analytics_report elif feed_name.startswith('apple_music'): return get_apple_fact_analytics_report(date, licensor) def _get_feed_name(flow_name, context): """Get feed name from flow and context.""" flow = get_flow(flow_name) feed_name = flow.contextified_feed_name(context) fact_analytics_report = get_fact_analytics_report( feed_name, context['context_date'], context['licensor']) if not fact_analytics_report: return [feed_name] fact_analytics_feed_name = f'{feed_name}_{fact_analytics_report}' return [feed_name, fact_analytics_feed_name] def schedule_tasks(flow, licensor, period_start, period_end, initial_context, reload, soft_reload): """Add tasks to table backload_tasks.""" start_date = datetime.strptime(period_start, DATE_FORMAT) assert start_date end_date = datetime.strptime(period_end, DATE_FORMAT) assert end_date if initial_context: common_context = json.loads(initial_context) else: common_context = {} if reload: common_context['reload'] = 'True' if soft_reload: common_context['soft_reload'] = 'True' common_context['licensor'] = licensor dynamodb = boto3.resource('dynamodb') dynamodb_table = dynamodb.Table(FEED_INGESTION_TABLE_NAME) with BackloadTasksExecutor(SF_CONFIG) as sf_executor: current_date = start_date while current_date <= end_date: date_str = current_date.strftime(DATE_FORMAT) context = common_context.copy() context['context_date'] = date_str if soft_reload: feed_names = _get_feed_name(flow, context) for feed_name in feed_names: existing_item_response = dynamodb_table.get_item( Key={ 'feed_name': feed_name, 'date': date_str, } ) if 'Item' in existing_item_response: status_item = existing_item_response['Item'] status_item = update_status_for_soft_reload( status_item) dynamodb_table.put_item(Item=status_item) sf_executor.insert_backload_tasks( licensor=licensor, date=current_date, flow=flow, context=json.dumps(context) ) current_date += timedelta(days=1) logging.info(f'Created task flow={flow}, context={context}') def cancel_tasks(licensor, flow, period_start, period_end): """Delete tasks from table backload_tasks.""" assert datetime.strptime(period_start, DATE_FORMAT) assert datetime.strptime(period_end, DATE_FORMAT) with BackloadTasksExecutor(SF_CONFIG) as sf_executor: sf_executor.delete_backload_tasks( licensor=licensor, flow=flow, period_start=period_start, period_end=period_end, ) def main(): """Execute main method.""" parser = argparse.ArgumentParser( description='Date Range Flow Execution Scheduler') parser.add_argument( '--licensor', default=os.environ.get('LICENSOR'), help='Licensor value') parser.add_argument( '--flow', default=os.environ.get('FLOW'), help='name of the workflow') parser.add_argument( '--period_start', type=str, default=os.environ.get('PERIOD_START'), help='Date in format YYYY-MM-DD') parser.add_argument( '--period_end', type=str, default=os.environ.get('PERIOD_END'), help='Date in format YYYY-MM-DD') parser.add_argument( '--reload', choices=['full', 'soft', 'None'], default=os.environ.get('RELOAD'), help='Reload value passed to flow context') parser.add_argument( '--cancel', action='store_true', default=bool((os.environ.get('CANCEL', '').lower() == 'true')), help='Cancel already scheduled tasks') parser.add_argument( '--context', default=os.environ.get('CONTEXT'), help='Shared initial context') args = parser.parse_args() if (not args.licensor or not args.flow or not args.period_start or not args.period_end): raise ValueError( 'Missing one of required parameters ' '(licensor,flow,period_start,period_end)') if args.cancel: cancel_tasks( licensor=args.licensor, flow=args.flow, period_start=args.period_start, period_end=args.period_end, ) else: schedule_tasks( licensor=args.licensor, flow=args.flow, period_start=args.period_start, period_end=args.period_end, initial_context=args.context, reload=args.reload == 'full', soft_reload=args.reload == 'soft', ) # TODO: remove in favour of in-feed soft-reload # functions: feed_ingestion.util.task_status.soft_reload_update_status_* def update_status_for_soft_reload(status_dynamodb_entry): """Prepare feed_ingestion status entry to soft reload. Args: status_dynamodb_entry: (dict) feed_ingestion_status table entry Returns: (dict): updated status entry ready to be inserted to the table """ result = dict(status_dynamodb_entry) if 'status' in status_dynamodb_entry: result['status'] = 'POPULATED_RAW_TABLE' completed_tasks = result['completed_tasks_status'].split(',') \ if 'completed_tasks_status' in result else [] clear_after_task = None tasks_to_remove = [] feed_name = status_dynamodb_entry.get('feed_name', '') if (feed_name.startswith('amazon_music_') or feed_name.startswith('spotify') or feed_name.startswith('apple_music') or feed_name.startswith('itunes') or feed_name.startswith('deezer') or feed_name.startswith('pandora')): # all tasks statuses from amazon_music feed: # set_overall_status_DOWNLOADED,staging_raw_table_tasks,set_overall_status_POPULATED_RAW_TABLE # ,load_aggregated_table,update_dim_tables,create_staging_fact,load_staging_fact # ,load_aggregated_skips_and_saves,load_fact_data,set_overall_status_INGESTED tasks_to_remove = [ 'create_staging_fact', 'load_staging_fact', 'load_aggregated_skips_and_saves', 'load_fact_data', 'set_overall_status_INGESTED', ] else: # all others feed_name if 'completed_tasks_status' in status_dynamodb_entry: # by default clear everything after task sets POPULATED_RAW_TABLE clear_after_task = 'set_overall_status_POPULATED_RAW_TABLE' if clear_after_task and clear_after_task in completed_tasks: last_entry_index = completed_tasks.index(clear_after_task) + 1 completed_tasks = completed_tasks[:last_entry_index] completed_tasks = [task for task in completed_tasks if task not in tasks_to_remove] result['completed_tasks_status'] = ','.join(completed_tasks) return result if __name__ == '__main__': main()