def copy_parts(event, context): topic_arn = event["Records"][0]["Sns"]["TopicArn"] msg = json.loads(event["Records"][0]["Sns"]["Message"]) blobstore_handle = dss.Config.get_blobstore_handle( platform_to_replica[msg["source_platform"]]) source_url = blobstore_handle.generate_presigned_GET_url( bucket=msg["source_bucket"], key=msg["source_key"]) futures = [] gs = dss.Config.get_native_handle(Replica.gcp) with ThreadPoolExecutor(max_workers=4) as executor: for part in msg["parts"]: logger.info(log_msg.format(part=part, **msg)) if msg["dest_platform"] == "s3": upload_url = "{host}/{bucket}/{key}?partNumber={part_num}&uploadId={mpu_id}".format( host=clients.s3.meta.endpoint_url, bucket=msg["dest_bucket"], key=msg["dest_key"], part_num=part["id"], mpu_id=msg["mpu"]) elif msg["dest_platform"] == "gs": assert len(msg["parts"]) == 1 dest_blob_name = "{}.part{}".format(msg["dest_key"], part["id"]) dest_blob = gs.get_bucket( msg["dest_bucket"]).blob(dest_blob_name) upload_url = dest_blob.create_resumable_upload_session( size=part["end"] - part["start"] + 1) futures.append( executor.submit(copy_part, upload_url, source_url, msg["dest_platform"], part)) for future in futures: future.result() if msg["dest_platform"] == "s3": mpu = resources.s3.Bucket(msg["dest_bucket"]).Object( msg["dest_key"]).MultipartUpload(msg["mpu"]) parts = list(mpu.parts.all()) elif msg["dest_platform"] == "gs": part_names = [ "{}.part{}".format(msg["dest_key"], p + 1) for p in range(msg["total_parts"]) ] parts = [ gs.get_bucket(msg["dest_bucket"]).get_blob(p) for p in part_names ] parts = [p for p in parts if p is not None] logger.info("Parts complete: {}".format(len(parts))) logger.info("Parts outstanding: {}".format(msg["total_parts"] - len(parts))) if msg["total_parts"] - len( parts) < parts_per_worker[msg["dest_platform"]] * 2: logger.info("Calling closer") send_sns_msg( ARN(topic_arn, resource=sns_topics["closer"][msg["dest_platform"]]), msg) logger.info("Called closer")
def reaper(event, context): queue_name = "dss-dlq-" + os.environ["DSS_DEPLOYMENT_STAGE"] # Get the queue queue = sqs.get_queue_by_name(QueueName=queue_name) message_count = 0 while context.get_remaining_time_in_millis() > 20000: # Long poll for messages for message in queue.receive_messages(MaxNumberOfMessages=10, AttributeNames=['All'], MessageAttributeNames=['All'], WaitTimeSeconds=10): try: logger.debug("Received a message for reprocessing: %s", str(message.body)) # re-process messages by sending them back to the SNS topic sns_message = json.loads(message.body)["Records"][0]["Sns"] topic_arn = sns_message["TopicArn"] msg = json.loads(sns_message["Message"]) logger.info( f"Received a message for reprocessing: %s SNS topic ARN: %s", str(msg), topic_arn) attrs = sns_message["MessageAttributes"] retry_count = int(attrs[DSS_REAPER_RETRY_KEY]['Value'] ) if DSS_REAPER_RETRY_KEY in attrs else 0 retry_count += 1 if retry_count < DSS_MAX_RETRY_COUNT: attrs = { DSS_REAPER_RETRY_KEY: { "DataType": "Number", "StringValue": str(retry_count) } } logger.info( "Incremented retry count: %d and resend SNS message", retry_count) send_sns_msg(topic_arn, msg, attrs) else: logger.critical( "Giving up on message: %s after %d attempts", msg, retry_count) except Exception as e: logger.error("Unable to process message: %s due to %s", str(message), str(e)) # Let the queue know that the message is processed logger.info('Deleting message from the queue') message.delete() message_count += 1 logger.info("Processed %d messages", message_count)
def step_functions_invoke(state_machine_name_template: str, execution_name: str, input, attributes=None) -> typing.Any: """ Invoke a step functions state machine. The name of the state machine to be invoked will be derived from `state_machine_name_template`, with string formatting to replace {stage} with the dss deployment stage. """ message = { SFN_TEMPLATE_KEY: state_machine_name_template, SFN_EXECUTION_KEY: execution_name, SFN_INPUT_KEY: json.dumps(input) } logger.debug('Sending message: %s', str(message)) response = send_sns_msg(sfn_sns_topic_arn, message, attributes) return response