import numpy as np import pandas as pd import logging import boto3 from botocore.exceptions import ClientError import time from io import StringIO import os from config import config import sentry_sdk from constants import knr_consts from utils.formatters import format_dataframe, format_date_for_error_file from utils.validation import handle_validation_errors, validate_mandatory_columns, validate_non_mandatory_columns, validate_columns_with_set_length, validate_contributor_columns, validate_isrc_column log = config.get_logger() sentry_sdk.init(config.SENTRY_DSN) def main(env): chunk_size = os.environ.get("CHUNK_SIZE") file_name = os.environ.get("FILE_NAME") log.info(f"Converting {file_name} to Ownership Ingest csv") start_time = time.perf_counter() # Process file raw_df = xls_to_dataframe(file_name) ingest_df = df_to_ingest_template_format(raw_df, env, file_name) dataframe_to_csv_chunks(ingest_df, chunk_size, file_name) time_taken = calculate_time_taken(start_time, time.perf_counter()) log.info(f'{len(ingest_df.index)} valid rows of {len(raw_df.index)} original rows processed in {time_taken}.') def xls_to_dataframe(file_name): log.info(f'Converting {file_name} to dataframe....') df = pd.read_excel(f'{file_name}', header=2, dtype=object) # Reset row indices to start from 0 after drop df.reset_index(drop=True, inplace=True) return df def df_to_ingest_template_format(df: pd.DataFrame, env, file_name): # Columns in ingest template and corresponding KNR columns. Excludes any columns/counterparts which don't yet exist in knr audio files (Content Type, Product SubType, Product Version, Associated Audio ISRC). These are inserted further down. try: log.info(f'Converting dataframe to ingest format...') # Create new df with required knr columns ingest_template_df = df[knr_consts.CORRESPONDING_KNR_COLUMNS] # Rewrite ownership ingest headers ingest_template_df.columns = knr_consts.OWNERSHIP_INGEST_COLUMNS # Dataframe that will hold errors error_df = pd.DataFrame(columns=knr_consts.OWNERSHIP_INGEST_COLUMNS) errors = [] rows_to_drop = [] # Insert missing columns so output exactly matches template ingest_template_df.insert(6, 'Content Type', 'Digital') ingest_template_df.insert(9, 'Product SubType', '') ingest_template_df.insert(10, 'Product Version', '') ingest_template_df.insert(21, 'Associated Audio ISRC', '') # VALIDATE MANDATORY COLUMNS ingest_template_df, error_df, validation_errors, error_row_indexes = validate_mandatory_columns(ingest_template_df, error_df) errors = errors + validation_errors rows_to_drop = rows_to_drop + error_row_indexes ingest_template_df, error_df, invalid_length_errs, cell_length_error_indexes = validate_columns_with_set_length(ingest_template_df, error_df, knr_consts.COLUMN_LENGTHS) errors = errors + invalid_length_errs rows_to_drop = rows_to_drop + cell_length_error_indexes for col_name in knr_consts.MANDATORY_COLUMNS_NO_VALIDATION: missing_vals = ingest_template_df.loc[ingest_template_df.loc[:, col_name].isna()] error_df, errs = handle_validation_errors(missing_vals, col_name, error_df) for x in missing_vals.index: rows_to_drop.append(x) errors = errors + errs # VALIDATE NON-MANDATORY COLUMNS ingest_template_df, error_df, non_mandatory_errors, non_man_err_indexes = validate_non_mandatory_columns(ingest_template_df, error_df) errors = errors + non_mandatory_errors rows_to_drop = rows_to_drop + non_man_err_indexes # VALIDATE CONTRIBUTOR COLUMNS (do not drop from ingest_template_df as values replaced with empty strings) ingest_template_df, error_df, contributor_errors = validate_contributor_columns(ingest_template_df, error_df) errors = errors + contributor_errors # VALIDATE ISRC COLUMN ingest_template_df, error_df, isrc_errors, isrc_row_indexes = validate_isrc_column(ingest_template_df, error_df) errors = errors + isrc_errors rows_to_drop = rows_to_drop + isrc_row_indexes # DROP ROWS TO BE EXCLUDED FROM MAIN FILE ingest_template_df.drop(sorted(list(set(rows_to_drop))), inplace=True) # WRITE ERROR FILE TO S3 write_error_file(error_df, env, file_name, errors) # FORMAT VALID VALUES ingest_template_df = format_dataframe(ingest_template_df) # Reset all indexes (after dropping invalid) ingest_template_df.reset_index(drop=True, inplace=True) return ingest_template_df except Exception as e: log.error(f'Error converting dataframe: {e}') pass def write_error_file(file, env, file_name, errors): # Format dates to avoid printing with timestamp file = file.assign(**{ 'Sales Start Date': format_date_for_error_file(file.loc[:,'Sales Start Date']), 'Collection Ownership Start Date': format_date_for_error_file(file.loc[:,'Collection Ownership Start Date']), 'Collection Ownership End Date': format_date_for_error_file(file.loc[:,'Collection Ownership End Date']), }) # Rename columns to knr terminology updated_knr_columns = knr_consts.CORRESPONDING_KNR_COLUMNS updated_knr_columns.pop() file.columns = updated_knr_columns + ["External ID", "Content Type", "Product SubType", "Product Version", "Associated Audio ISRC"] file.insert(0, "Error", errors) file.sort_index(inplace=True) file.index.name = "Original Row Number" xls_indices = list(map(lambda i: i + 4, file.index.tolist())) # Add 4 to each index, so that it maps to original excel (to refactor) file["Original Row Number"] = xls_indices file.set_index("Original Row Number", inplace=True) if(len(file.index >= 1)): s3_bucket = f'{env.lower()}-neighbouring-rights' csv_buffer = StringIO() file.to_csv(csv_buffer) s3_client = boto3.client('s3') s3_client.put_object(Body=csv_buffer.getvalue(), Bucket=s3_bucket, Key=f'{knr_consts.S3_KEY_PREFIX}errors/{file_name.split(".xls")[0]}-error-rows.csv') def dataframe_to_csv_chunks(df: pd.DataFrame, chunk_size, file_name): s3_bucket = f'{env.lower()}-neighbouring-rights' try: file_index = 1 # Dataframe to store current chunk of rows to write (is reset after each chunk written) chunk_to_write = pd.DataFrame() for upc_chunk in df.groupby(["UPC"]): # If the chunk of UPCs is larger than the chunk_size set, write a file for all that UPCs rows if (len(upc_chunk[1]) > int(chunk_size)): write_output_file_to_s3(upc_chunk[1], file_name, file_index) file_index += 1 else: # If number of rows in chunks_to_write + rows in upc_chunk is greater than chunk_size, write existing chunk_to_write, then reset chunk_to_write if (len(chunk_to_write) + len(upc_chunk[1]) > int(chunk_size)): write_output_file_to_s3(chunk_to_write, file_name, file_index) chunk_to_write = chunk_to_write[0:0] file_index += 1 # Add current upc_chunk to chunk_to_write chunk_to_write = pd.concat([chunk_to_write, upc_chunk[1]]) # Write final chunk_to_write (as will likely be some data left after last iteration) if len(chunk_to_write) > 0: write_output_file_to_s3(chunk_to_write, file_name, file_index) except ClientError as e: logging.error(e) pass def write_output_file_to_s3(chunk, file_name, file_index): s3_bucket = f'{env.lower()}-neighbouring-rights' log.info(f'Writing {file_name} to S3....') csv_buffer = StringIO() chunk.to_csv(csv_buffer, index=False) s3_client = boto3.client('s3') s3_client.put_object( Body=csv_buffer.getvalue(), Bucket=s3_bucket, Key=f'{knr_consts.S3_KEY_PREFIX}{file_name.split(".xls")[0]}-{file_index}.csv', ContentType='text/csv', Metadata={ 'profileid': '135936', 'filename': file_name, 'profileuuid': '8bd388f6-0072-4198-9c4f-f142e2685793', 'identityid': '20f15d6a-8ea6-4826-ae4e-ec14f88ce818' } ) def calculate_time_taken(start_time, end_time): time_taken = end_time - start_time if(time_taken < 60) : return (f'{round(time_taken, 3)}s') if(time_taken >= 60): minutes = int(time_taken / 60) return (f'{int(time_taken / 60)}m {int(time_taken - (minutes * 60))}s') if __name__ == "__main__": env = os.environ.get("ENV", 'qa') try: main(env) except Exception as e: log.error(e)