"""Tasks for Spotify Demographics Insights pre-aggregation workflow.""" import datetime from datetime import timedelta import itertools 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.spotify_demographics import config from analytics_aggregation.util import common as common_utils from analytics_aggregation.util import snowflake_db _sql_loader = snowflake_db.SQLLoader(config.QUERY_ROOT) def _last_day_spotify_ingested(days_back): """Get last ingested day date within provided days_back period. Args: days_back (int): Number of last days for check. Returns: str: The last ingested Spotify data day or None. """ sql_template = _sql_loader.load_query('get_max_spotify_raw_data_date') with snowflake_db.SnowflakeSQLExecutor( base_config.SNOWFLAKE_MICROSERVICE_CONFIG) as sf_executor: sql, non_identifier_params = sf_executor.validator.format_identifiers( sql_template, {'days_back': days_back}) result = sf_executor.fetchone(sql, non_identifier_params) last_day = result[0] if last_day: return last_day.strftime('%Y-%m-%d') return last_day def _last_dates_spotify_ingested(days_back): """Get last ingested dates within provided days_back period. Args: days_back (int): Number of last days for check. Returns: set: The last ingested Spotify dates or empty set. """ sql_template = _sql_loader.load_query('get_latest_ingested_spotify_dates') with snowflake_db.SnowflakeSQLExecutor( base_config.SNOWFLAKE_MICROSERVICE_CONFIG) as sf_executor: sql, non_identifier_params = sf_executor.validator.format_identifiers( sql_template, {'days_back': days_back}) result = sf_executor.fetchall(sql, non_identifier_params) if result: return set(row[0].strftime('%Y-%m-%d') for row in result) return set() @task.decorate(timeout=600) def bootstrap(activity, feed_name, context_date_range=None, reload=False): """Bootstrap ETL task. Determine date range for the ingestion and checks raw data availability. Args: activity (ActivityWorker): The activity worker. feed_name (str): Name of the feed. context_date_range (str): Date range from context (context_date). reload (str|bool): Flag to force reload and ignore all the statuses. """ 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: # in case date_range provided explicitly start_date, end_date = context_date_range.split('_') # check if raw Spotify data is available last_day_ingested = _last_day_spotify_ingested( config.DAYS_BACK_FOR_LAST_SPOTIFY_INGESTION) or '1970-01-01' if end_date > last_day_ingested: return common_utils.exit_message( 'Raw Spotify available only till {}'.format(last_day_ingested)) else: # if no date_range provided 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)) return { 'reload': reload, 'date_range_as_str': start_date + '_' + end_date, 'date_range': { 'start_date': start_date, 'end_date': end_date}} def _find_range_to_process(feed_name): """Determine range of dates that need to be processing. Args: feed_name (str): Name of the feed. Returns: dict: { 'found_days': Boolean - is there any days in the range, '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') monitoring_start_date = ( datetime.date.today() - timedelta(days=config.DAYS_BACK_FOR_LAST_SPOTIFY_INGESTION) ).strftime('%Y-%m-%d') result['week_range'] = monitoring_start_date + '_' + today last_spotify_ingested = ( _last_dates_spotify_ingested( config.DAYS_BACK_FOR_LAST_SPOTIFY_INGESTION) or {'1970-01-01'}) day_range = [] for day in common_utils.list_of_dates(monitoring_start_date, today): is_day_processed = feed_status.get_overall_status( feed_name, day) == feed_status.STATUS_INGESTED if day in last_spotify_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 def _cleanup_fact_demographics(executor, date_range): """Cleanup fact_demographics for the given range. Args: executor (SnowflakeSQLExecutor): Snowflake query executor. date_range (dict[str]): Dict with 2 keys start_date and end_date. """ sql_template = _sql_loader.load_query('cleanup_fact_demographics') params = { 'schema': base_config.SNOWFLAKE_CONFIG['schema']} params.update(date_range) sql, non_identifier_params = executor.validator.format_identifiers( sql_template, params) executor.execute(sql, non_identifier_params) def _populate_fact_demographics(executor, date_range): """Populate fact_demographics with Spotify data. Args: executor (SnowflakeSQLExecutor): Snowflake query executor. date_range (dict[str]): Dict with 2 keys start_date and end_date. """ sql_template = _sql_loader.load_query('populate_fact_demographics') params = { 'schema': base_config.SNOWFLAKE_CONFIG['schema']} params.update(date_range) sql, non_identifier_params = executor.validator.format_identifiers( sql_template, params) executor.execute(sql, non_identifier_params) @task.decorate(timeout=43200) def populate_fact_demographics(activity, date_range): """Populate fact_demographics with Spotify data. In the same transaction clean the range before population Args: activity (ActivityWorker): The activity worker. date_range (dict[str]): Dict with 2 keys start_date and end_date. """ with snowflake_db.SnowflakeSQLExecutor( base_config.SNOWFLAKE_CONFIG) as sf_executor: _cleanup_fact_demographics(sf_executor, date_range) activity.logger.info('fact_demographics is cleared up.') _populate_fact_demographics(sf_executor, date_range) activity.logger.info('fact_demographics is populated.') @task.decorate(timeout=3600) def warm_up_snowflake_demographics_cache(activity, attempt_limit): """Warm up snowflake demographics cache. Args: activity (garcon.activity.Activity): Activity. attempt_limit (int): Max number of warm up requests. """ api_connector = snowflake_db.SnowflakeMetadataConnector() api_connector.authenticate() today = datetime.date.today() year_ago = today - datetime.timedelta(days=366) # Extra day for leap year date_range = { 'start_date': year_ago, 'end_date': today} params = { 'schema': base_config.SNOWFLAKE_CONFIG['schema']} params.update(date_range) sql_template = _sql_loader.load_query('get_cohorts') with snowflake_db.SnowflakeSQLExecutor( base_config.SNOWFLAKE_MICROSERVICE_CONFIG) as sf_executor: sql, non_identifier_params = sf_executor.validator.format_identifiers( sql_template, params) for i in itertools.count(start=1): if i > attempt_limit: activity.logger.warning( 'Snowflake warm up attempt limit exceeded.') return result = sf_executor.execute(sql, non_identifier_params) if api_connector.get_query_scan_bytes_number(result) == 0: activity.logger.info( 'Demographic insight data was successfully retrieved to ' 'Snowflake local storage.') break