VED-956: Audit table update (#1026)

Author: JamesW1-NHS

lambdas/ack_backend/src/ack_processor.py

@@ -2,6 +2,7 @@
 
 import json
 
+from audit_table import increment_records_failed_count
 from common.batch.eof_utils import is_eof_message
 from convert_message_to_ack_row import convert_message_to_ack_row
 from logging_decorators import ack_lambda_handler_logging_decorator
@@ -49,6 +50,7 @@ def lambda_handler(event, _):
                 break
 
             ack_data_rows.append(convert_message_to_ack_row(message, created_at_formatted_string))
+            increment_records_failed_count(message_id)
 
     update_ack_file(file_key, created_at_formatted_string, ack_data_rows)
 

lambdas/ack_backend/src/audit_table.py

@@ -1,46 +1,114 @@
 """Add the filename to the audit table and check for duplicates."""
 
-from typing import Optional
-
 from common.clients import dynamodb_client, logger
 from common.models.errors import UnhandledAuditTableError
 from constants import AUDIT_TABLE_NAME, AuditTableKeys, FileStatus
 
+CONDITION_EXPRESSION = "attribute_exists(message_id)"
+
 
 def change_audit_table_status_to_processed(file_key: str, message_id: str) -> None:
     """Updates the status in the audit table to 'Processed' and returns the queue name."""
     try:
         # Update the status in the audit table to "Processed"
-        dynamodb_client.update_item(
+        response = dynamodb_client.update_item(
             TableName=AUDIT_TABLE_NAME,
             Key={AuditTableKeys.MESSAGE_ID: {"S": message_id}},
             UpdateExpression="SET #status = :status",
             ExpressionAttributeNames={"#status": "status"},
             ExpressionAttributeValues={":status": {"S": FileStatus.PROCESSED}},
-            ConditionExpression="attribute_exists(message_id)",
+            ConditionExpression=CONDITION_EXPRESSION,
+            ReturnValues="UPDATED_NEW",
         )
-
+        result = response.get("Attributes", {}).get("status").get("S")
         logger.info(
             "The status of %s file, with message id %s, was successfully updated to %s in the audit table",
             file_key,
             message_id,
-            FileStatus.PROCESSED,
+            result,
         )
 
     except Exception as error:  # pylint: disable = broad-exception-caught
         logger.error(error)
         raise UnhandledAuditTableError(error) from error
 
 
-def get_record_count_by_message_id(event_message_id: str) -> Optional[int]:
-    """Retrieves full audit entry by unique event message ID"""
+def get_record_count_and_failures_by_message_id(event_message_id: str) -> tuple[int, int]:
+    """Retrieves total record count and total failures by unique event message ID"""
     audit_record = dynamodb_client.get_item(
         TableName=AUDIT_TABLE_NAME, Key={AuditTableKeys.MESSAGE_ID: {"S": event_message_id}}
     )
 
     record_count = audit_record.get("Item", {}).get(AuditTableKeys.RECORD_COUNT, {}).get("N")
+    failures_count = audit_record.get("Item", {}).get(AuditTableKeys.RECORDS_FAILED, {}).get("N")
+
+    return int(record_count) if record_count else 0, int(failures_count) if failures_count else 0
+
+
+def increment_records_failed_count(message_id: str) -> None:
+    """
+    Increment a counter attribute safely, handling the case where it might not exist.
+    From https://docs.aws.amazon.com/code-library/latest/ug/dynamodb_example_dynamodb_Scenario_AtomicCounterOperations_section.html
+    """
+    increment_value = 1
+    initial_value = 0
+
+    try:
+        # Use SET with if_not_exists to safely increment the counter attribute
+        dynamodb_client.update_item(
+            TableName=AUDIT_TABLE_NAME,
+            Key={AuditTableKeys.MESSAGE_ID: {"S": message_id}},
+            UpdateExpression="SET #attribute = if_not_exists(#attribute, :initial) + :increment",
+            ExpressionAttributeNames={"#attribute": AuditTableKeys.RECORDS_FAILED},
+            ExpressionAttributeValues={":increment": {"N": str(increment_value)}, ":initial": {"N": str(initial_value)}},
+            ConditionExpression=CONDITION_EXPRESSION,
+            ReturnValues="UPDATED_NEW",
+        )
+
+    except Exception as error:  # pylint: disable = broad-exception-caught
+        logger.error(error)
+        raise UnhandledAuditTableError(error) from error
 
-    if not record_count:
-        return None
 
-    return int(record_count)
+def set_audit_record_success_count_and_end_time(
+    file_key: str, message_id: str, success_count: int, ingestion_end_time: str
+) -> None:
+    """Sets the 'records_succeeded' and 'ingestion_end_time' attributes for the given audit record"""
+    update_expression = (
+        f"SET #{AuditTableKeys.INGESTION_END_TIME} = :{AuditTableKeys.INGESTION_END_TIME}"
+        f", #{AuditTableKeys.RECORDS_SUCCEEDED} = :{AuditTableKeys.RECORDS_SUCCEEDED}"
+    )
+    expression_attr_names = {
+        f"#{AuditTableKeys.INGESTION_END_TIME}": AuditTableKeys.INGESTION_END_TIME,
+        f"#{AuditTableKeys.RECORDS_SUCCEEDED}": AuditTableKeys.RECORDS_SUCCEEDED,
+    }
+    expression_attr_values = {
+        f":{AuditTableKeys.INGESTION_END_TIME}": {"S": ingestion_end_time},
+        f":{AuditTableKeys.RECORDS_SUCCEEDED}": {"N": str(success_count)},
+    }
+
+    try:
+        dynamodb_client.update_item(
+            TableName=AUDIT_TABLE_NAME,
+            Key={AuditTableKeys.MESSAGE_ID: {"S": message_id}},
+            UpdateExpression=update_expression,
+            ExpressionAttributeNames=expression_attr_names,
+            ExpressionAttributeValues=expression_attr_values,
+            ConditionExpression=CONDITION_EXPRESSION,
+        )
+    except Exception as error:  # pylint: disable = broad-exception-caught
+        logger.error(error)
+        raise UnhandledAuditTableError(error) from error
+
+    logger.info(
+        "ingestion_end_time for %s file, with message id %s, was successfully updated to %s in the audit table",
+        file_key,
+        message_id,
+        ingestion_end_time,
+    )
+    logger.info(
+        "records_succeeded for %s file, with message id %s, was successfully updated to %s in the audit table",
+        file_key,
+        message_id,
+        str(success_count),
+    )

lambdas/ack_backend/src/constants.py

@@ -38,6 +38,9 @@ class AuditTableKeys:
     RECORD_COUNT = "record_count"
     STATUS = "status"
     TIMESTAMP = "timestamp"
+    INGESTION_END_TIME = "ingestion_end_time"
+    RECORDS_SUCCEEDED = "records_succeeded"
+    RECORDS_FAILED = "records_failed"
 
 
 ACK_HEADERS = [

lambdas/ack_backend/src/logging_decorators.py

@@ -91,6 +91,7 @@ def wrapper(*args, **kwargs):
                 "message": message_for_logs,
             }
             generate_and_send_logs(STREAM_NAME, start_time, base_log_data, additional_log_data)
+
         return result
 
     return wrapper

lambdas/ack_backend/src/update_ack_file.py

@@ -1,10 +1,15 @@
 """Functions for uploading the data to the ack file"""
 
+import time
 from io import BytesIO, StringIO
 
 from botocore.exceptions import ClientError
 
-from audit_table import change_audit_table_status_to_processed, get_record_count_by_message_id
+from audit_table import (
+    change_audit_table_status_to_processed,
+    get_record_count_and_failures_by_message_id,
+    set_audit_record_success_count_and_end_time,
+)
 from common.aws_s3_utils import move_file
 from common.clients import get_s3_client, logger
 from constants import (
@@ -71,15 +76,22 @@ def complete_batch_file_process(
         get_source_bucket_name(), f"{BATCH_FILE_PROCESSING_DIR}/{file_key}", f"{BATCH_FILE_ARCHIVE_DIR}/{file_key}"
     )
 
-    total_ack_rows_processed = get_record_count_by_message_id(message_id)
+    total_ack_rows_processed, total_failures = get_record_count_and_failures_by_message_id(message_id)
     change_audit_table_status_to_processed(file_key, message_id)
 
+    # Consider creating time utils and using datetime instead of time
+    ingestion_end_time = time.strftime("%Y%m%dT%H%M%S00", time.gmtime())
+    successful_record_count = total_ack_rows_processed - total_failures
+    set_audit_record_success_count_and_end_time(file_key, message_id, successful_record_count, ingestion_end_time)
+
     return {
         "message_id": message_id,
         "file_key": file_key,
         "supplier": supplier,
         "vaccine_type": vaccine_type,
         "row_count": total_ack_rows_processed,
+        "success_count": successful_record_count,
+        "failure_count": total_failures,
     }