"""Bulk Asset Scheduler.""" from datetime import datetime import logging import os from sys import exit as sysexit from time import sleep from constants import db from connectors.snowflake import close_connection from utils import ( sfn_utils, snowflake_utils ) import config if not config.SFN_ARN: sysexit('You must export SFN_ARN in your environment with the SFN ARN.') if not config.SFN_NAME: sysexit('You must export SFN_NAME in your environment with the SFN name.') SLEEP_DELAY = 10 # create logger with 'spam_application' logger = logging.getLogger('test') logger.setLevel(logging.DEBUG) # create console handler with a higher log level ch = logging.StreamHandler() ch.setLevel(logging.INFO) # create formatter and add it to the handlers formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') # noqa ch.setFormatter(formatter) # add the handlers to the logger logger.addHandler(ch) # Get "today" at runtime TODAY = datetime.now().date().isoformat() def main(): """Main method.""" # TODO: Check the number of instances of SFN_NAME / SFN_ARN # TODO: Ensure a similar execution is not running # Get the SFN_ARN from the name sfn_arn = \ sfn_utils.get_sfn_arn_by_name(config.SFN_NAME)['stateMachineArn'] # Get the list of execution names execution_name_list = \ sfn_utils.get_full_execution_name_list_by_arn(sfn_arn) execution_count = sfn_utils.get_full_execution_count_by_arn(sfn_arn) # Create queue tables snowflake_utils.create_table_queue() snowflake_utils.create_sfn_queue() # Look for new master tables, make sub-tables, and queue them logger.info('Looking for source ingestion tables to split.') parent_table_list = snowflake_utils.get_all_tables_to_ingest() for parent_table in parent_table_list: parent_table_name = parent_table[db.SOURCE_TABLE] logger.info(f'Table found: {parent_table_name}') release_list = \ snowflake_utils.get_release_sizes(table_name=parent_table_name) logger.info( f'{parent_table_name} contains {len(release_list)} releases.') table_dict = process_release_list(release_list, parent_table_name) logger.info(f'Creating {len(table_dict)} new tables.') for table_name, upc_list in table_dict.items(): snowflake_utils.create_asset_sub_table( upc_list=upc_list, source_table=parent_table_name, sub_table=table_name) # Queue each table in the SFN queue logger.info( f'Queueing new SFN Jobs from {parent_table_name} sub tables.') for table_name in table_dict.keys(): snowflake_utils.add_queue_item( sfn_name=config.SFN_NAME, source_table=table_name, execution_name=None, status='QUEUED' ) # Delete table from table queue logger.info(f'Removing {parent_table_name} from table queue.') snowflake_utils.delete_table_queue_item( table_name=parent_table_name) # Clean up missing executions orphaned_executions = snowflake_utils.get_orphaned_ingestions() orphan_count = len(orphaned_executions) if orphan_count: logger.info(f'{orphan_count} orphaned items need execution names.') for execution in orphaned_executions: table_name = execution[db.SOURCE_TABLE] logger.info(f'Finding execution name for {table_name}.') # Get the list of execution names execution_name_list = \ sfn_utils.get_full_execution_name_list_by_arn(sfn_arn) # Check for matching executions. for exec_name in execution_name_list: if table_name in exec_name: logger.info( f'Found execution name {exec_name} for orphaned_table ' f'{table_name}.') snowflake_utils.update_queue_field_by_table( field_name='EXECUTION_NAME', value=exec_name, # noqa table_name=table_name ) # Find in-progress SFN executions in the queue logger.info('Checking in-progress ingestions.') in_progress_queue = snowflake_utils.get_in_progress_ingestions() if len(in_progress_queue): # Iterate through statuses done_status = ['SUCCEEDED', 'FAILED', 'TIMED_OUT', 'ABORTED'] for status in done_status: finished_name_list = [] # Find finished SFN executions using AWS API. finished_name_list.extend( sfn_utils.get_full_execution_name_list_by_arn( sfn_arn, status) or []) # Compare them. Any matches should have their queue updated. for i in in_progress_queue: if i[db.EXEC_NAME] in finished_name_list: logger.info( f'Marking execution {i[db.EXEC_NAME]} completed.') snowflake_utils.update_queue_field_by_exec( field_name='END_TIME', value=datetime.now(), # noqa execution_name=i[db.EXEC_NAME] ) snowflake_utils.update_queue_field_by_exec( field_name='STATUS', value=status, # noqa execution_name=i[db.EXEC_NAME] ) # Queue each table in the SFN queue logger.info('Finding TODO items in the queue.') todo_tables = snowflake_utils.get_todo_ingestions() logger.info(f'{len(todo_tables)} items waiting in the queue.') # todo_table_names = [t[db.SOURCE_TABLE] for t in todo_tables] # Drop a JSON for each table for table in todo_tables: table_name = table[db.SOURCE_TABLE] if execution_count < config.MAX_CONCURRENT_SFN: # Example ABOVEBOARD_ASSET_ALL_BATCH02_001 file_name = f'{config.SNOWFLAKE_SCHEMA}_{table_name}' bucket = config.JSON_DROP_BUCKET prefix = config.JSON_DROP_KEY_PREFIX # Execute state machine logger.info(f'Executing state machine for: {table_name}') snowflake_utils.update_queue_field_by_table( field_name='TRIGGER_JSON_KEY', value=f's3://{os.path.join(bucket, prefix, file_name)}', # noqa table_name=table_name ) snowflake_utils.update_queue_field_by_table( field_name='START_TIME', value=datetime.now(), # noqa table_name=table_name ) snowflake_utils.update_queue_field_by_table( field_name='STATUS', value='RUNNING', # noqa table_name=table_name ) snowflake_utils.drop_json_trigger( bucket=bucket, prefix=prefix, file_name=file_name, table_name=table_name ) execution_count += 1 else: break logger.info(f'Sleeping {SLEEP_DELAY} seconds.') # Sleep while executions start sleep(SLEEP_DELAY) logger.info(f'Finding execution name for {table_name}.') # Get the list of execution names execution_name_list = \ sfn_utils.get_full_execution_name_list_by_arn(sfn_arn) # execution_count = sfn_utils.get_full_execution_count_by_arn(sfn_arn) # Check for matching executions. for exec_name in execution_name_list: if table_name in exec_name: logger.info( f'Found execution name {exec_name} for {table_name}.') snowflake_utils.update_queue_field_by_table( field_name='EXECUTION_NAME', value=exec_name, # noqa table_name=table_name ) def process_release_list(release_list, table_name): """Convert the release list into chunked tables.""" release_count = 0 asset_count = 0 total_count = 0 table_num = 1 table_dict = {} upc_list = [] for release in release_list: release_count += 1 total_count += release[db.ASSET_COUNT] if asset_count+release[db.ASSET_COUNT] > config.ASSETS_PER_TABLE: table_dict[table_name+'_'+str(table_num).zfill(3)] = upc_list upc_list = [] table_num += 1 asset_count = 0 asset_count += release[db.ASSET_COUNT] upc_list.append(release[db.UPC]) if len(upc_list): table_dict[table_name+'_'+str(table_num).zfill(3)] = upc_list upc_list = [] logger.info(f'{total_count} assets chunked.') return table_dict if __name__ == '__main__': main() close_connection()