"""find_or_create_account.""" import json import uuid import boto3 import config from constants.general import PALM_TREE_RECORDS_IMPRINT_PREFIX from constants.snowflake_queries import \ GET_COUNTRY_ID_BY_PARENT_REPERTOIRE_OWNER_NAME from ddex_ingester_common.constants.swb_deal_types import \ SME_ANALYTICS_NOT_FOR_DISTRIBUTION from ddex_ingester_common.helpers.s3_ddex import load_ddex_json from ddex_ingester_common.lambda_exceptions import ( LambdaException, ProductExistUnderDifferentVendorException, RetryableException, VendorDoNotIngestException) from ddex_ingester_common.logging import utils as logging_utils from ddex_ingester_common.models.s3.body import Body as S3Body from ddex_ingester_common.models.state_machine.body import ( Body as StateMachineBody) from ddex_ingester_common.schemas.s3_schema import S3Schema from ddex_ingester_common.schemas.state_machine_schema import \ StateMachineSchema from helpers.graphql_requests import ( create_new_subaccount, create_new_vendor, get_product_by_upc, get_subaccount_for_vendor, get_vendor_for_rep_owner_code, update_external_identifier_1, update_vendor_country_id) from helpers.snowflake_utils import execute_snowflake_query from helpers.vendor_mapping import ( get_parent_repertoire_owner_mapping, get_vendor_and_subaccount_for_rep_owner_code_rows, insert_vendor_and_subaccount_mapping) logger = logging_utils.get_logger(config.app_logger) def handler(event, context): """find_or_create_account handler.""" logger.info(f'Triggered find_or_create_account: {event}') context = StateMachineSchema().load(event) s3_context = S3Schema().load(load_ddex_json(event)) correlation_id = context.correlation_id or str(uuid.uuid4()) context.correlation_id = correlation_id s3_context.correlation_id = correlation_id logging_utils.update_logger_correlation_id(logger, correlation_id) logging_utils.update_logger_with_message_ids( logger, context.message_id, context.message_thread_id, context.execution_name ) logger.info('Getting the vendor/subaccount.') try: # Verify if the product already exists under a different vendor than in the mapping table. # If so, update the existing product without modifying the mapping. # If upc is None then the product will be inserted as a new product # with a new upc using set_product as a placeholder upc if context.product.upc: vendor_id, subaccount_id = \ get_vendor_and_subaccount_for_existing_product(s3_context) else: vendor_id, subaccount_id = None, None # Check if the parent_rep_owner_name is needed to be remapped check_remapping_parent_rep_owner(s3_context) if vendor_id is None: # check if mapping exists in inbound_major_label_mapping table vendor_id, subaccount_id = get_vendor_and_subaccount_in_mapping_table( s3_context) # in case there is no mapping in the table if not vendor_id or not subaccount_id: # save the correlation_id in s3_context to use it in the graphql requests vendor_id = get_vendor_via_graphql(s3_context) if vendor_id: subaccount_id = get_subaccount_via_graphql(s3_context, vendor_id) # if vendor_id is not found, create a new vendor elif not vendor_id: vendor_id = create_vendor(s3_context) if vendor_id and not subaccount_id: # if vendor_id is found, create a new subaccount subaccount_id = create_subaccount(s3_context, vendor_id) if vendor_id and subaccount_id: logger.info( f'Inserting vendor {vendor_id} and subaccount {subaccount_id} into mapping table.') insert_vendor_and_subaccount_mapping( logger, vendor_id, subaccount_id, s3_context) if not vendor_id or not subaccount_id: logger.info( f'Vendor_id {vendor_id} and/or subaccount_id {subaccount_id} is empty.') raise LambdaException( f'Vendor_id {vendor_id} and/or subaccount_id {subaccount_id} is empty.') except LambdaException: # this will raise ProductExistUnderDifferentVendorException, VendorDoNotIngestException raise except Exception as err: # try to catch the error like timeout or connection error # and retry the lambda function in terraform logger.error(f'Error getting vendor/subaccount: {err}') raise RetryableException(f'Error getting vendor/subaccount: {err}') context.product.vendor_id = vendor_id context.product.subaccount_id = subaccount_id # Set not_for_distribution to SME_ANALYTICS_NOT_FOR_DISTRIBUTION context.product.not_for_distribution = SME_ANALYTICS_NOT_FOR_DISTRIBUTION s3_context.product.vendor_id = context.product.vendor_id s3_context.product.subaccount_id = context.product.subaccount_id save_s3_context(context, s3_context) return StateMachineSchema().dump(context) def get_vendor_and_subaccount_in_mapping_table(s3_context: S3Body): """Return the SME vendor/subaccount.""" vendor_id, subaccount_id = None, None rows = get_vendor_and_subaccount_for_rep_owner_code_rows( logger, s3_context) # There should only be one row for label mapping if len(rows) > 1: raise LambdaException( f'Multiple label mappings found in the mapping table: {rows}') if len(rows) == 1: do_not_ingest = rows[0].get('do_not_ingest', None) if do_not_ingest and \ not s3_context.product.imprint.startswith( PALM_TREE_RECORDS_IMPRINT_PREFIX): raise VendorDoNotIngestException( f'Parent rep owner code {s3_context.product.parent_repertoire_owner_code} ' f'and rep owner code {s3_context.product.repertoire_owner_code} ' f'are not for ingestion.' ) if s3_context.product.imprint.startswith( PALM_TREE_RECORDS_IMPRINT_PREFIX): logger.info('Found "Palm Tree Records" ignore do_not_ingest flag') vendor_id = rows[0].get('vendor_id', None) subaccount_id = rows[0].get('subaccount_id', None) return vendor_id, subaccount_id def get_vendor_via_graphql(s3_context: S3Body): """Get vendor via graphql request.""" logger.info('Getting the vendor via graphql.') vendors = get_vendor_for_rep_owner_code(logger, s3_context) # There should only be one row for label mapping if len(vendors) > 1: raise LambdaException(f'Multiple label mappings found: {vendors}') if len(vendors) == 1: return vendors[0].get('vendorId') def get_subaccount_via_graphql(s3_context: S3Body, vendor_id: int): """Get subaccount via graphql request.""" logger.info('Getting the subaccount via graphql.') subaccounts = get_subaccount_for_vendor(logger, s3_context, vendor_id) # There should only be one subaccount if len(subaccounts) > 1: raise LambdaException(f'Multiple subaccounts for vendor: {vendor_id}') if len(subaccounts) == 1: return subaccounts[0].get('subaccountId') def get_vendor_and_subaccount_for_existing_product(s3_context: S3Body): """Get vendor and subaccount if product exists.""" logger.info('Getting the vendor and subaccount if product exists.') product = get_product_by_upc(logger, s3_context) if product and product['notForDistribution'] == SME_ANALYTICS_NOT_FOR_DISTRIBUTION: return product['vendorId'], product['subaccountId'] if product and product['notForDistribution'] != SME_ANALYTICS_NOT_FOR_DISTRIBUTION: raise ProductExistUnderDifferentVendorException( f'Product with UPC {s3_context.product.upc} already exists under a different vendor. ' f'Vendor ID: {product["vendorId"]}, Subaccount ID: {product["subaccountId"]}' ) return None, None def check_remapping_parent_rep_owner(s3_context: S3Body): """Check if the parent_rep_owner_name is remapped.""" logger.info('Checking if the parent_rep_owner_name is needed to be remapped.') # Get parent repertoire owner mapping in REP_OWNER_HIERARCHY table check_mapping = get_parent_repertoire_owner_mapping(logger, s3_context) if not check_mapping: return if len(check_mapping) > 1: raise LambdaException( f'Multiple parent repertoire owner mappings found: {check_mapping}') remapped_parent_repertoire_owner_code = check_mapping[0]['REP_OWNER_PARENT_CD'] # Check if the parent_repertoire_owner_code is remapped if s3_context.product.parent_repertoire_owner_code != remapped_parent_repertoire_owner_code: s3_context.product.parent_repertoire_owner_code = remapped_parent_repertoire_owner_code s3_context.product.parent_repertoire_owner_name = check_mapping[0]['REP_OWNER_PARENT_NM'] def create_vendor(s3_context: S3Body): """Create vendor.""" logger.info('Creating vendor.') country_id = get_vendor_country_id(s3_context) vendor = create_new_vendor(logger, s3_context) if not vendor: raise LambdaException('Vendor creation failed.') vendor_id, vendor_uuid = vendor['vendorId'], vendor['uuid'] # Update external_identifier_1 for the new vendor update_external_identifier_1(logger, s3_context, vendor_uuid) if country_id: update_vendor_country_id(logger, country_id, vendor_uuid) return vendor_id def create_subaccount(s3_context: S3Body, vendor_id: int): """Create subaccount.""" logger.info('Creating subaccount.') subaccount = create_new_subaccount(logger, s3_context, vendor_id) if not subaccount: raise LambdaException('Subaccount creation failed.') return subaccount['subaccountId'] def save_s3_context(context: StateMachineBody, s3_context: S3Body): """Save s3 context.""" s3_client = boto3.client('s3') json_file_path = f'{context.key}parsed_ddex.json' s3_client.put_object( Bucket=context.bucket, Key=json_file_path, Body=json.dumps(S3Schema().dump(s3_context)).encode(encoding='UTF-8') ) def get_vendor_country_id(s3_context): """Get vendor country id.""" parent_repertoire_owner_name = s3_context.product.parent_repertoire_owner_name logger.info(f'Getting country id for {parent_repertoire_owner_name}') result = execute_snowflake_query( GET_COUNTRY_ID_BY_PARENT_REPERTOIRE_OWNER_NAME, {'parent_repertoire_owner_name': parent_repertoire_owner_name}) if len(result) == 0: logger.warning( f'Can not find country id for {parent_repertoire_owner_name}') return None return result[0]['ID']