Added ITOC log and minor fixups

Author: dlzhry2nhs

lambdas/ack_backend/src/audit_table.py

@@ -1,8 +1,5 @@
 """Add the filename to the audit table and check for duplicates."""
 
-import time
-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
@@ -36,63 +33,26 @@ def change_audit_table_status_to_processed(file_key: str, message_id: str) -> No
         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")
 
-    if not record_count:
-        return None
-
-    return int(record_count)
-
-
-def set_records_succeeded_count(message_id: str) -> None:
-    """Set the 'records_succeeded' item in the audit table entry"""
-    audit_record = dynamodb_client.get_item(
-        TableName=AUDIT_TABLE_NAME, Key={AuditTableKeys.MESSAGE_ID: {"S": message_id}}
-    )
-    record_count_item = audit_record.get("Item", {}).get(AuditTableKeys.RECORD_COUNT, {}).get("N")
-    records_failed_item = audit_record.get("Item", {}).get(AuditTableKeys.RECORDS_FAILED, {}).get("N")
-
-    record_count = int(record_count_item) if record_count_item else 0
-    records_failed = int(records_failed_item) if records_failed_item else 0
-    records_succeeded = record_count - records_failed
-
-    try:
-        response = dynamodb_client.update_item(
-            TableName=AUDIT_TABLE_NAME,
-            Key={AuditTableKeys.MESSAGE_ID: {"S": message_id}},
-            UpdateExpression="SET #attribute = :value",
-            ExpressionAttributeNames={"#attribute": AuditTableKeys.RECORDS_SUCCEEDED},
-            ExpressionAttributeValues={":value": {"N": str(records_succeeded)}},
-            ConditionExpression=CONDITION_EXPRESSION,
-            ReturnValues="UPDATED_NEW",
-        )
-        result = response.get("Attributes", {}).get(AuditTableKeys.RECORDS_SUCCEEDED).get("N")
-        logger.info(
-            "Attribute %s for message id %s set to %s in the audit table",
-            AuditTableKeys.RECORDS_SUCCEEDED,
-            message_id,
-            result,
-        )
-
-    except Exception as error:  # pylint: disable = broad-exception-caught
-        logger.error(error)
-        raise UnhandledAuditTableError(error) from error
+    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(
@@ -110,37 +70,45 @@ def increment_records_failed_count(message_id: str) -> None:
         raise UnhandledAuditTableError(error) from error
 
 
-def set_audit_table_ingestion_end_time(
-    file_key: str,
-    message_id: str,
-    complete_time: float,
+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 ingestion_end_time in the audit table to the requested time"""
-    # format the time
-    ingestion_end_time = time.strftime("%Y%m%dT%H%M%S00", time.gmtime(complete_time))
-
-    update_expression = f"SET #{AuditTableKeys.INGESTION_END_TIME} = :{AuditTableKeys.INGESTION_END_TIME}"
-    expression_attr_names = {f"#{AuditTableKeys.INGESTION_END_TIME}": AuditTableKeys.INGESTION_END_TIME}
-    expression_attr_values = {f":{AuditTableKeys.INGESTION_END_TIME}": {"S": ingestion_end_time}}
+    """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:
-        response = dynamodb_client.update_item(
+        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=f"attribute_exists({AuditTableKeys.MESSAGE_ID})",
-            ReturnValues="UPDATED_NEW",
-        )
-        result = response.get("Attributes", {}).get(AuditTableKeys.INGESTION_END_TIME).get("S")
-        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,
-            result,
+            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/logging_decorators.py

@@ -78,9 +78,10 @@ def wrapper(*args, **kwargs):
             "function_name": f"{PREFIX}_{func.__name__}",
             "date_time": str(datetime.now()),
         }
+        start_time = time.time()
 
         # NB this doesn't require a try-catch block as the wrapped function never throws an exception
-        result, ingestion_end_time = func(*args, **kwargs)
+        result = func(*args, **kwargs)
         if result is not None:
             message_for_logs = "Record processing complete"
             base_log_data.update(result)
@@ -89,7 +90,7 @@ def wrapper(*args, **kwargs):
                 "statusCode": 200,
                 "message": message_for_logs,
             }
