Setup batch filter lambda

Author: dlzhry2nhs

ack_backend/src/constants.py

@@ -3,9 +3,6 @@
 import os
 
 AUDIT_TABLE_NAME = os.getenv("AUDIT_TABLE_NAME")
-FILE_NAME_PROC_LAMBDA_NAME = os.getenv("FILE_NAME_PROC_LAMBDA_NAME")
-AUDIT_TABLE_FILENAME_GSI = "filename_index"
-AUDIT_TABLE_QUEUE_NAME_GSI = "queue_name_index"
 
 def get_source_bucket_name() -> str:
     """Get the SOURCE_BUCKET_NAME environment from environment variables."""

batch_processor_filter/poetry.lock

@@ -7,8 +7,11 @@ readme      = "README.md"
 packages    = [{include = "src"}]
 
 [tool.poetry.dependencies]
-python = "~3.11"
-coverage = "^7.9.1"
+python            = "~3.11"
+coverage          = "^7.9.1"
+aws-lambda-typing = "~2.20.0"
+boto3             = "~1.38.42"
+moto              = "^4"
 
 [build-system]
 requires      = ["poetry-core ~= 1.5.0"]

batch_processor_filter/pyproject.toml

@@ -0,0 +1,46 @@
+import boto3
+from boto3.dynamodb.conditions import Key
+
+from constants import AUDIT_TABLE_NAME, REGION_NAME, AUDIT_TABLE_FILENAME_GSI, AuditTableKeys, FileStatus, \
+    AUDIT_TABLE_QUEUE_NAME_GSI
+
+
+class BatchAuditRepository:
+    """Batch audit repository class."""
+    _DUPLICATE_CHECK_FILE_STATUS_CONDITION = (
+        Key(AuditTableKeys.STATUS).eq(FileStatus.PROCESSED)
+        | Key(AuditTableKeys.STATUS).eq(FileStatus.PREPROCESSED)
+        | Key(AuditTableKeys.STATUS).eq(FileStatus.PROCESSING)
+    )
+
+    def __init__(self):
+        self._batch_audit_table = boto3.resource("dynamodb", region_name=REGION_NAME).Table(AUDIT_TABLE_NAME)
+
+    def is_duplicate_file(self, file_key: str) -> bool:
+        matching_files = self._batch_audit_table.query(
+            IndexName=AUDIT_TABLE_FILENAME_GSI,
+            KeyConditionExpression=Key(AuditTableKeys.FILENAME).eq(file_key),
+            FilterExpression=self._DUPLICATE_CHECK_FILE_STATUS_CONDITION
+        ).get("Items", [])
+
+        return len(matching_files) > 0
+
+    def is_event_processing_for_supplier_and_vacc_type(self, supplier: str, vacc_type: str) -> bool:
+        queue_name = f"{supplier}_{vacc_type}"
+
+        files_in_processing = self._batch_audit_table.query(
+            IndexName=AUDIT_TABLE_QUEUE_NAME_GSI,
+            KeyConditionExpression=Key(AuditTableKeys.QUEUE_NAME).eq(queue_name) & Key(AuditTableKeys.STATUS)
+            .eq(FileStatus.PROCESSING)
+        ).get("Items", [])
+
+        return len(files_in_processing) > 0
+
+    def update_status(self, message_id: str, updated_status: str) -> None:
+        self._batch_audit_table.update_item(
+            Key={AuditTableKeys.MESSAGE_ID: message_id},
+            UpdateExpression="SET #status = :status",
+            ExpressionAttributeNames={"#status": "status"},
+            ExpressionAttributeValues={":status": updated_status},
+            ConditionExpression="attribute_exists(message_id)"
+        )

batch_processor_filter/src/batch_audit_repository.py

@@ -0,0 +1,10 @@
+from typing import TypedDict
+
+
+class BatchFileCreatedEvent(TypedDict):
+    message_id: str
+    vaccine_type: str
+    supplier: str
+    filename: str
+    permission: list[str]
+    created_at_formatted_string: str

batch_processor_filter/src/batch_file_created_event.py

