"""Tasks of the Chartmetric Charts Ingestion Workflow.""" from datetime import date from datetime import datetime from datetime import timedelta import time from typing import Dict from typing import List from typing import Optional from uuid import uuid4 from garcon import task from garcon.activity import Activity, ActivityWorker from kafka.errors import NoBrokersAvailable from kafka.errors import NodeNotReadyError from kafka.structs import TopicPartition from feed_ingestion.flows.chartmetric_charts import config from feed_ingestion.flows.chartmetric_charts.dates import date_range_to_dates from feed_ingestion.flows.chartmetric_charts.ows_charts \ import prime_chart_dates from feed_ingestion.flows.chartmetric_charts.snowflake_executor \ import SnowflakeExecutor from feed_ingestion.flows.helpers import get_sf_config from feed_ingestion.util.kafka.executor import KafkaExecutor from feed_ingestion.util.kafka.helpers import check_dlq_topic_records from feed_ingestion.util.kafka.helpers import compare_current_vs_latest_offsets from feed_ingestion.util.kafka.helpers import get_offsets_per_partiton from feed_ingestion.util.kafka.helpers import on_send_error_callback from feed_ingestion.util.kafka.helpers import on_send_success_callback from feed_ingestion.util.kafka.structs import AutoOffsetReset from feed_ingestion.util.sentry_util import send_error_or_warning @task.decorate(timeout=600) def bootstrap( activity, date, date_limit=None, platform_names=None, reload=None, days_back=None): """Bootstrap workflow by getting the correct configurations. Args: activity (ActivityWorker): The activity worker. date (str): The date to ingest from (YYYY-MM-DD). date_limit (str): The date limit to ingest to (YYYY-MM-DD). platform_names (str): A comma-separated list of platforms (e.g., "amazon,apple"). reload (str): A flag which indicates if we have to clear out the log table, and perform a force refresh. days_back (int): days back value when no 'date' arg specified. By default it is 2. Returns: dict: Context. """ # date is the date passed in or days_back's date if not platform_names: list_of_platforms = config.all_platform_names else: list_of_platforms = [] for platform_name in platform_names.split(','): if platform_name not in config.all_platform_names: raise ValueError(f'Unknown platform {platform_name}') list_of_platforms.append(platform_name) if date: date_obj = datetime.strptime(date, '%Y-%m-%d') else: if days_back: days_back = int(days_back) else: days_back = 2 date_obj = datetime.utcnow().date() - timedelta(days=days_back) date_obj_yt = date_obj - timedelta(days=2) platform_days_back = {platform: (date_obj_yt.strftime('%Y-%m-%d') if platform == 'youtube' else date_obj.strftime('%Y-%m-%d')) for platform in list_of_platforms} # date_limit is the date_limit passed in or a date very far in the future if date_limit: date_limit_obj = datetime.strptime(date_limit, '%Y-%m-%d') else: date_limit_obj = datetime.strptime('3000-01-01', '%Y-%m-%d') for key, value in platform_days_back.items(): activity.logger.info('Bootstrap flow for {} from {} to {}'.format( key, value, date_limit_obj)) return dict( feed_name=config.feed_name, date=date_obj.strftime('%Y-%m-%d'), date_limit=date_limit_obj.strftime('%Y-%m-%d'), platform_names=list_of_platforms, platform_days_back=platform_days_back, reload='True' if reload else 'False', ingestion_started_at=datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S') ) @task.decorate(timeout=3600) def clear_log_table( activity, feed_name, date, date_limit, platform_name, reload): """Create the table that will hold the spotify data to ingest. Args: activity (ActivityWorker): The activity worker. feed_name (str): The name of the feed. date (str): The date to ingest from (YYYY-MM-DD). date_limit (str): The date limit to ingest to (YYYY-MM-DD). platform_name (str): The name of the platform (e.g. spotify). reload (str): A flag which indicates if we have to clear out the log table, and perform a force refresh. Returns: dict: Context. """ activity.logger.info( 'Starting clear_log_table for {}'.format(platform_name)) if reload == 'True': with SnowflakeExecutor( get_sf_config(config.secrets_path)) as sf_executor: params = { 'platform_name': platform_name, 'log_table_name': config.log_table_name, 'ingest_date': date, 'ingest_date_limit': date_limit } sf_executor.execute_query('clear_log_table', **params) activity.logger.info( 'clear_log_table for {} finished'.format(platform_name)) elif reload == 'False': activity.logger.info( 'clear_log_table for {} not required, skip'.format(platform_name)) @task.decorate(timeout=7200) def insert_into_log_table( activity, feed_name, date, date_limit, ingestion_started_at): """Insert unique keys of ingested rows to the log table. Args: activity (ActivityWorker): The activity worker. feed_name (str): The name of the feed. date (str): The date to ingest from (YYYY-MM-DD). date_limit (str): The date limit to ingest to (YYYY-MM-DD). ingestion_started_at (str): Timestamp '%Y-%m-%d %H:%M:%S' when ETL has started. platform_name (str): The name of the platform (e.g. spotify). """ activity.logger.info( 'Insert into log table for all platforms') with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: query_name = 'insert_into_log_table' params = { 'table_name': config.STAGING_TABLE_NAME, 'log_table_name': config.log_table_name, 'ingest_date': date, 'ingest_date_limit': date_limit, 'ingestion_started_at': ingestion_started_at } sf_executor.execute_query(query_name, **params) @task.decorate(timeout=3600) def create_staging_fact(activity): """Create staging fact table that will hold the data for all platforms. Args: activity (ActivityWorker): The activity worker. Returns: dict: Context. """ with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: # (re-)create staging fact activity.logger.info('Starting create_staging_fact') params = { 'table_name': config.STAGING_TABLE_NAME, } sf_executor.execute_query('create_staging_fact', **params) sf_executor.grant_select_to_facts_db_prod_schema_read( 'STAGING_FACT_CHARTMETRIC_CHARTS') activity.logger.info('Finished create_staging_fact') @task.decorate(timeout=60 * 60) def load_staging_fact(activity, date, date_limit, platform_name): """Load fact table with data from raw tables. Args: activity (ActivityWorker): The activity worker. date (str): The date to ingest from (YYYY-MM-DD). date_limit (str): The date limit to ingest to (YYYY-MM-DD). platform_name (str): The name of the platform (e.g. spotify). Returns: dict: Context. """ activity.logger.info( 'Starting load_staging_fact for platform {}'.format(platform_name)) with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: params = { 'platform_name': platform_name, 'ingest_date': date, 'ingest_date_limit': date_limit } query_name = 'load_staging_fact_{}'.format(platform_name.lower()) sf_executor.execute_query(query_name, **params) # TODO: get number of inserted queries and raise error is it's zero activity.logger.info( 'Finished load_staging_fact for platform {}'.format(platform_name)) @task.decorate(timeout=60 * 60) def find_unchanged_charts(activity, date, date_limit, platform_name): """Find unchanged charts. Args: activity (ActivityWorker): The activity worker. date (str): The date to ingest from (YYYY-MM-DD). date_limit (str): The date limit to ingest to (YYYY-MM-DD). platform_name (str): The name of the platform (e.g. spotify). Returns: dict: Context. """ activity.logger.info( 'Starting find_unchanged_charts_ for platform {}'.format( platform_name)) with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: params = { 'platform_name': platform_name, 'ingest_date': date, 'ingest_date_limit': date_limit } query_name = 'find_unchanged_charts_{}'.format(platform_name.lower()) result = sf_executor.fetchall_query(query_name, **params) activity.logger.info( 'List of unchanged charts {} for platform {}' .format(result, platform_name)) return {'unchanged_charts': result} @task.decorate(timeout=14400) def push_metadata_to_kafka( activity: Activity, run_id: str, sf_query_name: str, kafka_topic_suffix: str, message_key: str, platforms_to_skip: List, platform_name: str) -> Dict: """Ingest data from the Snowflake table into Neo4j. Args: activity: The activity worker. run_id: The SWF execution run_id. sf_query_name: The name of the snowflake query file. kafka_topic_suffix: The suffix of the kafka topic to push into. platform_name: The name of the platform if used or None. Returns: dict: context. """ if (platform_name and platforms_to_skip and platform_name in platforms_to_skip): # skip all platforms not having labels and product info activity.logger.info('Skipping push to kafka for ' 'platform {}'.format(platform_name)) return {'skip_kafka_step': True} sf_config = get_sf_config(config.secrets_path) table_name = config.STAGING_TABLE_NAME sf_query_id = str(uuid4()) if sf_query_name == 'get_sf_productandlabel_relations': album_id_column = config.ALBUM_ID_COLUMN_MAPPING.get(platform_name) else: album_id_column = None topic_template = f'{config.KAFKA_TOPIC_PREFIX}{kafka_topic_suffix}' topic = topic_template.format(str(platform_name).lower()) activity.logger.info(f'Starting to push metadata to {topic}') latest_offsets = {} while True: try: producer = (KafkaExecutor(config.KAFKA_CLUSTER_NAME) .producer(client_id=run_id)) break except (NoBrokersAvailable, NodeNotReadyError) as err: activity.logger.warning( f'Kafka is not ready: {err}. Wait for another 60 sec.') # Wait till kafka is ready time.sleep(60) continue total_row_count = 0 with SnowflakeExecutor(sf_config) as sf_executor: sf_params = { 'table_name': table_name, 'platform': platform_name, 'album_id_column': album_id_column } result = sf_executor.fetchmany_dict_query( query_name=sf_query_name, size=config.sf_batch_size, **sf_params) for rows in result: row_count = len(rows) total_row_count += row_count activity.logger.info( f'Pushing metadata to {topic} ' f'with {row_count} rows to execute...') for row in rows: key = { message_key: str(row.get(message_key)), 'SF_QUERY_ID': sf_query_id } (producer .send(topic, key=key, value=row) .add_callback(on_send_success_callback, latest_offsets) .add_errback(on_send_error_callback, activity)) producer.flush() activity.logger.info(f'Done pushing {row_count} rows') activity.logger.info( f'Completed pushing {total_row_count} messages to {topic}') return dict( latest_offsets=latest_offsets, sf_query_id=sf_query_id) @task.decorate(timeout=8000) def check_completion_status_in_kafka( activity: Activity, run_id: str, kafka_topic_suffix: str, latest_offsets: dict, platform_name: str, skip_kafka_step: bool = False) -> None: """Get the completion event of metadata sink process from kafka. Args: activity: The executed Activity instance. run_id: The SWF execution run_id. kafka_topic_suffix: Topic suffix used in template. latest_offsets: Topic offsets o track neo4j_sink completion events. platform_name: The name of the platform (e.g. spotify). skip_kafka_step: Indicator to verify push was executed. """ # Skip waiting for consumer if there was no push step if skip_kafka_step: return start = time.perf_counter() topic_template = f'{config.KAFKA_TOPIC_PREFIX}{kafka_topic_suffix}' topic = topic_template.format(str(platform_name).lower()) activity.logger.info( 'Waiting for neo4j_sink kafka connector to process ' f'metadata in {topic}') while True: try: client = KafkaExecutor( config.KAFKA_CLUSTER_NAME).admin_client(client_id=run_id) break except (NoBrokersAvailable, NodeNotReadyError) as err: activity.logger.warning( f'Kafka is not ready: {err}. Wait for another 60 sec.') # Wait till kafka is ready time.sleep(60) continue UNFINISHED_PARTITIONS = {k for k in latest_offsets.keys()} FINISHED_PARTITIONS = set() while UNFINISHED_PARTITIONS: try: group_offsets = client.list_consumer_group_offsets( group_id=config.KAFKA_CONSUMER_GROUP ) except (NoBrokersAvailable, NodeNotReadyError) as err: activity.logger.warning( f'Kafka is not ready: {err}. Wait for another 60 sec.') # Wait till kafka is ready time.sleep(60) continue # Get sorted list of (partition, offset) namedtuples for the topic current_offsets = get_offsets_per_partiton(topic, group_offsets) # Compare current consumed latest offsets with the latest pushed offest compare_current_vs_latest_offsets( latest_offsets, current_offsets, FINISHED_PARTITIONS, UNFINISHED_PARTITIONS) # All partitions are done, break the loop if len(UNFINISHED_PARTITIONS) < 1: break # Wait for 1 minute and continue while loop # if there is unconsumed data left time.sleep(60) # Offests are beyond those sent, neo4j_sink done consuming end = time.perf_counter() activity.logger.info( f'Sink is completed for {topic} in {end - start} seconds') @task.decorate(timeout=8000) def check_for_errors_in_dlq( activity: Activity, workflow_id: str, run_id: str, sf_query_id: str, skip_kafka_step: bool = False) -> Optional[Dict]: """Get the completion event of metadata sink process from kafka. Args: activity: The executed Activity instance. workflow_id: The id of the workflow in SWF, run_id: The SWF execution run_id. sf_query_id: The synthetic id of the snowflake query. skip_kafka_step: Indicator to verify push was executed. """ # Skip DLQ checks if there was no push step if skip_kafka_step: return activity.logger.info('Starting to listen neo4j_sink DLQ topic in kafka') topic = config.KAFKA_DLQ_TOPIC while True: try: consumer = KafkaExecutor(config.KAFKA_CLUSTER_NAME).consumer( client_id=run_id, group_id=workflow_id, enable_auto_commit=False, auto_offset_reset=AutoOffsetReset.EARLIEST, ) break except (NoBrokersAvailable, NodeNotReadyError) as err: activity.logger.warning( f'Kafka is not ready: {err}. Wait for another 60 sec.') # Wait till kafka is ready time.sleep(60) continue partitions_id_set = consumer.partitions_for_topic(topic) if not partitions_id_set: activity.logger.error( f'Not able to find any partitions for: {topic}.' 'Please check that topic exists.') return {'stop': True} UNFINISHED_PARTITIONS = [] for partiotion_id in partitions_id_set: UNFINISHED_PARTITIONS.append( TopicPartition(topic, partiotion_id)) end_offsets = consumer.end_offsets(UNFINISHED_PARTITIONS) if not end_offsets: activity.logger.info( f'Have not found any records in: {topic}.' 'LGTM.') return consumer.assign(UNFINISHED_PARTITIONS) consumer.poll(timeout_ms=0) # Ensure group is rebalanced consumer.seek_to_beginning() while UNFINISHED_PARTITIONS: polled_records = consumer.poll(timeout_ms=200) result = check_dlq_topic_records( sf_query_id, polled_records, UNFINISHED_PARTITIONS, end_offsets) if result.get('stop'): record = result['record'] error_message = f'Neo4j error - {record.headers}, {record.value}' send_error_or_warning(Exception(error_message)) consumer.unsubscribe() consumer.close() activity.logger.error( 'neo4j_sink has raised an exception ' f'with Snowflake query id: {sf_query_id}') return {'stop': True} # All partitions are done, break the loop if len(UNFINISHED_PARTITIONS) < 1: break consumer.unsubscribe() # Close can produce 'Fetch to node %i failed: Cancelled' error # which is transient in current state until this PR is merged # https://github.com/dpkp/kafka-python/pull/2172/files consumer.close() @task.decorate(timeout=60 * 60) def update_isrc_to_label_participant_mapping(activity): """Update ISRC_TO_LABEL_PARTICIPANT_MAPPING table. Args: activity (ActivityWorker): The activity worker. Returns: dict: Context. """ with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: activity.logger.info( 'Starting update_isrc_to_label_participant_mapping') sf_executor.execute_query('update_isrc_to_label_participant_mapping') activity.logger.info( 'Finished update_isrc_to_label_participant_mapping') @task.decorate(timeout=60 * 60) def update_upc_to_label_participant_mapping(activity): """Update UPC_TO_LABEL_PARTICIPANT_MAPPING table. Args: activity (ActivityWorker): The activity worker. Returns: dict: Context. """ with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: activity.logger.info( 'Starting update_upc_to_label_participant_mapping') sf_executor.execute_query('update_upc_to_label_participant_mapping') activity.logger.info( 'Finished update_upc_to_label_participant_mapping') @task.decorate(timeout=60 * 60) def update_dim_tables(activity): """Update dim_tables. Args: activity (ActivityWorker): The activity worker. Returns: dict: Context. """ with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: activity.logger.info('Starting update_dim_chart') sf_executor.execute_query('update_dim_chart') activity.logger.info('Finished update_dim_tables') @task.decorate(timeout=60 * 60) def refresh_aggregated_table(activity): """Refresh aggregated charts table. Args: activity (ActivityWorker): The activity worker. Returns: dict: Context. """ with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: activity.logger.info('Starting temp_create_aggregated_table') sf_executor.execute_query('create_aggregated_table') activity.logger.info('Starting cloning aggregated_table') sf_executor.execute_query('clone_aggregated_table') activity.logger.info('Finished create_aggregated_table') @task.decorate(timeout=60 * 60) def refresh_charts_for_employee_aggregate(activity): """Refresh charts_for_employee_aggregate table. Args: activity (ActivityWorker): The activity worker. Returns: dict: Context. """ with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: activity.logger.info('Starting create_charts_for_employee_aggregate') sf_executor.execute_query('create_charts_for_employee_aggregate') activity.logger.info( 'Starting cloning create_charts_for_employee_aggregate') sf_executor.execute_query('clone_charts_for_employee_aggregate') activity.logger.info('Finished create_charts_for_employee_aggregate') @task.decorate(timeout=60 * 60) def refresh_fact_charts_filtered_recording(activity): """Refresh fact_charts_filtered_recording table. Args: activity (ActivityWorker): The activity worker. Returns: dict: Context. """ with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: activity.logger.info('Starting fact_charts_filtered_recording') sf_executor.execute_query('create_fact_charts_filtered_recording') activity.logger.info( 'Starting cloning create_fact_charts_filtered_recording') sf_executor.execute_query('clone_fact_charts_filtered_recording') activity.logger.info('Finished create_fact_charts_filtered_recording') @task.decorate(timeout=60 * 60) def update_staging_fact_sound_recording(activity): """Update staging_fact table with globalSoundRecording ID. Args: activity (ActivityWorker): The activity worker. Returns: dict: Context. """ # We need to wait until Kafka Data Highway propagates # all neo4j changes from previous task into snowflake tables WAITING_TIMEOUT = 2 * 60 activity.logger.info( f'Waiting {WAITING_TIMEOUT} seconds for neo4j to snowflake sync') time.sleep(WAITING_TIMEOUT) with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: activity.logger.info('Starting update_staging_fact_sound_recording') # we update within 2 queries because they work much fatser rather # single UPDATE query containing OR statement. sf_executor.execute_query( 'update_staging_fact_sound_recording_by_isrc') sf_executor.execute_query( 'update_staging_fact_sound_recording_by_track_id') activity.logger.info('Finished update_staging_fact_sound_recording') @task.decorate(timeout=60 * 60) def update_staging_fact_public_product(activity): """Update staging_fact table with PublicProduct ID. Args: activity (ActivityWorker): The activity worker. Returns: dict: Context. """ with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: activity.logger.info('Starting update_staging_fact_public_product') sf_executor.execute_query( 'update_staging_fact_public_product') activity.logger.info('Finished update_staging_fact_public_product') @task.decorate(timeout=60 * 60) def load_fact_data(activity, date, date_limit, platform_name): """Load fact table with data from staging_fact table. Args: activity (ActivityWorker): The activity worker. date (str): The date to ingest from (YYYY-MM-DD). date_limit (str): The date limit to ingest to (YYYY-MM-DD). platform_name (str): The name of the platform (e.g. spotify). Returns: dict: Context. """ activity.logger.info( 'Starting load_fact_data for platform {}'.format(platform_name)) with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: params = { 'platform_name': platform_name, 'ingest_date': date, 'ingest_date_limit': date_limit } activity.logger.info(f'Deleting existing data for {platform_name}...') query_name = 'delete_from_fact_charts' sf_executor.execute_query(query_name, **params) query_name = 'load_fact_charts' activity.logger.info(f'Loading new data for {platform_name}...') sf_executor.execute_query(query_name, **params) activity.logger.info( 'Finished load_fact_data for platform {}'.format(platform_name)) @task.decorate(timeout=60 * 60) def update_dim_chart_latest_chart_date(activity): """Update all dim_chart set latest_chart_date with the latest fact_charts. Args: activity (ActivityWorker): The activity worker. Returns: dict: Context. """ activity.logger.info('Start update_dim_chart_latest_chart_date') with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: query_name = 'update_dim_chart_latest_chart_date' sf_executor.execute_query(query_name) activity.logger.info('Finished update_dim_chart_latest_chart_date') @task.decorate(timeout=60 * 60) def update_position_change(activity, date, platform_days_back): """Update position_change column in fact table. Args: activity (ActivityWorker): The activity worker. Returns: dict: Context. """ activity.logger.info('Start update_position_change') if 'youtube' not in platform_days_back: youtube_ingest_date = date else: youtube_ingest_date = platform_days_back['youtube'] with SnowflakeExecutor(get_sf_config(config.secrets_path)) as sf_executor: query_name = 'update_position_change_fact_table' params = { 'ingest_date': date, 'youtube_ingest_date': youtube_ingest_date } sf_executor.execute_query(query_name, **params) activity.logger.info('Finished update_position_change') @task.decorate(timeout=60 * 60) def ows_charts_cache_prime( activity: ActivityWorker, date_limit: str, platform_days_back: Dict[str, str], ) -> None: """Prime the ows-charts cache for one platform's charts. Args: activity (ActivityWorker): The activity worker. date_limit (str): The date limit to ingest to (YYYY-MM-DD). platform_days_back (dict): A mapping of platforms to the earliest ingest date for the platform. """ activity.logger.info('Start ows_charts_cache_prime') parsed_date_limit = date.fromisoformat(date_limit) updated_chart_dates = [] sf_config = get_sf_config(config.secrets_path) for platform_name, date_from in platform_days_back.items(): with SnowflakeExecutor(sf_config) as sf_executor: query_name = 'get_platform_chart_ids' params = { 'platform': platform_name, 'date': date_from, } platform_chart_ids = sf_executor.fetchall_query( query_name, **params, ) platform_updated_chart_dates = [ { 'chart_id': chart_id, 'chart_date': chart_date.isoformat(), } for chart_id, chart_latest_date in platform_chart_ids for chart_date in date_range_to_dates( date.fromisoformat(date_from), min( chart_latest_date + timedelta(days=1), parsed_date_limit, ) ) ] activity.logger.info( f'Found {len(platform_updated_chart_dates)} {platform_name}' ' chart dates to prime.', ) updated_chart_dates.extend(platform_updated_chart_dates) if not updated_chart_dates: activity.logger.warning('No charts to prime.') return batches = list( range(0, len(updated_chart_dates), config.OWS_CHARTS_BATCH_SIZE), ) activity.logger.info( f'Priming {len(updated_chart_dates)} chart dates' f' in {len(batches)} batches.' ) for batch_id, batch_start in enumerate(batches): activity.logger.info(f'Priming Batch {batch_id + 1}.') batch_end = batch_start + config.OWS_CHARTS_BATCH_SIZE batch = updated_chart_dates[batch_start:batch_end] result = prime_chart_dates(batch) if result.status != 200: message = result.errors['message'] activity.logger.warning( f'Call to ows-charts failed ({result.status}): {message}', ) activity.logger.info('Finished ows_charts_cache_prime')