-            generate_and_send_logs(STREAM_NAME, ingestion_end_time, base_log_data, additional_log_data)
+            generate_and_send_logs(STREAM_NAME, start_time, base_log_data, additional_log_data)
 
         return result
 

lambdas/ack_backend/src/update_ack_file.py

@@ -7,9 +7,8 @@
 
 from audit_table import (
     change_audit_table_status_to_processed,
-    get_record_count_by_message_id,
-    set_audit_table_ingestion_end_time,
-    set_records_succeeded_count,
+    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
@@ -67,7 +66,7 @@ def complete_batch_file_process(
     vaccine_type: str,
     created_at_formatted_string: str,
     file_key: str,
-) -> tuple[dict, float]:
+) -> dict:
     """Mark the batch file as processed. This involves moving the ack and original file to destinations and updating
     the audit table status"""
     ack_filename = f"{file_key.replace('.csv', f'_BusAck_{created_at_formatted_string}.csv')}"
@@ -77,21 +76,23 @@ 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)
-    set_records_succeeded_count(message_id)
 
-    ingestion_end_time = time.time()
-    set_audit_table_ingestion_end_time(file_key, message_id, ingestion_end_time)
+    # 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)
 
-    result = {
+    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,
     }
-    return result, ingestion_end_time
 
 
 def obtain_current_ack_content(temp_ack_file_key: str) -> StringIO:

lambdas/ack_backend/tests/test_audit_table.py

@@ -1,5 +1,5 @@
 import unittest
-from unittest.mock import patch
+from unittest.mock import call, patch
 
 import audit_table
 from common.models.errors import UnhandledAuditTableError
@@ -37,82 +37,29 @@ def test_change_audit_table_status_to_processed_raises(self):
         self.assertIn("fail!", str(ctx.exception))
         self.mock_logger.error.assert_called_once()
 
-    def test_get_record_count_by_message_id_returns_the_record_count(self):
-        """Test that get_record_count_by_message_id retrieves the integer value of the total record count"""
+    def test_get_record_count_and_failures_by_message_id_returns_the_record_count_and_failures(self):
+        """Test that get_record_count_by_message_id retrieves the integer values of the total record count and
+        failures"""
         test_message_id = "1234"
-
         self.mock_dynamodb_client.get_item.return_value = {
-            "Item": {"message_id": {"S": test_message_id}, "record_count": {"N": "1000"}}
+            "Item": {"message_id": {"S": test_message_id}, "record_count": {"N": "1000"}, "records_failed": {"N": "5"}}
         }
 
-        self.assertEqual(audit_table.get_record_count_by_message_id(test_message_id), 1000)
-
-    def test_get_record_count_by_message_id_returns_none_if_record_count_not_set(self):
-        """Test that if the record count has not yet been set on the audit item then None is returned"""
-        test_message_id = "1234"
-
-        self.mock_dynamodb_client.get_item.return_value = {"Item": {"message_id": {"S": test_message_id}}}
+        record_count, failed_count = audit_table.get_record_count_and_failures_by_message_id(test_message_id)
 
-        self.assertIsNone(audit_table.get_record_count_by_message_id(test_message_id))
+        self.assertEqual(record_count, 1000)
+        self.assertEqual(failed_count, 5)
 
-    def test_set_records_succeeded_count(self):
+    def test_get_record_count_and_failures_by_message_id_returns_zero_if_values_not_set(self):
+        """Test that if the record count has not yet been set on the audit item then zero is returned"""
         test_message_id = "1234"
