"""Ingest Bulk Upload files programmatically.""" from collections import OrderedDict from datetime import datetime from time import sleep import os import sys # for sysexit() import uuid import pandas as pd from botocore.exceptions import BotoCoreError from integration_scripts import db_utils as orch_db from integration_scripts import logger as log from integration_scripts import s3_backoff_utils as s3_utils from integration_scripts import xlsx_to_pandas_s3 as xp from integration_scripts.connectors import sentry from integration_scripts.connectors.snowflake import snow_db_config from integration_scripts.connectors.s3 import s3_bucket_name from integration_scripts.general_use import escape_bad_chars from integration_scripts.orch_utils import get_bulk_upc_from_release_ids from integration_scripts.common_config import AWS_ACCESS_KEY_ID, \ AWS_SECRET_ACCESS_KEY from constants import database as db_consts from constants import s3 as s3_const from constants import spreadsheet as sheet_consts from constants import participants as part_consts from constants.artist_fields import ARTIST_FIELDS, CLASSICAL_ARTIST_FIELDS from constants.data_headers import data_labels from constants.data_headers import data_headers, FORCED_TABLE from constants.database import RELEASE_LOG_TABLE, TRACK_LOG_TABLE from constants.globals import ( excel_line_number, keepalive_file ) from constants.formats import INTEGRATIONS_INPUT_FILE_EXTENSION, \ SME_INPUT_FILE_EXTENSION, STOPFILE_EXTENSION, SME_LABEL_COPY, BULK_UPLOAD from constants.ingestion_key_id import ingestion_key_field from constants.illegal_chars import illegal_chars from constants.mysql import RELEASE_CHUNK_REQUESTS from constants.na_values import na_values from constants.pandas import ( CONVERTORS, LEADING_ZERO_CONVERTORS ) from sql.snowflake.select_bulk_upload_labels_from_table \ import SELECT_BULK_UPLOAD_LABELS_FROM_TABLE from sql.snowflake.select_bulk_upload_table_by_orchlabelid \ import SELECT_BULK_UPLOAD_TABLE_BY_ORCHLABELID from sql.snowflake.select_bulk_upload_table_by_orchlabelid_omit_releases \ import SELECT_BULK_UPLOAD_TABLE_BY_ORCHLABELID_OMIT_RELEASES from sql.snowflake.select_bulk_upload_label_by_ingestion_key \ import SELECT_BULK_UPLOAD_LABEL_BY_INGESTION_KEY from sql.snowflake.select_all_distinct_artists_in_label import \ SELECT_ALL_DISTINCT_ARTISTS_IN_LABEL from sql.snowflake.select_all_distinct_classical_artists_in_label import \ SELECT_ALL_DISTINCT_CLASSICAL_ARTISTS_IN_LABEL from sql.snowflake.select_all_distinct_artists_in_label_resumed import \ SELECT_ALL_DISTINCT_ARTISTS_IN_LABEL_RESUMED from sql.snowflake.select_all_distinct_classical_artists_in_label_resumed \ import SELECT_ALL_DISTINCT_CLASSICAL_ARTISTS_IN_LABEL_RESUMED from sql.snowflake.select_distinct_participants \ import SELECT_DISTINCT_PARTICIPANTS from util import db_utils from util.exceptions import StopFileError import config from model.bulk_upload_tracking import update_bulk_upload_row to_ingest_key = s3_const.to_ingest_key.format( config.S3_FOLDER, config.ENVIRONMENT.lower()) been_ingested_key = s3_const.been_ingested_key.format( config.S3_FOLDER, config.ENVIRONMENT.lower()) staging_key = s3_const.staging_key.format( config.S3_FOLDER, config.ENVIRONMENT.lower()) def chunk_artists( product_data, primary_artist_name_field, new_artist_chunk_dict, withheld_artist_list, embargo_artist_list, request_data, chunk_list): """Chunk artist data into genre and embargo groups.""" # Break condition needs_chunking = True genre_field = data_headers['genre'] release_withheld_artist_list = list() for track_data in product_data: # Shorthand Artist list. primary_artist_name = \ track_data[primary_artist_name_field].strip().lower() if primary_artist_name not in new_artist_chunk_dict.keys(): new_artist_chunk_dict[primary_artist_name] = list() # Chunk first appearance of each embargoed artist separately. for check_field in ARTIST_FIELDS: artist_name_field = data_headers[check_field] # Shorthand Artist list. artist_name_list = [ a.lower().strip() for a in track_data[artist_name_field].split('|') if a != '' ] # Catch repeats within releases and across releases. for artist_name in artist_name_list: # Check if artist name in embargo list if artist_name in embargo_artist_list: # Edge case pass if check_field in CLASSICAL_ARTIST_FIELDS and \ not track_data[ genre_field].strip().lower() == 'classical': continue # if non-primary artist & primary artist is on this release if primary_artist_name in release_withheld_artist_list: # Save name in release list release_withheld_artist_list.append(artist_name) # Remove name from whole list embargo_artist_list.remove(artist_name) # Remove name # artist_name = None elif primary_artist_name not in withheld_artist_list: # Create a chunk in the new artist queue if needs_chunking: new_artist_chunk_dict[primary_artist_name].append( request_data.copy()) needs_chunking = False withheld_artist_list.append(artist_name) release_withheld_artist_list.append(artist_name) embargo_artist_list.remove(artist_name) # Remove name else: if needs_chunking: new_artist_chunk_dict[primary_artist_name].append( request_data.copy()) needs_chunking = False release_withheld_artist_list.append( artist_name) release_withheld_artist_list.append( primary_artist_name) # Magic embargo_artist_list.remove(artist_name) # Reset clean # artist_name = None withheld_artist_list = list( set(withheld_artist_list + release_withheld_artist_list) ) # No embargoed artist match if needs_chunking: # Create a chunk and reset the container chunk_list.append(request_data.copy()) return withheld_artist_list.copy(), \ embargo_artist_list.copy(), \ new_artist_chunk_dict.copy(), \ chunk_list.copy() @sentry.sentry_wrap def chunk_requests(bulk_upload_rows, embargo_artist_list): """Collate db response into chunked dict for requests. Chunk artists for Bulk Upload attempt embargoing. Args: bulk_upload_rows (list): A list of dicts to collate. embargo_artist_list (list): A list of artists Returns: list: A collated list of dicts """ # Init vars project_code = None product_code = None request_data = OrderedDict() product_data_list = list() # the list of tracks for a project row_number = 2 # Excel spreadsheets start at row 2 total_rows = 0 chunk_list = list() # each chunk is a request_data track_data_dict = dict() # track data in a dictionary withheld_artist_list = list() # New artist list new_artist_chunk_dict = dict() # for check_name in artist_fields: # new_artist_chunk_dict[check_name] = list() # Shorthand artist field primary_artist_name_field = data_headers['release_artists_primary_artist'] # Iterate over DB result rows for row in bulk_upload_rows: # each row (multiple projects) # Checks for initial app state if project_code is None: project_code = row[db_consts.project_code] if product_code is None: product_code = row[db_consts.product_code] # Check for change in product_code values if (product_code and product_code != row[db_consts.product_code]) or \ (project_code and project_code != row[db_consts.project_code]): # Store finished product request_data['{}_____{}'.format( project_code, product_code)] = product_data_list # Chunk based on artist withheld_artist_list, embargo_artist_list, new_artist_chunk_dict, \ chunk_list = chunk_artists( product_data_list, primary_artist_name_field, new_artist_chunk_dict, withheld_artist_list, embargo_artist_list, request_data, chunk_list) new_art_len = sum( [len(a) for k, a in new_artist_chunk_dict.items()]) release_count = len(chunk_list) + new_art_len # Debug log.info('Release {} prepared.'.format(release_count)) total_rows += row_number row_number = 2 # Reset spreadsheet request_data = OrderedDict() # start a new request # Reset product dict product_data_list = list() # Update project and product codes product_code = row[db_consts.product_code] project_code = row[db_consts.project_code] # Loop through columns and populate request_data fields for data_col_name, data_col_val in data_headers.items(): if data_col_name in row.keys(): track_data_dict[data_col_val] = row[data_col_name] # For debug purposes ----------------------------------------- # Repair values inline if config.REPAIR_VALUES: pass # For ad-hoc field value substitution # For debug purposes ----------------------------------------- # Data Repair ------------------------------------------------ # Repair dates # if 'date' in data_col_name and row[data_col_name] != '': # date = datetime.datetime.strptime( # row[data_col_name], '%M/%d/%Y') # track_data_dict[data_col_val] = \ # datetime.datetime.strftime(date, '%Y-%M-%d') # Escape bad chars in product_code if data_col_name == db_consts.product_code: track_data_dict[data_col_val] = escape_bad_chars( row[data_col_name]) # Data Manipulation ------------------------------------------ # Ensure 3rd Party Publisher is not blank if data_col_name == db_consts.third_party_publisher and \ row[data_col_name] != 'Yes': track_data_dict[data_col_val] = 'No' # Remove MR Ownership info if config.SEND_BLANK_MR_OWNERSHIP and data_col_name == \ db_consts.ownership_for_this_sound_recording: track_data_dict[data_col_val] = '' # Remove all lyrics if config.STRIP_LYRICS and \ data_col_name == db_consts.track_lyrics: track_data_dict[data_col_val] = None else: log.warning( 'DB field \'{}\' not in returned row.'.format( data_col_name)) # Add row number field track_data_dict[excel_line_number] = row_number # Append product to project product_data_list.append(track_data_dict.copy()) # Reset track track_data_dict = dict() row_number += 1 # Chunking - for labels with fewer rows than 1 chunk's worth if row_number > len(bulk_upload_rows) + 1: # Append product to project # Store finished product request_data['{}_____{}'.format( project_code, product_code)] = product_data_list withheld_artist_list, embargo_artist_list, new_artist_chunk_dict, \ chunk_list = chunk_artists( product_data_list, primary_artist_name_field, new_artist_chunk_dict, withheld_artist_list, embargo_artist_list, request_data, chunk_list) new_art_len = sum( [len(a) for k, a in new_artist_chunk_dict.items()]) release_count = len(chunk_list) + new_art_len log.info('Release {} prepared.'.format(release_count)) # else: # log.log(25, 'Chunk {} created.'.format(len(chunk_list))) request_data = None product_data_list = None # Chunking - Append final product to request if more than one chunk. if request_data or product_data_list: request_data[ '{}_____{}'.format(project_code, product_code)] = \ product_data_list # Checking for embargoed_artist here may be overkill due to # logical constraints. Can the last chunk ever be a dupe artist risk? withheld_artist_list, embargo_artist_list, new_artist_chunk_dict, \ chunk_list = chunk_artists( product_data_list, primary_artist_name_field, new_artist_chunk_dict, withheld_artist_list, embargo_artist_list, request_data, chunk_list) new_art_len = sum( [len(a) for k, a in new_artist_chunk_dict.items()]) release_count = len(chunk_list) + new_art_len log.info('Release {} prepared.'.format(release_count)) # else: # log.log(25, 'Chunk {} created.'.format(len(chunk_list))) # End Chunking Edge Conditions --------------------------------------- if len(embargo_artist_list): sys.exit( 'Something did not get queued: {}'.format(embargo_artist_list)) return new_artist_chunk_dict, chunk_list @sentry.sentry_wrap def chunk_with_token(access_token, chunk_list): """Add token to chunked data shape.""" label_chunk_list = list() # Init list of requests corpus = dict() for chunk_num, chunk in enumerate(chunk_list, start=1): corpus['access_token'] = access_token corpus['data'] = chunk # each a request_data dict label_chunk_list.append(corpus.copy()) corpus = dict() return label_chunk_list # End Looping over labels -------------------------------------------- @sentry.sentry_wrap def get_input_label_list(table_name=None): """Get the list of input labels.""" # Use label list if passed if config.BULK_UPLOAD_LABELS: # To limit which labels are checked return_label_list = config.BULK_UPLOAD_LABELS else: if not table_name: table_name = config.SNOWFLAKE_SOURCE_TABLE # Default value key_field = ingestion_key_field # Use a given single ingestion_key_id if config.INGESTION_KEY_VALUE: if config.INGESTION_KEY_FIELD: # Update default value if provided key_field = config.INGESTION_KEY_FIELD params = { 'ingestion_key_value': config.INGESTION_KEY_VALUE, 'ingestion_key_field': key_field, 'table_name': table_name, 'snowflake_database': snow_db_config['database'], 'snowflake_schema': snow_db_config['schema'], } # Pull Phonofile metadata from bulk upload view fetched_rows = db_utils.get_snowflake_results( sql=SELECT_BULK_UPLOAD_LABEL_BY_INGESTION_KEY, params=params) if not len(fetched_rows): sys.exit('No release matches key - {}: {}.'.format( key_field, config.INGESTION_KEY_VALUE)) else: params = { 'table_name': table_name, 'snowflake_database': snow_db_config['database'], 'snowflake_schema': snow_db_config['schema'] } fetched_rows = db_utils.get_snowflake_results( sql=SELECT_BULK_UPLOAD_LABELS_FROM_TABLE, params=params) if not len(fetched_rows): sys.exit( 'No row matches found in ' 'SNOWFLAKE_SOURCE_TABLE: {}.'.format(table_name)) if not len(fetched_rows): sys.exit('No labels found.') return_label_list = [item['orchlabelid'] for item in fetched_rows] return return_label_list # , omit_release_list @sentry.sentry_wrap def get_all_input_artist_names(source_table, label_id, session_id): """Return a set of all distinct artists for a single label in a single table. Args: source_table: (str) The table to query. label_id: (int) The label_id to query. session_id: (uuid) The session_id to query. Returns: set: A set of all distinct artists for a single label in a single table. """ # Create param dict and init repeated val params = { 'table_name': source_table, 'label_id': label_id, 'snowflake_database': snow_db_config['database'], 'snowflake_schema': snow_db_config['schema'], } log_table = TRACK_LOG_TABLE.upper() # Create return list distinct_artist_list = list() # Iterate over list of artist fields to check for field in ARTIST_FIELDS: # Query Distinct field values params['field_name'] = field # Split query logic based on whether we are resuming if config.RESUME_SESSION_ID: params = { **params, 'ingestion_key_field': ingestion_key_field, 'log_table': log_table, 'session_id': session_id } if field in CLASSICAL_ARTIST_FIELDS: sql = SELECT_ALL_DISTINCT_CLASSICAL_ARTISTS_IN_LABEL_RESUMED else: sql = SELECT_ALL_DISTINCT_ARTISTS_IN_LABEL_RESUMED else: if field in CLASSICAL_ARTIST_FIELDS: sql = SELECT_ALL_DISTINCT_CLASSICAL_ARTISTS_IN_LABEL else: sql = SELECT_ALL_DISTINCT_ARTISTS_IN_LABEL artist_list = db_utils.get_snowflake_results( sql=sql, params=params) artist_list = [a[field] for a in artist_list] all_artists_list = [] for artist in artist_list: if '|' in artist: for a in artist.split('|'): test_artist = a.strip() if test_artist: all_artists_list.append(test_artist.lower()) else: test_artist = artist.strip() if test_artist: all_artists_list.append(test_artist.lower()) distinct_artist_list = distinct_artist_list + all_artists_list return set(distinct_artist_list) def get_label_rows(source_table, lbl_id): """Return bulk upload rows for given label_id.""" # Pull label metadata from source_table label_rows = db_utils.get_snowflake_results( sql=SELECT_BULK_UPLOAD_TABLE_BY_ORCHLABELID, params={ 'label_id': lbl_id, 'ingestion_key_field': ingestion_key_field, 'table_name': source_table, 'snowflake_database': snow_db_config['database'], 'snowflake_schema': snow_db_config['schema'], }) return label_rows def get_file_participant_rows(source_table): """Return participant info from a single source file.""" participant_params = { 'source_db': part_consts.source_db, 'source_schema': part_consts.source_schema, 'source_table': source_table, 'participant_db': part_consts.participant_db, 'participant_schema': 'PROD' if config.ENVIRONMENT.lower() == 'prod' else 'QA', 'release_participant_table': part_consts.release_participant_table, 'track_participant_table': part_consts.track_participant_table, 'reporting_db': part_consts.reporting_db, 'reporting_schema': part_consts.reporting_schema, 'reporting_table': part_consts.reporting_table, } # Pull label metadata from source_table label_rows = db_utils.get_snowflake_results( sql=SELECT_DISTINCT_PARTICIPANTS, params=participant_params) return label_rows def get_ingested_upcs(session_id): """Return a list of upc's derived from an input table.""" release_id_list = db_utils.get_logged_release_ids(session_id) upc_list = [] # Chunk art relations requests to prevent deadlocks. for i in range(0, len(release_id_list), RELEASE_CHUNK_REQUESTS): upc_list += get_bulk_upc_from_release_ids( release_id_list=release_id_list[0 + i:i + RELEASE_CHUNK_REQUESTS]) return list(OrderedDict.fromkeys(upc_list)) def get_sme_project_names(session_id): """Return a list of project_names and project_id's derived from an input upc list. Args: session_id: (uuid) A UUID which correlates to a specific input set Returns: dict: A dict of project_id's as keys and project_names as values. """ # Get product_codes, project_id's from product_code_project_id_list = db_utils.get_logged_product_and_project( session_id) # Make product_code: product_id map product_project_dict = { x.product_code: x.project_id for x in product_code_project_id_list } # Strip only product_codes product_code_list = product_project_dict.keys() # Get project_names from product_codes results = db_utils.get_sme_project_names_from_product_codes( product_code_list=product_code_list) # Make dict of product_code: project_name product_code_project_name_dict = { x.product_code: x.project_name for x in results } output_dict = {} # Map project_id's to project_names for k, v in product_project_dict.items(): if not v: log.warning('Product code {} has no project_id.'.format(k)) continue output_dict = { **output_dict, v: product_code_project_name_dict.get(k) } return output_dict def resume_label_rows(source_table, lbl_id, ses_id): """Return unprocessed bulk upload rows from in-progress ingestion.""" log_table = RELEASE_LOG_TABLE.upper() label_rows = db_utils.get_snowflake_results( sql=SELECT_BULK_UPLOAD_TABLE_BY_ORCHLABELID_OMIT_RELEASES, params={ 'label_id': lbl_id, 'ingestion_key_field': ingestion_key_field, 'log_table': log_table, 'session_id': ses_id, 'table_name': source_table, 'snowflake_database': snow_db_config['database'], 'snowflake_schema': snow_db_config['schema'], }) return label_rows def get_response_length(response_dict): """Return the summed length of all rows in the response dict. """ return sum([ len(o) for o in [ y for y in [ n for a, n in response_dict.items()]]]) def backup_key(bucket_name, source_key, target_key): """Backup keys and collect source keys for deletion.""" try: # Copy to backup s3_utils.copy_from_backoff(bucket_name, source_key, target_key) except BotoCoreError: # Copy fails return False else: return True def backup_keys_in_list(bucket_name, source_list, target_key): """Backup keys and collect source keys for deletion.""" success_list = list() failure_list = list() for key in source_list: target = s3_utils.get_s3_file_key(target_key, key) if backup_key(bucket_name, key, target): success_list.append(key) else: failure_list.append(key) return success_list.copy(), failure_list.copy() def get_session_id(): """Return an existing or new session id """ if config.RESUME_SESSION_ID: session_id = config.RESUME_SESSION_ID else: session_id = str(uuid.uuid4().hex) return session_id def get_ingestion_file_ext(): """Return the appropriate file extension for the current INPUT_FORMAT""" if config.INPUT_FORMAT == SME_LABEL_COPY: file_ext = SME_INPUT_FILE_EXTENSION elif config.INPUT_FORMAT == BULK_UPLOAD: file_ext = INTEGRATIONS_INPUT_FILE_EXTENSION elif not config.INPUT_FORMAT: raise EnvironmentError( 'INPUT_FORMAT must be set to begin ingestion.') else: raise EnvironmentError( 'INPUT_FORMAT \'{}\' is not recognized.'.format( config.INPUT_FORMAT)) return file_ext @sentry.sentry_wrap def load_s3_bucket(copy_file_count=10): """Prepare a bucket for preprocessing by loading files if needed. Args: file_count: (int) The number of files to move at once. """ msg = 'Checking for files in {}/{}'.format(s3_bucket_name, to_ingest_key) log.info(msg) # Check the count of files in a bucket. input_files = s3_utils.filter_file_keys( to_ingest_key, s3_bucket_name, file_ext=get_ingestion_file_ext()) input_file_count = len(input_files) # if files, jump out. if input_file_count: log.info('{} files found!'.format(input_file_count)) return msg = 'No files found. Checking for files in {}/{}'.format( s3_bucket_name, staging_key) log.info(msg) # Look for staging files. staging_files = s3_utils.filter_file_keys( staging_key, s3_bucket_name, file_ext=get_ingestion_file_ext()) staged_file_count = len(staging_files) # Move any files from staging if staged_file_count: file_count = min(staged_file_count, copy_file_count) log.info( 'Moving {} files from staging to ingestion folder.'.format( file_count)) for i in range(0, file_count): s3_utils.move_backoff( os.path.join(staging_key, staging_files[i]), os.path.join(to_ingest_key, staging_files[i]), s3_bucket_name) else: log.warning('No files to ingest.') raise FileNotFoundError('No files to ingest.') @sentry.sentry_wrap def preprocess_files(): """Return the programmatically created table name or a provided name.""" table_file_index = dict() time_now = datetime.strftime(datetime.now(), '%Y_%m_%d_%H_%M_%S') ingest_path = to_ingest_key files_only = s3_utils.filter_file_keys( ingest_path, s3_bucket_name, file_ext=get_ingestion_file_ext()) # Warn on empty S3 metadata folder if not len(files_only): log.warning('There are no files in {}', s3_bucket_name) else: # Ingest from S3 table_file_index = get_source_table_names(files_only, time_now) return table_file_index, time_now @sentry.sentry_wrap def check_for_stopfile(): """Check for the existence of a stopfile.""" ingest_path = to_ingest_key stopfiles_only = s3_utils.filter_file_keys( ingest_path, s3_bucket_name, file_ext=STOPFILE_EXTENSION) if len(stopfiles_only): msg = 'Stopfile found! Terminating' log.warning(msg) raise StopFileError(msg) @sentry.sentry_wrap def preprocess_table(): """Return the programmatically created table name or a provided name.""" # stage_name = None table_file_index = dict() time_now = datetime.strftime(datetime.now(), '%Y_%m_%d_%H_%M_%S') # Raise exception on missing value if not len(config.SNOWFLAKE_SOURCE_TABLE): raise EnvironmentError( 'SNOWFLAKE_SOURCE_TABLE must be set when S3_INGEST is false.') # Check for illegal chars in given table name for char in illegal_chars: # Get filename without illegal chars if char in config.SNOWFLAKE_SOURCE_TABLE: raise EnvironmentError( 'SNOWFLAKE_SOURCE_TABLE must not contain any of ' '[\'{}\']'.format('\', \''.join(illegal_chars))) table_file_index[FORCED_TABLE] = { 'table_name': config.SNOWFLAKE_SOURCE_TABLE, 'code': generate_file_code(), 'label_tracking': dict() } return table_file_index, time_now @sentry.sentry_wrap def get_bulk_upload_rows(session_id, table_name, label_id): """Return a list of labels with active rows to ingest.""" if not config.RESUME_SESSION_ID: label_bulk_upload_rows = get_label_rows(table_name, label_id) else: log.info('Continuing session: {}'.format(session_id)) label_bulk_upload_rows = resume_label_rows( table_name, label_id, session_id) return label_bulk_upload_rows def generate_file_code(): """Generate a new file code.""" return str(uuid.uuid4().hex).upper()[:6] def get_source_table_names(files_only, start_time): """Ingest each file to a separate table.""" source_table_name_dict = dict() time_now = start_time # Loop through leaf keys for k in files_only: # Strip file code for SME_LABEL_COPY; for others, generate a new one if config.INPUT_FORMAT == 'SME_LABEL_COPY': code = os.path.splitext(k)[0][-6:] else: code = generate_file_code() s3_table_name = s3_const.s3_source_template + '{}_{}'.format( code, time_now) source_table_name_dict = { **source_table_name_dict, k: { 'table_name': s3_table_name, 'code': code, 'label_tracking': dict() } } return source_table_name_dict def convert_s3_file_to_table(xlsx_key, s3_stage_name, s3_table_name): """Ingest a single file into a source table""" schema = snow_db_config['schema'] db = snow_db_config['database'] ingest_path = to_ingest_key # Create fully qualified keys source_key = s3_utils.get_s3_file_key(ingest_path, xlsx_key) # Read XLSX to dataframe log.info('Converting file to DataFrame: {}.', source_key) if config.INPUT_FORMAT == 'SME_LABEL_COPY': dframe = xp.read_csv_to_dataframe_s3( source_key, s3_bucket_name, override_na_values=na_values) else: # Create converters to retain leading zeroes for relevant fields if config.RETAIN_LEADING_ZEROES: converters = { **LEADING_ZERO_CONVERTORS, **CONVERTORS, } else: converters = CONVERTORS dframe = xp.read_xlsx_to_dataframe_s3( source_key, s3_bucket_name, override_na_values=na_values, converters=converters) log.info('{} converted to Dataframe.', source_key) if config.REMOVE_NAN: # Remove all nan values dframe.fillna('', inplace=True) # Get all columns file_columns = [n for n in dframe.columns] # IF labelID field is not in the dframe columns - throw error if sheet_consts.orchlabelid not in file_columns: log.error('{} not found in Excel header row.'.format( sheet_consts.orchlabelid)) raise RuntimeError('Missing Label ID in Header.') # -- Prune irrelevant fields. -------------------------------------------- # Make a comparison map drop_map = {x.lower(): x for x in file_columns} # Normalize the comparison lists to lowercase file_columns_lower = [x.lower() for x in file_columns] data_labels_lower = [x.lower() for x in data_labels] # Find the bad columns drop_labels = list(set(file_columns_lower) - set(data_labels_lower)) # Restore the original cases drop_labels = [v for k, v in drop_map.items() if k in drop_labels] # Prune if len(drop_labels): log.info('Pruning irrelevant DataFrame columns.') dframe.drop(columns=drop_labels, inplace=True) log.info('DataFrame pruned.') if config.STRIP_WHITESPACE: # Strip all whitespace from eligible fields. cols = dframe.select_dtypes(['object']).columns dframe[cols] = dframe[cols].apply(lambda x: x.astype(str).str.strip()) # Add extra needed columns log.info('Adding order and completion state columns.') dframe.insert(0, 'ROW_NUMBER', range(1, 1 + len(dframe.index))) dframe.insert(len(dframe.columns), 'PROGRESS', 0) # Fix potential value issues log.info('Ensuring UPC\'s are proper variable type.') # Shorten var names upc_field = sheet_consts.digital_upc man_upc_field = sheet_consts.manufacturers_upc # Conform values to column spec # splitext() removes decimals if excel has added them in number formatting # Prevent casting MANUFACTURER_UPC to int (and removing leading zeroes if config.RETAIN_LEADING_ZEROES: dframe[upc_field] = dframe[upc_field].astype('str').apply( lambda x: os.path.splitext(x)[0]).apply( lambda x: '' if x in ['nan', ''] else x) dframe[man_upc_field] = dframe[man_upc_field].astype('str').apply( lambda x: os.path.splitext(x)[0]).apply( lambda x: '' if x in ['nan', ''] else x) else: dframe[upc_field] = dframe[upc_field].astype('str').apply( lambda x: os.path.splitext(x)[0]).apply( lambda x: '' if x in ['nan', ''] else '%.i' % int(x)) dframe[man_upc_field] = dframe[man_upc_field].astype('str').apply( lambda x: os.path.splitext(x)[0]).apply( lambda x: '' if x in ['nan', ''] else '%.i' % int(x)) # Pad short UPC's (less than 12 chars) if config.PAD_DIGITAL_UPC: dframe[upc_field] = dframe[upc_field].astype( 'str').apply(lambda x: x.zfill(12) if len(x) > 9 else x) dframe[man_upc_field] = dframe[man_upc_field].astype( 'str').apply(lambda x: x.zfill(12) if len(x) > 9 else x) # Sort if valid sort order is passed. if type(sheet_consts.sort_field_list) in [str, list]: sort_string = sheet_consts.sort_field_list if isinstance(sort_string, list): sort_string = ', '.join(sort_string) log.info('Sorting DataFrame by : {}.', sort_string) dframe.sort_values(by=sheet_consts.sort_field_list, inplace=True) log.info('DataFrame sorted.') output_file_key = os.path.splitext(os.path.split(source_key)[1])[0] output_path_and_key = s3_utils.get_s3_file_key( ingest_path, output_file_key) csv_file_key = xp.convert_dataframe_to_csv_s3( dframe, output_path_and_key, sheet_consts.release_name) log.info('Creating source table: {}.'.format(s3_table_name)) # Create Source Table db_utils.create_s3_source_table(s3_table_name=s3_table_name) # Copy into Source Table log.info('Copying rows from {} to {}.'.format( s3_stage_name, s3_table_name)) orch_db.copy_into_table_from_s3(db=db, schema=schema, s3_stage_name=s3_stage_name, s3_table_name=s3_table_name) log.info('Deleting CSV: {}.'.format(csv_file_key)) s3_utils.delete_backoff(csv_file_key, s3_bucket_name) @sentry.sentry_wrap def ingest_file(xlsx_key, s3_stage_name, s3_table_name): """Stage and load files from given bucket name.""" if not s3_bucket_name or not config.SNOWFLAKE_FILE_FORMAT: raise RuntimeError('S3_BUCKET and SNOWFLAKE_FILE_FORMAT must be set ' 'in the environment.') db = snow_db_config['database'] schema = snow_db_config['schema'] bucket = s3_bucket_name file_format = config.SNOWFLAKE_FILE_FORMAT aws_key = AWS_ACCESS_KEY_ID aws_secret = AWS_SECRET_ACCESS_KEY ingest_path = to_ingest_key log.info('Creating S3 ingestion file format....') # Check and Create BULK_UPLOAD_FILE_FORMAT db_utils.create_file_format(format_type='ingest', db=db, schema=schema) # Check and Create BULK_UPLOAD_TEMPLATE log.info('Checking existence of the Bulk Upload template table....') db_utils.create_upload_template() log.info('Creating S3 stage....') # Create S3 Stage db_utils.create_s3_stage(bucket=bucket, path=ingest_path, db=db, schema=schema, s3_stage_name=s3_stage_name, file_format=file_format, aws_key=aws_key, aws_secret=aws_secret) convert_s3_file_to_table(xlsx_key, s3_stage_name, s3_table_name) # Drop Stage log.info('Dropping S3 stage.') orch_db.drop_s3_stage(db=db, schema=schema, s3_stage_name=s3_stage_name) def move_ingested_files(file_ext=None): """Move files from the ingested location to a backup location.""" # Delete CSV's try: s3_utils.delete_all_files_in_folder( to_ingest_key, s3_bucket_name, 'csv') except Exception as e: log.info( 'Failed to delete CSV files in {}: {}'.format( to_ingest_key, str(e))) # Move files try: s3_utils.move_all_files_in_folder( to_ingest_key, been_ingested_key, s3_bucket_name, file_ext) except Exception as e: log.info( 'Failed to move files in {}: {}'.format(to_ingest_key, str(e))) @sentry.sentry_wrap def move_ingested_file(file_name): """Move a single file from the ingested location to a backup location.""" log.info('Moving ingested S3 files to {}....'.format(been_ingested_key)) source_key = s3_utils.get_s3_file_key(to_ingest_key, file_name) backup_target_key = s3_utils.get_s3_file_key(been_ingested_key, file_name) # Copy backup file try: s3_utils.copy_from_backoff( s3_bucket_name, source_key, backup_target_key) except Exception as e: log.info( 'Failed to move files in {}: {}'.format(to_ingest_key, str(e))) # Delete File try: s3_utils.delete_backoff(source_key, s3_bucket_name) except Exception as e: log.info( 'Failed to delete CSV files in {}: {}'.format( to_ingest_key, str(e))) @sentry.sentry_wrap def copy_ingested_file(file_name, env, session_id, code): """Move a single file from the ingested location to a backup location.""" log.info('Copying ingested S3 files to output folder') root_file_name = os.path.splitext(file_name)[0] for char in illegal_chars: # Get filename without illegal chars root_file_name = root_file_name.replace(char, '_') sig = root_file_name + '_' + code source_key = s3_utils.get_s3_file_key(to_ingest_key, file_name) output_key = s3_const.output_key.format(config.S3_FOLDER) output_target_key = s3_utils.get_s3_file_key( output_key, env, session_id, sig, file_name) # Copy output file try: s3_utils.copy_from_backoff( s3_bucket_name, source_key, output_target_key) except Exception as e: log.info( 'Failed to move files in {}: {}'.format(to_ingest_key, str(e))) log.info('Input file copied to {}'.format(output_target_key)) def unload_result_csv(session_id, input_file_root, s3_stage_name, s3_table_name, code, label_id=None): """Joins the ingestion data set to the ingestion log, and creates a CSV.""" track_log_table = TRACK_LOG_TABLE.upper() if label_id: db_utils.join_log_and_unload_label_data( session_id=session_id, s3_stage_name=s3_stage_name, s3_table_name=s3_table_name, input_file_root=input_file_root, label_id=label_id, track_log_table=track_log_table, code=code) else: db_utils.join_log_and_unload_all_data( session_id=session_id, s3_stage_name=s3_stage_name, s3_table_name=s3_table_name, input_file_root=input_file_root, track_log_table=track_log_table, code=code) def unload_reporting_csv(session_id, input_file_root, source_table_name, s3_stage_name, code, start_time, label_id=None): """Joins the ingestion data set to the ingestion log, and creates a CSV.""" release_log_table = RELEASE_LOG_TABLE.upper() track_log_table = TRACK_LOG_TABLE.upper() if label_id: return db_utils.join_log_and_unload_label_reports( session_id=session_id, s3_stage_name=s3_stage_name, input_file_root=input_file_root, label_id=label_id, release_log_table=release_log_table, track_log_table=track_log_table, code=code ) else: return db_utils.join_log_and_unload_all_reports( session_id=session_id, s3_stage_name=s3_stage_name, input_file_root=input_file_root, source_table_name=source_table_name, release_log_table=release_log_table, track_log_table=track_log_table, code=code, start_time=start_time ) def export_label_results(session_id, table_name, label_id, file_name, code, start_time): """Export the output of the run to XLSX files.""" # Create output Stage bucket = s3_bucket_name path = s3_const.output_key.format(config.S3_FOLDER) source_schema = snow_db_config['schema'] source_db = snow_db_config['database'] file_format = config.SNOWFLAKE_FILE_FORMAT_OUTPUT aws_key = AWS_ACCESS_KEY_ID aws_secret = AWS_SECRET_ACCESS_KEY env = config.ENVIRONMENT log.info( 'Exporting results and reports for {} - {}....'.format( label_id, file_name)) # Get datetime to create unique stage and source for operation time_now = start_time s3_stage_output_name = \ s3_const.s3_stage_output_template + '{}_{}'.format(code, time_now) # Check and Create BULK_UPLOAD_FILE_FORMAT db_utils.create_file_format( format_type='output', db=source_db, schema=source_schema) # Create S3 Stage db_utils.create_s3_stage(bucket=bucket, path=path, db=source_db, schema=source_schema, s3_stage_name=s3_stage_output_name, file_format=file_format, aws_key=aws_key, aws_secret=aws_secret) # Clean illegal characters input_file_root = os.path.splitext(file_name)[0] for char in illegal_chars: # Get filename without illegal chars input_file_root = input_file_root.replace(char, '_') # Export Summary Files try: # Create a joined table of config.SNOWFLAKE_TRACK_LOG_TABLE with # config.SNOWFLAKE_SOURCE_TABLE unload_result_csv(session_id=session_id, input_file_root=input_file_root, s3_stage_name=s3_stage_output_name, s3_table_name=table_name, code=code, label_id=label_id) except Exception as e: log.info( 'Unload to CSV on \'{}\' failed.'.format(s3_stage_output_name)) raise e try: unload_reporting_csv( session_id=session_id, label_id=label_id, input_file_root=input_file_root, source_table_name=table_name, s3_stage_name=s3_stage_output_name, code=code, start_time=start_time) except Exception as e: log.info( 'Unload to CSV on \'{}\' failed.'.format(s3_stage_output_name)) raise e sig = input_file_root + '_' + code working_folder = os.path.join(path, env, session_id, sig) # Convert CSV to XLSX converted_list, new_file_list = xp.convert_all_csv_to_xlsx_s3( working_folder, ignore_folders=s3_const.ignore_folders) put_session_id(session_id=session_id, env=env, path=working_folder) # Remove all output csv s3_utils.delete_file_list(converted_list, bucket) # Drop S3 Stage orch_db.drop_s3_stage( db=source_db, schema=source_schema, s3_stage_name=s3_stage_output_name) return new_file_list # TODO - COULD NOT PROCESS - SERVER ERROR does not show up in output. # This is due to NOTICES being returned from the server, instead of ERRORs. # These are not being processed correctly. def export_file_results(session_id, table_name, file_name, code, start_time): # Create output Stage bucket = s3_bucket_name path = s3_const.output_key.format(config.S3_FOLDER) source_schema = snow_db_config['schema'] source_db = snow_db_config['database'] file_format = config.SNOWFLAKE_FILE_FORMAT_OUTPUT aws_key = AWS_ACCESS_KEY_ID aws_secret = AWS_SECRET_ACCESS_KEY env = config.ENVIRONMENT log.info('Exporting results and reports for {}....'.format(file_name)) # Get datetime to create unique stage and source for operation time_now = start_time s3_stage_output_name = \ s3_const.s3_stage_output_template + '{}_{}'.format(code, time_now) # Check and Create BULK_UPLOAD_FILE_FORMAT db_utils.create_file_format( format_type='output', db=source_db, schema=source_schema) # Create S3 Stage db_utils.create_s3_stage(bucket=bucket, path=path, db=source_db, schema=source_schema, s3_stage_name=s3_stage_output_name, file_format=file_format, aws_key=aws_key, aws_secret=aws_secret) # Clean illegal characters input_file_root = os.path.splitext(file_name)[0] for char in illegal_chars: # Get filename without illegal chars input_file_root = input_file_root.replace(char, '_') # Export Summary Files try: # Create a joined table of config.SNOWFLAKE_TRACK_LOG_TABLE with # config.SNOWFLAKE_SOURCE_TABLE unload_result_csv(session_id=session_id, input_file_root=input_file_root, s3_stage_name=s3_stage_output_name, s3_table_name=table_name, code=code) except Exception as e: log.info( 'Unload to CSV on \'{}\' failed.'.format(s3_stage_output_name)) raise e try: audio_asset_table, cover_asset_table \ = unload_reporting_csv( session_id=session_id, input_file_root=input_file_root, source_table_name=table_name, s3_stage_name=s3_stage_output_name, code=code, start_time=start_time) except Exception as e: log.info( 'Unload to CSV on \'{}\' failed.'.format(s3_stage_output_name)) raise e sig = input_file_root + '_' + code working_folder = os.path.join(path, env, session_id, sig) # Convert CSV to XLSX converted_list, new_file_list = xp.convert_all_csv_to_xlsx_s3( working_folder, top_level=True, ignore_folders=s3_const.ignore_folders) # Output asset ingestion table names to file for next asset ingestion step ingestion_table_filename = '{}_asset_tables_{}_{}.txt'.format( input_file_root, code, env) table_string = audio_asset_table + '\n' + cover_asset_table file_path = os.path.join(working_folder, ingestion_table_filename) s3_utils.put_string_to_s3(file_path, table_string, s3_bucket_name) new_file_list.append(file_path) # Remove all output csv s3_utils.delete_file_list(converted_list, bucket) # Drop S3 Stage orch_db.drop_s3_stage( db=source_db, schema=source_schema, s3_stage_name=s3_stage_output_name) return new_file_list @sentry.sentry_wrap def put_session_id(session_id, env, path=None): """Put the session Id in a file in the folder.""" output_key = s3_const.output_key.format(config.S3_FOLDER) if not path: path = os.path.join(output_key, env, session_id) session_id_filename = os.path.join(path, 'session_id.txt') s3_utils.put_string_to_s3(session_id_filename, session_id, s3_bucket_name) def poll_fivetran(session_id, expected_track_count, expected_release_count): """Poll Snowflake to ensure Fivetran sync has completed. Args: session_id (str): The session_id of the current run expected_track_count (int): The expected number of track rows expected_release_count (int): The expected number of release rows Returns (bool): True if the number of rows matches. """ # Check Release table row count release_count = db_utils.select_count_of_snowflake_rows_by_session( db=config.SNOWFLAKE_SYNC_TARGET_DB, schema=config.SNOWFLAKE_SYNC_TARGET_SCHEMA, table_name=db_consts.RELEASE_LOG_TABLE, session_id=session_id) # Check Track table row count track_count = db_utils.select_count_of_snowflake_rows_by_session( db=config.SNOWFLAKE_SYNC_TARGET_DB, schema=config.SNOWFLAKE_SYNC_TARGET_SCHEMA, table_name=db_consts.TRACK_LOG_TABLE, session_id=session_id) log.info('Number of release rows synced via Fivetran: {}'.format( release_count )) log.info('Number of release rows expected: {}'.format( expected_release_count )) log.info('Number of track rows synced via Fivetran: {}'.format( track_count )) log.info('Number of track rows expected: {}'.format( expected_track_count )) count_success = track_count >= expected_track_count \ and release_count >= expected_release_count return count_success @sentry.sentry_wrap def output_files_with_polling( session_id, table_list, total_rows, total_releases, start_time): """Write output files iff Snowflake polling shows FiveTran sync complete. Args: session_id (str): The session_id of the current run table_list (dict): Data regarding tables, files, and labels total_rows (int): Total number of rows processed total_releases (int): Total number of releases processed start_time: Returns: """ fivetran_sync_complete = False try_count = 0 # Async check for while not fivetran_sync_complete and \ try_count < config.FIVETRAN_SYNC_TRIES: # Report sleep msg = 'Sleeping for {} minutes while Fivetran sync completes.' log.info(msg.format(int(config.FIVETRAN_SYNC_WAIT / 60))) # Wait for sync sleep(config.FIVETRAN_SYNC_WAIT) # Check for sync fivetran_sync_complete = poll_fivetran( session_id, total_rows, total_releases) # Burn a try try_count += 1 # Report if failure if not fivetran_sync_complete: msg = 'Fivetran sync was not complete. {} tries left.' tries_left = config.FIVETRAN_SYNC_TRIES - try_count + 1 log.info(msg.format(tries_left)) # Touch the local environment file to prevent timeout with open('tmp/{}'.format(keepalive_file), 'w') as f: f.write('') # Check ultimate success if fivetran_sync_complete or config.FORCE_WRITE_OUTFILES: # Output reports for file_name, elements in table_list.items(): for label_id, tracking in elements['label_tracking'].items(): new_file_list = export_label_results( session_id, elements['table_name'], label_id, file_name, elements['code'], start_time) tracking['output_file_names'] = new_file_list tracking = update_bulk_upload_row(**tracking).message # Write out last whole file. export_file_results(session_id, elements['table_name'], file_name, elements['code'], start_time) else: msg = 'Fivetran sync was not completed after {} tries over ' \ '{} minutes.' log.error(msg.format( config.FIVETRAN_SYNC_TRIES + 1, config.FIVETRAN_SYNC_WAIT * config.FIVETRAN_SYNC_TRIES + 1))