"""Automate pulling assets by ISRC.""" import argparse import json import os import sys import time import boto3 import snowflake.connector from sqlalchemy import create_engine from sqlalchemy.sql import text # function that copies from S3 to AVL LAMBDA_FUNCTION_NAME = 'lambda-support-copy-assets-prod' # bucket files are copied into S3_BUCKET = 'prod-product-support-assets-from-avl' # could be dynamic ASSET_TYPE = 1 ASSET_LOCATION = 1 # initialized at start before calling main() SNOWFLAKE_CONN = None DD_CONN = None LABMDA_CLIENT = None S3_CLIENT = None def parse_args(): parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument( 'isrcs', type=str, nargs='+', help='ISRCs to identify tracks with') parser.add_argument( '-s', action='store_true', help='skip lambda invoke step, go straight to download') parser.add_argument( '--dir', type=str, required=False, default=os.getcwd(), help='local directory to copy files into') return parser.parse_args() def main(args): isrcs = sys.argv[1:] # get metadata about files we want files = get_file_metadata(isrcs) # start lambda functions to move assets to S3 if not args.s: invoke_functions(files) # get files from s3 download_files(files, args.dir) def get_file_metadata(isrcs): # get tracks and filenames by ISRCs tracks = tracks_by_isrcs(isrcs) print(tracks, '\n') # group filenames from tracks filenames = { x['FILENAME']: { 'upc': x['UPC'], 'isrc': x['ISRC'], 'duration': x['DURATION'] } for x in tracks } print(filenames, '\n') return filenames def tracks_by_isrcs(isrcs): return SNOWFLAKE_CONN.execute( """ SELECT x.ID, r.RELEASE_NAME, x.TRACK_NAME, x.CD, x.TRACK_ID, x.UPC, ald.filename, x.ISRC, (x.LENGTH_MINUTE * 60) + x.LENGTH_SECONDS AS duration FROM ORCHARD_APP_REPORTING_V2.ART_RELATIONS_PROD_ART_RELATIONS.TRACK AS x JOIN ORCHARD_APP_REPORTING_V2.ART_RELATIONS_PROD_ART_RELATIONS.RELEASES AS r ON r.upc = x.upc JOIN ORCHARD_APP_REPORTING_V2.REPORTSDD_DIRECT_DELIVERY.ASSET_LOCATION_DETAIL AS ald ON ald.filename = concat(x.upc,'_',x.cd,'_',x.track_id,'.wav') JOIN ORCHARD_APP_REPORTING_V2.REPORTSDD_DIRECT_DELIVERY.ASSET_LOCATION al ON ald.asset_location_id = al.asset_location_id JOIN ORCHARD_APP_REPORTING_V2.REPORTSDD_DIRECT_DELIVERY.ASSET a ON a.asset_id = al.asset_id and a.asset_type_id = %s JOIN ORCHARD_APP_REPORTING_V2.REPORTSDD_DIRECT_DELIVERY.STORAGE_DRIVE sd ON sd.storage_drive_id = al.storage_drive_id JOIN ORCHARD_APP_REPORTING_V2.REPORTSDD_DIRECT_DELIVERY.STORAGE s ON s.storage_id = sd.storage_id and s.physical_location_id = %s WHERE x.ISRC in (%s) ORDER BY x.UPC, x.TRACK_ID ASC """, # noqa:E501 (ASSET_TYPE, ASSET_LOCATION, isrcs,) ).fetchall() def files_by_upcs(upcs): # https://github.com/theorchard/lambda-support-copy-assets/blob/master/lambda/support_copy_assets/sql/queries.py query = text(""" SELECT CASE WHEN ta.file_type = 'file' THEN CONCAT(REPLACE(sdaf.`initial_folder`, '/', ''), '/', SUBSTRING(a.`upc`, 1, 6), '/', SUBSTRING(a.`upc`, 7, 3),'/',ald.`filename`) ELSE CONCAT(REPLACE(sdaf.`initial_folder`, '/', ''), '/', SUBSTRING(a.`upc`, 1, 6), '/', SUBSTRING(a.`upc`, 7, 3), '/',a.`upc`,'/',ald.`filename`) END AS file_path, a.`upc`, ta.file_type, INET_NTOA(s.`local_ip`) AS ip FROM `asset` a INNER JOIN `asset_location` al ON a.asset_id = al.asset_id INNER JOIN `storage_drive` sd ON sd.storage_drive_id = al.storage_drive_id INNER JOIN `storage` s ON s.storage_id = sd.storage_id LEFT JOIN `asset_location_detail` ald ON ald.asset_location_id = al.asset_location_id INNER JOIN `asset_type` ta ON a.asset_type_id = ta.asset_type_id INNER JOIN `storage_drive_asset_folder` sdaf ON sd.storage_drive_id = sdaf.storage_drive_id AND sdaf.asset_type_id = a.asset_type_id WHERE a.upc IN :upcs AND s.physical_location_id IN :asset_locations AND a.asset_type_id IN :asset_types ORDER BY a.upc """) # noqa:E501 return DD_CONN.execute( query, upcs=tuple(upcs), asset_locations=tuple([ASSET_LOCATION]), asset_types=tuple([ASSET_TYPE]) ).fetchall() def invoke_functions(files): # do same query as lambda to calculate offsets upcs = set([ v['upc'] for _, v in files.items() ]) print(upcs, '\n') file_rows = files_by_upcs(upcs) print(file_rows, '\n') # calculate offsets for running lambda indexes = [ idx for idx, file_row in enumerate(file_rows) if file_row['file_path'].split('/')[-1] in files.keys() ] print(indexes, '\n') # invoke lambda using offsets to fetch exact track(s) for idx in indexes: payload = json.dumps({ 'upcs': str(upcs), 'asset_type': ASSET_TYPE, 'offset': idx, 'limit': 1 }) print(payload) response = LAMBDA_CLIENT.invoke( FunctionName=LAMBDA_FUNCTION_NAME, Payload=payload, InvocationType='Event' ) print(response['ResponseMetadata']['HTTPStatusCode']) print('\n') def download_files(files, dirname, timeout=960): start_time = time.time() if not dirname.endswith(os.path.sep): dirname += os.path.sep # try to download files until timeout or all done while True: finished_files = 0 # check each file for filename, details in files.items(): local_filename = f"{dirname}{details['isrc']}_{details['duration']}_{filename}" # noqa:E501 # skip if already downloaded if not os.path.isfile(local_filename): response = S3_CLIENT.list_objects(Bucket=S3_BUCKET, Prefix=str(details['upc'])) # noqa:E501 # match bucket contents with filename if 'Contents' in response: for s3_file in response['Contents']: s3_filename = s3_file['Key'].split('/')[-1] if s3_filename == filename: print(s3_file) S3_CLIENT.download_file(S3_BUCKET, s3_file['Key'], local_filename) # noqa:E501 print(local_filename) finished_files += 1 else: print(local_filename) finished_files += 1 # timeout or sleep and try again if finished_files >= len(files.keys()): print('done') return elif time.time() - start_time > timeout: print('timeout downloading files') return else: print('sleeping...') time.sleep(60) if __name__ == '__main__': args = parse_args() if not os.path.isdir(args.dir): exit(f'{args.dir} is not a valid directory') SNOWFLAKE_CONN = snowflake.connector.connect( user=os.getenv('SNOWFLAKE_USER'), password=os.getenv('SNOWFLAKE_PASSWORD'), account=os.getenv('SNOWFLAKE_ACCOUNT'), warehouse=os.getenv('SNOWFLAKE_WAREHOUSE') ).cursor(snowflake.connector.DictCursor) SNOWFLAKE_CONN.execute('SELECT 1') print('Snowflake initialized...') user = os.getenv('DD_USER') password = os.getenv('DD_PASSWORD') hostname = os.getenv('DD_HOSTNAME') DD_CONN = create_engine(f'mysql+pymysql://{user}:{password}@{hostname}/direct_delivery') # noqa:E501 DD_CONN.execute('SELECT 1') print('Direct Delivery initialized...') if not args.s: LAMBDA_CLIENT = boto3.client('lambda') LAMBDA_CLIENT.get_function(FunctionName=LAMBDA_FUNCTION_NAME) print('Lambda Client initialized...') S3_CLIENT = boto3.client('s3') S3_CLIENT.head_bucket(Bucket=S3_BUCKET) print('S3 Client initialized...') print('\n') main(args)