-        self.mock_dynamodb_client.get_item.return_value = {
-            "Item": {"message_id": {"S": test_message_id}, "record_count": {"N": "1000"}, "records_failed": {"N": "42"}}
-        }
-        audit_table.set_records_succeeded_count(test_message_id)
-        self.mock_dynamodb_client.get_item.assert_called_once()
-        self.mock_dynamodb_client.update_item.assert_called_once_with(
-            TableName=AUDIT_TABLE_NAME,
-            Key={AuditTableKeys.MESSAGE_ID: {"S": test_message_id}},
-            UpdateExpression="SET #attribute = :value",
-            ExpressionAttributeNames={"#attribute": AuditTableKeys.RECORDS_SUCCEEDED},
-            ExpressionAttributeValues={":value": {"N": "958"}},
-            ConditionExpression="attribute_exists(message_id)",
-            ReturnValues="UPDATED_NEW",
-        )
-        self.mock_logger.info.assert_called_once()
 
-    def test_set_records_succeeded_count_no_failures(self):
-        test_message_id = "1234"
-        self.mock_dynamodb_client.get_item.return_value = {
-            "Item": {"message_id": {"S": test_message_id}, "record_count": {"N": "1000"}}
-        }
-        audit_table.set_records_succeeded_count(test_message_id)
-        self.mock_dynamodb_client.get_item.assert_called_once()
-        self.mock_dynamodb_client.update_item.assert_called_once_with(
-            TableName=AUDIT_TABLE_NAME,
-            Key={AuditTableKeys.MESSAGE_ID: {"S": test_message_id}},
-            UpdateExpression="SET #attribute = :value",
-            ExpressionAttributeNames={"#attribute": AuditTableKeys.RECORDS_SUCCEEDED},
-            ExpressionAttributeValues={":value": {"N": "1000"}},
-            ConditionExpression="attribute_exists(message_id)",
-            ReturnValues="UPDATED_NEW",
-        )
-        self.mock_logger.info.assert_called_once()
-
-    def test_set_records_succeeded_count_no_records(self):
-        test_message_id = "1234"
         self.mock_dynamodb_client.get_item.return_value = {"Item": {"message_id": {"S": test_message_id}}}
-        audit_table.set_records_succeeded_count(test_message_id)
-        self.mock_dynamodb_client.get_item.assert_called_once()
-        self.mock_dynamodb_client.update_item.assert_called_once_with(
-            TableName=AUDIT_TABLE_NAME,
-            Key={AuditTableKeys.MESSAGE_ID: {"S": test_message_id}},
-            UpdateExpression="SET #attribute = :value",
-            ExpressionAttributeNames={"#attribute": AuditTableKeys.RECORDS_SUCCEEDED},
-            ExpressionAttributeValues={":value": {"N": "0"}},
-            ConditionExpression="attribute_exists(message_id)",
-            ReturnValues="UPDATED_NEW",
-        )
-        self.mock_logger.info.assert_called_once()
 
-    def test_set_records_succeeded_count_raises(self):
-        self.mock_dynamodb_client.update_item.side_effect = Exception("fail!")
-        with self.assertRaises(UnhandledAuditTableError) as ctx:
-            audit_table.set_records_succeeded_count("msg1")
-        self.assertIn("fail!", str(ctx.exception))
-        self.mock_logger.error.assert_called_once()
+        record_count, failed_count = audit_table.get_record_count_and_failures_by_message_id(test_message_id)
+
+        self.assertEqual(record_count, 0)
+        self.assertEqual(failed_count, 0)
 
     def test_increment_records_failed_count(self):
         """Checks audit table correctly increments the records_failed count"""
@@ -135,29 +82,61 @@ def test_increment_records_failed_count_raises(self):
         self.assertIn("fail!", str(ctx.exception))
         self.mock_logger.error.assert_called_once()
 
-    def test_set_audit_table_ingestion_end_time(self):
-        """Checks audit table correctly sets ingestion_end_time to the requested value"""
+    def test_set_audit_record_success_count_and_end_time(self):
+        """Checks audit table correctly sets ingestion_end_time and success count to the requested value"""
         test_file_key = "RSV_Vaccinations_v5_X26_20210730T12000000.csv"
         test_message_id = "1234"