@@ -0,0 +1,54 @@
+"""Batch processor filter service module"""
+import boto3
+import json
+import logging
+
+from batch_audit_repository import BatchAuditRepository
+from batch_file_created_event import BatchFileCreatedEvent
+from constants import REGION_NAME, FileStatus, QUEUE_URL
+from exceptions import EventAlreadyProcessingForSupplierAndVaccTypeError
+from send_log_to_firehose import send_log_to_firehose
+
+logging.basicConfig(level="INFO")
+logger = logging.getLogger()
+logger.setLevel("INFO")
+
+
+class BatchProcessorFilterService:
+    """Batch processor filter service class. Provides the business logic for the Lambda function"""
+    def __init__(self, audit_repo: BatchAuditRepository = BatchAuditRepository()):
+        self._batch_audit_repository = audit_repo
+        self._queue_client = boto3.client('sqs', region_name=REGION_NAME)
+
+    def _is_duplicate_file(self, file_key: str) -> bool:
+        """Checks if a file with the same name has already been processed or marked for processing"""
+        return self._batch_audit_repository.is_duplicate_file(file_key)
+
+    def apply_filter(self, batch_file_created_event: BatchFileCreatedEvent) -> None:
+        filename = batch_file_created_event["filename"]
+        message_id = batch_file_created_event["message_id"]
+        supplier = batch_file_created_event["supplier"]
+        vaccine_type = batch_file_created_event["vaccine_type"]
+
+        if self._is_duplicate_file(filename):
+            # Mark as processed and return without error so next event will be picked up from queue
+            logger.info("A duplicate file has already been processed. Filename: %s", filename)
+            self._batch_audit_repository.update_status(message_id, FileStatus.DUPLICATE)
+            return
+
+        if self._batch_audit_repository.is_event_processing_for_supplier_and_vacc_type(supplier, vaccine_type):
+            # Raise error so event is returned to queue and retried again later
+            logger.info("Batch event already being processed for supplier and vacc type. Filename: %s", filename)
+            raise EventAlreadyProcessingForSupplierAndVaccTypeError(f"Batch event already processing for supplier: "
+                                                                    f"{supplier} and vacc type: {vaccine_type}")
+
+        self._queue_client.send_message(
+            QueueUrl=QUEUE_URL,
+            MessageBody=json.dumps(batch_file_created_event),
+            MessageGroupId=f"{supplier}_{vaccine_type}"
+        )
+        self._batch_audit_repository.update_status(message_id, FileStatus.PROCESSING)
+
+        successful_log_message = "File forwarded for processing by ECS"
+        logger.info(successful_log_message)
+        send_log_to_firehose({**batch_file_created_event, "message": successful_log_message})

batch_processor_filter/src/batch_processor_filter_service.py

@@ -0,0 +1,29 @@
+import os
+from enum import StrEnum
+
+REGION_NAME = "eu-west-2"
+AUDIT_TABLE_NAME = os.getenv("AUDIT_TABLE_NAME")
+AUDIT_TABLE_FILENAME_GSI = os.getenv("FILE_NAME_GSI")
+AUDIT_TABLE_QUEUE_NAME_GSI = os.getenv("QUEUE_NAME_GSI")
+QUEUE_URL = os.getenv("QUEUE_URL")
+SPLUNK_FIREHOSE_STREAM_NAME = os.getenv("SPLUNK_FIREHOSE_NAME")
+
+
+class FileStatus(StrEnum):
+    """File status constants"""
+
+    QUEUED = "Queued"
+    PROCESSING = "Processing"
+    PREPROCESSED = "Preprocessed"
+    PROCESSED = "Processed"
+    DUPLICATE = "Not processed - duplicate"
+
+
+class AuditTableKeys(StrEnum):
+    """Audit table keys"""
+
+    FILENAME = "filename"
+    MESSAGE_ID = "message_id"
+    QUEUE_NAME = "queue_name"
+    STATUS = "status"
+    TIMESTAMP = "timestamp"

batch_processor_filter/src/constants.py

@@ -0,0 +1,11 @@
+"""Exceptions for the batch processor filter Lambda function"""
+
+
+class InvalidBatchSizeError(Exception):
+    """Raised when the SQS event batch size is anything other than 1"""
+    pass
+
+
+class EventAlreadyProcessingForSupplierAndVaccTypeError(Exception):
+    """Raised when there is already a batch processing event in flight for the same supplier and vaccine type"""
+    pass

batch_processor_filter/src/exceptions.py

@@ -1,4 +1,23 @@
-def lambda_handler(event, _):
-    print("hello world")
-    print(event)
-    return True
+import json
+from aws_lambda_typing import events, context
+
+from batch_file_created_event import BatchFileCreatedEvent
+from batch_processor_filter_service import BatchProcessorFilterService
+from exceptions import InvalidBatchSizeError
+
+
+service = BatchProcessorFilterService()
+
+
+def lambda_handler(event: events.SQSEvent, _: context):
+    event_records = event.get("Records", [])
+
+    # Terraform is configured so this Lambda will get a batch size of 1. We are using SQS FIFO with the message group
+    # id set to {supplier}_{vaccine_type} so we will only want batches of 1 at a time.
+    # Lambda will scale out to handle multiple message groups in parallel:
+    # https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/fifo-queue-lambda-behavior.html
+    if len(event_records) != 1:
+        raise InvalidBatchSizeError(f"Received {len(event_records)} records, expected 1")
+
+    batch_file_created_event: BatchFileCreatedEvent = json.loads(event_records[0].get("body"))
+    service.apply_filter(batch_file_created_event)

batch_processor_filter/src/lambda_handler.py

@@ -0,0 +1,15 @@
+import json
+import boto3
+
+from constants import SPLUNK_FIREHOSE_STREAM_NAME, REGION_NAME
+
+firehose_client = boto3.client("firehose", region_name=REGION_NAME)
+
+
+# Not keen on including blocking calls in function code for log data
+# Consider simply logging and setting up CW subscription filters to forward to Firehose
+# https://docs.aws.amazon.com/firehose/latest/dev/writing-with-cloudwatch-logs.html
+def send_log_to_firehose(log_data: dict) -> None:
+    """Sends the log_message to Firehose"""
+    record = {"Data": json.dumps({"event": log_data}).encode("utf-8")}
+    firehose_client.put_record(DeliveryStreamName=SPLUNK_FIREHOSE_STREAM_NAME, Record=record)