PR-7502 Refactored to grab SQS from lambda

Author: benbart

ctorm/ctorm/cnm_sender.py

@@ -19,7 +19,11 @@ def __init__(
         self.provider = provider
         self.cnm_s_generator = CtormCnmSGenerator()
 
-    def send_all(self):
+    def send_all(self) -> bool:
         for granule in self.granules:
             cnms = self.cnm_s_generator(granule, self.provider)
             log.debug("Sending %s", cnms)
+            # TODO: actually send
+
+        # TODO: determine successfulllness
+        return True

ctorm/ctorm/lambda_run.py

@@ -5,6 +5,8 @@
 import json
 import logging
 
+from aws_lambda_typing import context as context_
+
 from ctorm.load_tester import lambda_handler
 
 log = logging.getLogger()
@@ -19,14 +21,15 @@
 
 
 # mock event data
-with open("../../data/lambda_event.json", "r") as f:
+with open("../../data/lambda-event-eventbridge.json", "r") as f:
     mock_event = json.load(f)
 
 
 # Mock context object
-class MockContext:
+class MockContext(context_.Context):
     function_name = "local_test"
     memory_limit_in_mb = 128
+    aws_request_id = "LOLOCAL"
 
 
 # Run it

ctorm/ctorm/load_tester.py

@@ -1,20 +1,29 @@
 import json
 import logging
 import os
+from functools import cache
 from typing import List
 
+import boto3
 from aws_lambda_typing import context as context_
 from aws_lambda_typing import events
 from cnm_sender import CnmSender
 
 from ctorm.config import (
+    AWS_REGION,
     CtormConfig,
     CtormPreparedGranule,
 )
 
 log = logging.getLogger(__name__)
 
 
+@cache
+def get_sqs_client():
+    return boto3.client("sqs", region_name=AWS_REGION)
+
+
+@cache
 def configure_cfg():
     cfg = CtormConfig.from_file(
         cfg_file=os.getenv("CFG_FILE", "./ctorm.cfg"),
@@ -43,9 +52,46 @@ def configure_logging() -> None:
         )
 
 
-def load_test(cfg: CtormConfig, granule_list: List[CtormPreparedGranule]):
-    c_sender = CnmSender(cfg, granule_list)
-    c_sender.send_all()
+def get_granule_list(event: events.EventBridgeEvent) -> List[CtormPreparedGranule]:
+    # TODO: Implement
+    return []
+
+
+def load_test(
+    cfg: CtormConfig,
+    gr_queue_url: str,
+):
+    # fetch 10 messages from SQS
+    messages = []
+    while len(messages) < 10:
+        # Since we're trying to send x granules per invocation, we will do what we can
+        # to get a full load of 10 messages.
+        sqs_resp = get_sqs_client().receive_message(
+            QueueUrl=gr_queue_url,
+            MaxNumberOfMessages=min(10 - len(messages), 10),
+            WaitTimeSeconds=5,
+            VisibilityTimeout=3600,
+        )
+
+        batch = sqs_resp.get("Messages", [])
+        if not batch:
+            break
+
+        messages.extend(batch)
+
+    for message in sqs_resp.get("Messages", []):
+        body = json.loads(message["Body"])
+        receipt_handle = message["ReceiptHandle"]
+        granule_list = body.get("granules", [])
+
+        c_sender = CnmSender(cfg, granule_list)
+        success = c_sender.send_all()
+        if success:
+            log.debug("Deleting message %s", receipt_handle)
+            # TODO: uncomment after some dev'ing
+            # get_sqs_client().delete_message(
+            #     QueueUrl=gr_queue_url, ReceiptHandle=receipt_handle
+            # )
 
 
 def lambda_handler(event: events.EventBridgeEvent, context: context_.Context):
@@ -61,9 +107,10 @@ def lambda_handler(event: events.EventBridgeEvent, context: context_.Context):
 
     try:
         log.debug("Received event: %s", event)
-        g_list = event["Records"].pop().get("body")
-        g_list = json.loads(g_list).get("granules")
-        load_test(cfg, g_list)
+        # g_list = event["Records"].pop().get("body")
+        # g_list = json.loads(g_list).get("granules")
+
+        load_test(cfg, event["granules_queue_url"])
 
         log.info("CNM sender invocation completed")
         return {"ok": True}

ctorm/ctorm/prepare.py

@@ -10,7 +10,7 @@
 from botocore.exceptions import ClientError
 from mypy_boto3_s3.type_defs import ListObjectsV2OutputTypeDef
 
-from ctorm.config import AWS_REGION, CtormConfig, CtormPipeline
+from ctorm.config import AWS_REGION, MD5_CHECKSUM_PATTERN, CtormConfig, CtormPipeline
 
 log = getLogger(__name__)
 
@@ -164,8 +164,9 @@ def process_ummg(self, ummg: dict, ct_bkt: CtormPipeline) -> dict:
                 ]:
                     if r_urls["URL"].endswith(distr_file["Name"]):
                         # We handily have the md5 and size in the ummg
-                        filedict[K.MD5] = distr_file["Checksum"]["Value"]
+                        # TODO: double-check this test is correct and we're not unnecessarily HEADing too many files.
                         filedict[K.SIZE] = distr_file["SizeInBytes"]
+                        filedict[K.MD5] = distr_file["Checksum"]["Value"]
                         break
                     else:
                         # We must look to S3 for the size and md5
@@ -176,17 +177,17 @@ def process_ummg(self, ummg: dict, ct_bkt: CtormPipeline) -> dict:
                             )
                             log.debug("head_object: %s", headobj)
                             filedict[K.SIZE] = headobj["ContentLength"]
-                            filedict[K.MD5] = headobj["ETag"].replace('"', "")
-                            if filedict[K.MD5].endswith("-1"):
-                                # This was a multipart upload. We'll have to do something clever to get the MD5 of it.
-
-                                if self.cfg.calc_md5:
-                                    log.debug(
-                                        "multipart upload detected for %s", objloc
-                                    )
-                                    filedict[K.MD5] = self.get_real_md5(
-                                        ct_bkt, bucket, objloc
-                                    )
+
+                            md5 = headobj["ETag"].replace('"', "")
+                            if MD5_CHECKSUM_PATTERN.fullmatch(md5):
+                                filedict[K.MD5] = md5
+                            elif self.cfg.calc_md5:
+                                log.debug("multipart upload detected for %s", objloc)
+                                filedict[K.MD5] = self.get_real_md5(
+                                    ct_bkt, bucket, objloc
+                                )
+                            else:
+                                log.debug("no need to calculate md5 for %s", objloc)
                         except ClientError as e:
                             log.error("head_object failed: %s", e)
                             # TODO: trash entire message for this granule?