import argparse import os import uuid import logging import sys import time import boto3 from sqlalchemy.exc import IntegrityError from sqlalchemy.exc import SQLAlchemyError import config from connectors import mysql, sentry from constants import asset_types from models import asset_upload from models import release BUCKET = 'pf-thumb-migrate' RAW_BUCKET = '{}-orcd-raw-assets'.format(config.ENVIRONMENT) TOKEN = config.TOKEN USER_ID = 'alw:43807' PREFIX = 'round2' IMAGE_TYPES = [ asset_types.TYPE_FILE_JPG, asset_types.TYPE_FILE_JPEG, asset_types.TYPE_FILE_TIF, asset_types.TYPE_FILE_TIFF ] root_logger = logging.getLogger() root_logger.setLevel(logging.INFO) handler_stdout = logging.StreamHandler(sys.stdout) handler_stdout.setLevel(logging.INFO) root_logger.addHandler(handler_stdout) def _generate_filename(): unique_name = uuid.uuid4() underscored_name = str(unique_name).replace('-', '_') return underscored_name def main(): # parse commandline args parser = argparse.ArgumentParser(prog='thumbnail backfill') parser.add_argument('-start', '--starting_token', required=False, help='starting point') args = parser.parse_args() # Create boto connections client = boto3.client('s3', region_name='us-east-1') s3 = boto3.resource('s3', region_name='us-east-1') # Set paginator offset if passed paginator_config = { 'PageSize': 100 } if args.starting_token: paginator_config['StartingToken'] = args.starting_token # Create paginator paginator = client.get_paginator('list_objects') page_iterator = paginator.paginate( Bucket=BUCKET, Prefix=PREFIX, PaginationConfig=paginator_config ) product_not_found_file = open("not_found.txt", "a+") product_has_artwork_file = open("has_artwork.txt", "a+") file_already_processed_file = open("processed.txt", "a+") # iterate through all objects. for page in page_iterator: tifs = page['Contents'] # One page of image files. with mysql.ar_db_session() as ar_session, mysql.au_db_session() as rds_session: # iterate over image files. for tif in tifs: key = tif['Key'] # Edge case: folder if key[-1] == '/': root_logger.log( msg='Skipping: {}. It is a folder).'.format(key), level=logging.INFO ) continue upc, ext = os.path.splitext(key) # Edge case: non-image if ext.strip('.').upper() not in IMAGE_TYPES: root_logger.log( msg='Skipping: {}. Not an image.'.format(key), level=logging.INFO ) continue # Get upc upc = int(os.path.split(upc)[-1]) # Get source object metadata try: response = client.head_object(Bucket=BUCKET, Key=key) except Exception as e: root_logger.log( msg='Skipping: {}. Could not find source file.'.format( key), level=logging.INFO ) continue metadata = response['Metadata'] # Check if file is already processed. env_processed_key = '{}_processed'.format(config.ENVIRONMENT) is_new_source_file = env_processed_key not in metadata or \ metadata[env_processed_key] == '0' # If new file if is_new_source_file: try: # get release from art_relation product = release.get_release_instance(upc, ar_session) if product: release_id = product.release_id root_logger.log( msg="processing product UPC: {} - " "Release Id: {}".format( upc, release_id), level=logging.INFO) if config.OVERWRITE_EXISTING: has_images = False else: # check RDS for assets v2 record assets = \ asset_upload.\ get_asset_uploads_by_product_id( release_id, rds_session) # check if any of the v2 assets are images has_images = any( [ asset.asset_type in IMAGE_TYPES for asset in assets ] ) # this product has an image asset, skip if has_images: root_logger.log( msg="skip product, already has artwork", level=logging.INFO) product_has_artwork_file.write( "{}\n".format(upc)) # this product has no image assets. else: # generate the unique filename unique_filename = _generate_filename() root_logger.log( msg="create image file {}".format( unique_filename), level=logging.INFO) # Create record on RDS asset_upload_record = \ asset_upload.create_asset_upload( user_id=USER_ID, filename=unique_filename, token=TOKEN, session=rds_session ) # if asset was successfully created in RDS if asset_upload_record: root_logger.log( msg="Committing Row to db.", level=logging.INFO) try: rds_session.commit() except Exception as e: root_logger.log( msg=str(e), level=logging.INFO) # copy file from source to destination copy_source = { 'Bucket': BUCKET, 'Key': key } # Add needed metadata with logging s3_metadata = { 'asset_type': 'TIF', 'product_id': str(release_id), 'original_filename': key, 'upc': str(upc), 'track_unique_id': '0', 'is_correction': '0' } msg = 'copy_source:\n{}'.format( copy_source) root_logger.log( msg=msg, level=logging.INFO) msg = 's3_metadata:\n{}'.format( s3_metadata) root_logger.log( msg=msg, level=logging.INFO) msg = 'RAW_BUCKET: {}'.format(RAW_BUCKET) root_logger.log(msg=msg, level=logging.INFO) unique_tif_filename = '{}.tif'.format( unique_filename) root_logger.log( msg='Unique filename: {}'.format( unique_tif_filename), level=logging.INFO) try: # Copy in place to attach metadata s3.meta.client.copy( copy_source, BUCKET, key, ExtraArgs={ 'Metadata': s3_metadata, 'ContentType': 'image/tiff', 'MetadataDirective': 'REPLACE' }, SourceClient=client ) # Copy to raw assets bucket s3.meta.client.copy( copy_source, RAW_BUCKET, unique_tif_filename, ExtraArgs={ 'ContentType': 'image/tiff', 'Metadata': s3_metadata }, SourceClient=client ) except Exception as e: root_logger.log( msg=str(e), level=logging.INFO) # Check that the file was copied try: head_test = client.head_object( Bucket=RAW_BUCKET, Key=unique_tif_filename ) except Exception as e: root_logger.log(msg=str(e), level=logging.INFO) msg = '{} not found on {}'.format( unique_tif_filename, RAW_BUCKET ) root_logger.log(msg=msg, level=logging.INFO) else: msg = '{} found on {}'.format( unique_tif_filename, RAW_BUCKET ) root_logger.log(msg=msg, level=logging.INFO) msg = 'head_object response:\n' \ '{}'.format(head_test) root_logger.log(msg=msg, level=logging.INFO) # update source file metadata in-place s3.meta.client.copy( copy_source, BUCKET, key, ExtraArgs={ 'Metadata': { env_processed_key: '1', **s3_metadata }, 'MetadataDirective': 'REPLACE' }, SourceClient=client ) root_logger.log( msg="Source asset marked as processed", level=logging.INFO) else: # if not found for some reason, save it to file product_not_found_file.write("{}\n".format(upc)) except IntegrityError as e: root_logger.log(msg=e, level=logging.INFO) if sentry.sentry_client: sentry.sentry_client.captureException() except SQLAlchemyError as e: root_logger.log(msg=e, level=logging.INFO) if sentry.sentry_client: sentry.sentry_client.captureException() else: root_logger.log( msg="File {} already processed".format(key), level=logging.INFO) file_already_processed_file.write(key) time.sleep(2) product_not_found_file.close() product_has_artwork_file.close() if __name__ == "__main__": main()