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