-        test_end_time = 1627647000
-        audit_table.set_audit_table_ingestion_end_time(test_file_key, test_message_id, test_end_time)
+        test_end_time = "20251208T14430000"
+        test_success_count = 5
+
+        audit_table.set_audit_record_success_count_and_end_time(
+            test_file_key, test_message_id, test_success_count, test_end_time
+        )
+
         self.mock_dynamodb_client.update_item.assert_called_once_with(
             TableName=AUDIT_TABLE_NAME,
             Key={AuditTableKeys.MESSAGE_ID: {"S": test_message_id}},
-            UpdateExpression=f"SET #{AuditTableKeys.INGESTION_END_TIME} = :{AuditTableKeys.INGESTION_END_TIME}",
-            ExpressionAttributeNames={f"#{AuditTableKeys.INGESTION_END_TIME}": AuditTableKeys.INGESTION_END_TIME},
-            ExpressionAttributeValues={f":{AuditTableKeys.INGESTION_END_TIME}": {"S": "20210730T12100000"}},
+            UpdateExpression=(
+                f"SET #{AuditTableKeys.INGESTION_END_TIME} = :{AuditTableKeys.INGESTION_END_TIME}"
+                f", #{AuditTableKeys.RECORDS_SUCCEEDED} = :{AuditTableKeys.RECORDS_SUCCEEDED}"
+            ),
+            ExpressionAttributeNames={
+                f"#{AuditTableKeys.INGESTION_END_TIME}": AuditTableKeys.INGESTION_END_TIME,
+                f"#{AuditTableKeys.RECORDS_SUCCEEDED}": AuditTableKeys.RECORDS_SUCCEEDED,
+            },
+            ExpressionAttributeValues={
+                f":{AuditTableKeys.INGESTION_END_TIME}": {"S": test_end_time},
+                f":{AuditTableKeys.RECORDS_SUCCEEDED}": {"N": str(test_success_count)},
+            },
             ConditionExpression="attribute_exists(message_id)",
-            ReturnValues="UPDATED_NEW",
         )
-        self.mock_logger.info.assert_called_once()
+        self.mock_logger.info.assert_has_calls(
+            [
+                call(
+                    "ingestion_end_time for %s file, with message id %s, was successfully updated to %s in the audit table",
+                    "RSV_Vaccinations_v5_X26_20210730T12000000.csv",
+                    "1234",
+                    "20251208T14430000",
+                ),
+                call(
+                    "records_succeeded for %s file, with message id %s, was successfully updated to %s in the audit table",
+                    "RSV_Vaccinations_v5_X26_20210730T12000000.csv",
+                    "1234",
+                    "5",
+                ),
+            ]
+        )
 
-    def test_set_audit_table_ingestion_end_time_throws_exception_with_invalid_id(self):
+    def test_set_audit_record_success_count_and_end_time_throws_exception_with_invalid_id(self):
         test_file_key = "RSV_Vaccinations_v5_X26_20210730T12000000.csv"
         test_message_id = "1234"
-        test_end_time = 1627647000
-        self.mock_dynamodb_client.update_item.side_effect = Exception("fail!")
+        test_end_time = "20251208T14430000"
+        test_success_count = 5
+        self.mock_dynamodb_client.update_item.side_effect = Exception("Unhandled error")
+
         with self.assertRaises(UnhandledAuditTableError) as ctx:
-            audit_table.set_audit_table_ingestion_end_time(test_file_key, test_message_id, test_end_time)
-        self.assertIn("fail!", str(ctx.exception))
+            audit_table.set_audit_record_success_count_and_end_time(
+                test_file_key, test_message_id, test_success_count, test_end_time
+            )
+        self.assertIn("Unhandled error", str(ctx.exception))
         self.mock_logger.error.assert_called_once()