"""YouTube Neo4j Workflow.""" from datetime import date as date_module from garcon import task from garcon_contrib.dynamo_feed_status import garcon_feed_status from feed_ingestion.flows.helpers import get_neo4j_config from feed_ingestion.flows.helpers import get_sf_config from feed_ingestion.flows.youtube_neo4j import config from feed_ingestion.flows.youtube_neo4j.snowflake_executor import \ YouTubeNeo4J from feed_ingestion.tasks import bootstrap from feed_ingestion.util.neo4j.neo4j_executor import Neo4jExecutor STOP_RESPONSE = {'stop': True} @task.decorate(timeout=1000) @bootstrap.reset_dynamodb_status_on_reload(config.feed_name) def bootstrap(activity, date, dw_config=None): """Bootstrap workflow by injecting initial context from config. Args: activity (ActivityWorker): The Garcon activity worker. date (str): Reporting date (YYYY-MM-DD). dw_config (dict): Dictionary of data warehouse config options. Returns: dict: Initial context for the workflow. """ date = date or date_module.today().strftime('%Y-%m-%d') return { 'feed_name': config.feed_name, 'date': date} @task.decorate(timeout=1000) def check_staging_status(activity, date): """Check if all required reports is available. This function checks if all tables for 'date' have been populated based on DynamoDB statuses. Args: activity (ActivityWorker): The Garcon activity worker. date (str): Reporting date (YYYY-MM-DD). Returns: dict: {} or STOP_RESPONSE if not all reports are available. """ def _is_staging_available(date): """Check if all raw data is available. Args: date (str): Reporting date (YYYY-MM-DD). Returns: bool: If data is available. """ return all(map( lambda r: garcon_feed_status.get_overall_status(r, date) == garcon_feed_status.STATUS_INGESTED, config.report_dynamo_status_names)) if not _is_staging_available(date): activity.logger.info('Staging reports are not ready') return STOP_RESPONSE return {} @task.decorate(timeout=3600) def clear_log_table( activity, feed_name, date, reload): """Clear 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). 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') if reload == 'True': with YouTubeNeo4J(get_sf_config(config.secrets_path)) as sf_executor: params = { 'log_table_name': config.snowflake['log_table_name'], 'ingest_date': date, } sf_executor.execute_query('clear_log_table', **params) activity.logger.info('Finished clear_log_table') else: activity.logger.info('clear_log_table not required, skip') @task.decorate(timeout=16000) def ingest_youtube_data(activity, feed_name, date, query_type, temp_table): """Ingest aggregated data from snowflake into Neo4j. Args: activity (ActivityWorker): The activity worker. feed_name (str): The name of the feed. date (str): The date from which to ingest data (YYYY-MM-DD) query_type (str): Name of the query. One of (asset, channel, video) temp_table (str): Temporary table name. Returns: dict: Context. """ activity.logger.info(f'Starting youtube {query_type} ingestion') sf_config = get_sf_config(feed_name) query_name = f'ingest_youtube_{query_type}' activity.logger.info(f'Using query {query_name}') offset = 0 limit = config.row_limit_per_batch[query_type] while True: with YouTubeNeo4J(sf_config) as sf_executor: sf_params = { 'table_name': temp_table, 'offset': offset, 'limit': limit } rows = sf_executor.fetchall_dict_query( query_name, **sf_params) offset += limit if not rows: break with Neo4jExecutor(feed_name, get_neo4j_config(feed_name)) \ as neo4j_executor: neo4j_params = {'rows': rows} activity.logger.info( f'Neo4j Starting {query_name} ' f'with {len(rows)} to execution...') neo4j_executor.execute_write_query( query_name, neo4j_params) activity.logger.info(f'Neo4j Done {query_name}') activity.logger.info(f'Finished youtube {query_type} ingestion') @task.decorate(timeout=3600) def create_table(activity, feed_name, date, query_name, temp_table, full_owner_changed_refresh): """Create the table that will hold 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). query_name (str): Name of the query. temp_table (str): Temporary table name. Returns: dict: Context. """ activity.logger.info(f'Starting create_table for {query_name}') if full_owner_changed_refresh is None: full_owner_changed_refresh = False with YouTubeNeo4J(get_sf_config(config.secrets_path)) as sf_executor: params = { 'temp_table_name': temp_table, 'ar_db': config.snowflake['ar_db'], 'ar_schema': config.snowflake['ar_schema'], 'log_table_name': config.snowflake['log_table_name'], 'ingest_date': date, 'full_owner_changed_refresh': full_owner_changed_refresh } sf_executor.execute_query( f'create_temp_{query_name}_table', **params) activity.logger.info(f'Finished create_table for {query_name}') @task.decorate(timeout=7200) def insert_into_log_table( activity, feed_name, date, query_name, temp_table): """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). query_name (str): Name of the query. temp_table (str): Temporary table name. """ activity.logger.info(f'Insert into log table for {query_name}...') with YouTubeNeo4J(get_sf_config(config.secrets_path)) as sf_executor: query_name = 'insert_into_log_table' params = { 'table_name': temp_table, 'log_table_name': config.snowflake['log_table_name'], 'ingest_date': date, } sf_executor.execute_query(query_name, **params)