import argparse import boto3 import botocore import decimal import json import requests import time import uuid from utils.timer import Timer import field_const import value_const from connectors import dynamodb def get_current_ownership(isrc): item = active_table.get_item(Key={field_const.ISRC: isrc}) ownership_from_dynamo = item.get('Item', {}) return ownership_from_dynamo def get_claimed_territories(ownership_info): return [ territory for territory in ownership_info.get(field_const.TERRITORIES)] def create_initial_record(active_table, isrc): try: active_table.put_item( Item={ field_const.ISRC: isrc, field_const.TERRITORIES: {}, field_const.LOCKED_TERRITORIES: {}, field_const.TIMESTAMP: decimal.Decimal( str(time.time() * 1000))}, ConditionExpression='attribute_not_exists({})'.format( field_const.ISRC)) except botocore.exceptions.ClientError as e: if e.response['Error']['Code'] != 'ConditionalCheckFailedException': raise e def update_aggregated_record(active_table, isrc, territories, value, field, key): expressions = [] attribute_names = {} attribute_values = {} for index, territory in enumerate(territories): expressions.append( '{field}.#t{index}=:tr{index}'.format(field=field, index=index)) attribute_names.update( {'#t{index}'.format(index=index): territory}) attribute_values.update( {':tr{index}'.format(index=index): {key: value}}) if dynamodb.batch_update_item_fields( active_table, {field_const.ISRC: isrc}, 'SET', expressions, attribute_names, attribute_values): expressions = [] attribute_names = {} attribute_values = {} dynamodb.batch_update_item_fields( active_table, {field_const.ISRC: isrc}, 'SET', expressions, attribute_names, attribute_values, force=True) def create_audit_record( audit_table, opcode, isrc, territories, correlation_id, user, tuid): if opcode == value_const.ADD_OPCODE: territories = {a2_code: tuid for a2_code in territories} record = { field_const.ISRC: '{isrc}'.format(isrc=isrc), field_const.TIMESTAMP: decimal.Decimal(str(time.time() * 1000)), field_const.USER: value_const.BACKFILL_USER, field_const.OPCODE: opcode, field_const.TERRITORIES: territories, field_const.CORRELATION_ID: str(correlation_id) } else: record = { field_const.ISRC: '{isrc}'.format(isrc=isrc), field_const.TIMESTAMP: decimal.Decimal(str(time.time() * 1000)), field_const.USER: value_const.BACKFILL_USER, field_const.OPCODE: opcode, field_const.TERRITORIES: territories, field_const.CORRELATION_ID: str(correlation_id), field_const.LOCK_REASON: 'Registry Backfill MEAT' } audit_table.put_item(Item=record) return record def create_active_record(active_table, isrc, value, territories, field, key): create_initial_record(active_table, isrc) update_aggregated_record(active_table, isrc, territories, value, field, key) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('-q', '--queue', required=True) parser.add_argument('-d', '--dynamodb_table', required=True) parser.add_argument('-a', '--dynamodb_table_audit', required=True) parser.add_argument('-c', '--consumer_duration', default=90, type=int) parser.add_argument('-n', '--messages', default=10, type=int) parser.add_argument('-w', '--wait_time', default=20, type=int) parser.add_argument('-v', '--vis_timeout', default=20, type=int) parser.add_argument('-e', '--env', default='dev', type=str) args = parser.parse_args() sqs_resource = boto3.resource('sqs') dynamodb_resource = boto3.resource('dynamodb') timer = Timer(args.consumer_duration) num_messages = args.messages wait_time_seconds = args.wait_time visibility_timeout = args.vis_timeout mr_backfill_queue = sqs_resource.get_queue_by_name( QueueName='{}-{}'.format(args.env, args.queue)) active_table = dynamodb_resource.Table( '{}-{}'.format(args.env, args.dynamodb_table)) audit_table = dynamodb.dynamodb_resource.Table( '{}-{}'.format(args.env, args.dynamodb_table_audit)) reason = {field_const.REASON: 'Registry Backfill MEAT'} headers = {'Correlation-Id': str(uuid.uuid1())} ows_territories_url = 'http://ows-territories.theorchard.io/territory/{}' territories_standard = 'ISO_3166_1_2016' r = requests.get( ows_territories_url.format(territories_standard), headers=headers) territories_response = r.json().get('items') territories_to_lock = set() for response in territories_response: territories_to_lock.add(response.get('territory_code_a2')) while not timer.is_expired(): messages = mr_backfill_queue.receive_messages( MaxNumberOfMessages=num_messages, WaitTimeSeconds=wait_time_seconds, VisibilityTimeout=visibility_timeout) entries = [] for message in messages: ownership = json.loads(message.body) isrc = ownership[field_const.ISRC] is_locked = int(ownership.get(field_const.IS_LOCKED)) derived_tuid = ownership.get(field_const.DERIVED_TUID) derived_ownership = ownership.get( field_const.DERIVED_OWNERSHIP, []) if is_locked: current_ownership = get_current_ownership(isrc) # lock if only the registry has this isrc and does not already # have locked territories if current_ownership: territories_to_lock_copy = territories_to_lock.copy() claimed_territories = get_claimed_territories( current_ownership) unclaimed_locked_territories = territories_to_lock_copy.difference( claimed_territories) if unclaimed_locked_territories: update_aggregated_record( active_table, isrc, unclaimed_locked_territories, value_const.BACKFILL_REASON, field_const.LOCKED_TERRITORIES, field_const.REASON) create_audit_record( audit_table, value_const.LOCK_OPCODE, isrc, list(unclaimed_locked_territories), uuid.uuid1(), value_const.BACKFILL_USER, None) elif derived_ownership: tuid = int(derived_tuid) create_active_record( active_table, isrc, tuid, derived_ownership.split(','), field_const.TERRITORIES, field_const.TUID) create_audit_record( audit_table, value_const.ADD_OPCODE, isrc, derived_ownership.split(','), uuid.uuid1(), value_const.BACKFILL_USER, tuid) else: # if no ownership found, then make a blank record create_initial_record(active_table, isrc) entries.append({ 'Id': str(uuid.uuid4()), 'ReceiptHandle': message.receipt_handle}) if entries: mr_backfill_queue.delete_messages(Entries=entries)