"""CLI interface for faster dev.""" import argparse import datetime import json import os import sys import time import boto3 # Apply local-stack overrides before db.py opens connections if '--local-stack' in sys.argv: os.environ['LOCAL_STACK'] = 'true' os.environ['DD_DB_HOST'] = os.environ['TEST_DD_DB_HOST'] os.environ['DD_USER'] = os.environ['TEST_DD_DB_USERNAME'] os.environ['DD_PASS'] = os.environ['TEST_DD_DB_PASSWORD'] os.environ['AR_DB_HOST'] = os.environ['TEST_AR_DB_HOST'] os.environ['AR_DB_USERNAME'] = os.environ['TEST_AR_DB_USERNAME'] os.environ['AR_DB_PASSWORD'] = os.environ['TEST_AR_DB_PASSWORD'] os.environ['S3_XML_BACKUP_BUCKET'] = 'local-vector-audit' import db sqs_client = boto3.client( 'sqs', endpoint_url='http://moto:5000', region_name='us-east-1', aws_access_key_id='abc', aws_secret_access_key='123' ) _s3_backup_bucket = os.environ.get('S3_XML_BACKUP_BUCKET', 'qa-vector-audit') _s3_backup_path = os.environ.get('S3_XML_BACKUP_PATH', '/metadata').strip('/') if '--local-stack' in sys.argv: s3_backup_client = boto3.client( 's3', endpoint_url='http://moto:5000', region_name='us-east-1', aws_access_key_id='abc', aws_secret_access_key='123' ) else: s3_backup_client = boto3.client('s3') allowed_encoder_ids = os.environ['ALLOWED_ENCODER_IDS'].split(',') allowed_priorities = os.environ['ENCODING_PRIORITIES'].split(',') def main() -> None: """Entrypoint.""" args = load_args() _print('Hello! Welcome to the VECTOR worker computer aided enrichment center!') # noqa:E501 _print(args) if args.local_stack: db.seed_local_stack() try: s3_backup_client.create_bucket(Bucket=_s3_backup_bucket) _print(f'Created S3 backup bucket => {_s3_backup_bucket}') except s3_backup_client.exceptions.BucketAlreadyOwnedByYou: pass if args.command == 'encode': eqd_id = encode_job( args.upc, args.dms_id, args.encoder_id, not args.a, args.priority ) if args.d: deliver_job(eqd_id, args.encoder_id) if args.command == 'deliver': deliver_job(args.eqd_id, args.encoder_id) def load_args(): """Parse and set cli args.""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter ) parser.add_argument( '--encoder_id', type=int, default=allowed_encoder_ids[0], choices=allowed_encoder_ids, help='encoder id use in queue namespaces' ) shared = argparse.ArgumentParser(add_help=False) shared.add_argument( '--local-stack', action='store_true', default=False, help='use local DBs and infrastructure instead of QA' ) subparser = parser.add_subparsers(dest='command') subparser.required = True encode_parser = subparser.add_parser( 'encode', parents=[shared], help='setup and encode job by upc', formatter_class=argparse.ArgumentDefaultsHelpFormatter ) encode_parser.add_argument( 'upc', type=str, help='product UPC to create job with => https://www.notion.so/Encoding-Workers-4614a8e498034b4984eba275c3709c21?pvs=4#4a9f2a8ff4d04661a8dd5f37fdc67229' # noqa:E501 ) encode_parser.add_argument( '-a', action='store_true', help='package assets in addition to metadata' ) encode_parser.add_argument( '-d', action='store_true', help='deliver package after encoding' ) encode_parser.add_argument( '--priority', type=int, default=allowed_priorities[0], choices=allowed_priorities, help='dpm priority of job' ) encode_parser.add_argument( '--dms_id', type=int, default=1896, help='dms store id to create job with' ) deliver_parser = subparser.add_parser( 'deliver', parents=[shared], help='deliver job by encoding_queue_detail_id' ) deliver_parser.add_argument( 'eqd_id', type=int, help='encoding_queue_detail_id to deliver' ) return parser.parse_args() def encode_job( upc: int, dms_id: int, encoder_id: int, meta_update: bool, priority: int) -> int: """Insert db records and sqs event for encoding order.""" db.select_product(upc) # validates UPC exists in local seed before queuing encoding_queue = db.insert_encoding_queue( encoder_id, meta_update, priority ) encoding_queue_detail = db.insert_encoding_queue_detail( encoding_queue['encoding_queue_id'], upc, 'ready_to_encode', dms_id ) eqd_id = encoding_queue_detail['encoding_queue_detail_id'] _print(f'Created encoding queue detail => {eqd_id}') queue_payload = { 'encoding_queue_detail_id': eqd_id, } queue_name = f'dev-encoding{encoder_id}_e{str(encoding_queue_detail["dpm_priority"]).zfill(7)}' # noqa:E501 queue_url = sqs_client.create_queue(QueueName=queue_name)['QueueUrl'] sqs_client.purge_queue(QueueUrl=queue_url) sqs_client.send_message( QueueUrl=queue_url, MessageBody=json.dumps(queue_payload) ) _print(f'Inserted job into queue => {queue_name}') watch_job(eqd_id, ['ready_to_encode', 'encoding'], 'encoded') return eqd_id def deliver_job( eqd_id: int, encoder_id: int) -> None: """Deliver encoded job.""" db.update_encoding_queue_detail_status( eqd_id, 'queued_for_delivery' ) encoding_queue_detail = db.select_encoding_queue_detail( eqd_id ) dms_delivery_spec = db.select_delivery_spec( encoding_queue_detail['dms_master_master_id'] ) delivery_batch_id = None if dms_delivery_spec['batch_delivery'] == 'Y': delivery_batch_id = db.insert_batch(eqd_id) _print(f'Created batch => {delivery_batch_id}') # https://github.com/theorchard/lambda-vector/blob/master/lambda/throttler/src/index.py#L156-L160 queue_payload = { 'eqd_id': eqd_id, 'dms_id': encoding_queue_detail['dms_master_master_id'], 'delivery_batch_id': delivery_batch_id } queue_name = f'dev-delivery{encoder_id}_e{str(encoding_queue_detail["dpm_priority"]).zfill(7)}' # noqa:E501 queue_url = sqs_client.create_queue(QueueName=queue_name)['QueueUrl'] sqs_client.purge_queue(QueueUrl=queue_url) sqs_client.send_message( QueueUrl=queue_url, MessageBody=json.dumps(queue_payload) ) _print(f'Inserted job into queue => {queue_name}') watch_job(eqd_id, ['queued_for_delivery', 'delivering'], 'delivered') if delivery_batch_id: db.update_batch_status(delivery_batch_id, 'closed') _verify_xml_backup(eqd_id) def watch_job( eqd_id: int, progress_statuses: list[str], end_status: str) -> None: """Poll for status changes of job.""" last_job_status = None while True: time.sleep(1) encoding_queue_detail = db.select_encoding_queue_detail(eqd_id) job_status = encoding_queue_detail['status'] if job_status != last_job_status: _print(f'Job ({eqd_id}) : status = {job_status}') last_job_status = job_status if job_status in progress_statuses: continue elif job_status == end_status: break else: _print(f"ERROR => {encoding_queue_detail['error_log']}") exit(1) def _verify_xml_backup(eqd_id: int) -> None: """Check that at least one XML backup file was uploaded to S3 for this EQD.""" import hashlib prefix = f'{_s3_backup_path}/{hashlib.md5(str(eqd_id).encode()).hexdigest()}' try: response = s3_backup_client.list_objects_v2( Bucket=_s3_backup_bucket, Prefix=prefix, ) found = response.get('Contents', []) if found: for obj in found: _print(f'XML backup verified => s3://{_s3_backup_bucket}/{obj["Key"]}') else: _print(f'WARNING: no XML backup found in s3://{_s3_backup_bucket}/{prefix}*') except Exception as e: _print(f'WARNING: could not verify XML backup: {e}') def _print(message: str) -> None: print(f'{datetime.datetime.now()} {message}') if __name__ == '__main__': main()