ingestion_complete

Author: JamesW1-NHS

lambdas/ack_backend/src/audit_table.py

@@ -1,5 +1,6 @@
 """Add the filename to the audit table and check for duplicates."""
 
+import time
 from typing import Optional
 
 from common.clients import dynamodb_client, logger
@@ -116,3 +117,39 @@ def increment_records_failed_count(message_id: str) -> None:
     except Exception as error:  # pylint: disable = broad-exception-caught
         logger.error(error)
         raise UnhandledAuditTableError(error) from error
+
+
+def set_audit_table_ingestion_complete(
+    file_key: str,
+    message_id: str,
+    complete_time: float,
+) -> None:
+    """Sets the ingestion_complete in the audit table to the requested time"""
+    # format the time
+    ingestion_complete = time.strftime("%Y%m%dT%H%M%S00", time.gmtime(complete_time))
+
+    update_expression = f"SET #{AuditTableKeys.INGESTION_COMPLETE} = :{AuditTableKeys.INGESTION_COMPLETE}"
+    expression_attr_names = {f"#{AuditTableKeys.INGESTION_COMPLETE}": AuditTableKeys.INGESTION_COMPLETE}
+    expression_attr_values = {f":{AuditTableKeys.INGESTION_COMPLETE}": {"S": ingestion_complete}}
+
+    try:
+        response = 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_COMPLETE).get("S")
+        logger.info(
+            "ingestion_complete for %s file, with message id %s, was successfully updated to %s in the audit table",
+            file_key,
+            message_id,
+            result,
+        )
+
+    except Exception as error:  # pylint: disable = broad-exception-caught
+        logger.error(error)
+        raise UnhandledAuditTableError(error) from error

lambdas/ack_backend/src/logging_decorators.py

@@ -5,6 +5,7 @@
 from datetime import datetime
 from functools import wraps
 
+from audit_table import set_audit_table_ingestion_complete
 from common.log_decorator import generate_and_send_logs
 
 PREFIX = "ack_processor"
@@ -78,7 +79,7 @@ def wrapper(*args, **kwargs):
             "function_name": f"{PREFIX}_{func.__name__}",
             "date_time": str(datetime.now()),
         }
-        start_time = time.time()
+        complete_time = time.time()
 
         # NB this doesn't require a try-catch block as the wrapped function never throws an exception
         result = func(*args, **kwargs)
@@ -90,10 +91,11 @@ def wrapper(*args, **kwargs):
                 "statusCode": 200,
                 "message": message_for_logs,
             }
+            file_key = base_log_data.get("filename")
+            message_id = base_log_data.get("message_id")
+            set_audit_table_ingestion_complete(file_key, message_id, complete_time)
 
-            # here: add start_time to audit table as "ingestion_complete" for message_id
-
-            generate_and_send_logs(STREAM_NAME, start_time, base_log_data, additional_log_data)
+            generate_and_send_logs(STREAM_NAME, complete_time, base_log_data, additional_log_data)
         return result
 
     return wrapper

lambdas/ack_backend/tests/test_audit_table.py

@@ -115,6 +115,7 @@ def test_set_records_succeeded_count_raises(self):
         self.mock_logger.error.assert_called_once()
 
     def test_increment_records_failed_count(self):
+        """Checks audit table correctly increments the records_failed count"""
         test_message_id = "1234"
         audit_table.increment_records_failed_count(test_message_id)
         self.mock_dynamodb_client.update_item.assert_called_once_with(
@@ -134,3 +135,30 @@ def test_increment_records_failed_count_raises(self):
             audit_table.increment_records_failed_count("msg1")
         self.assertIn("fail!", str(ctx.exception))
         self.mock_logger.error.assert_called_once()
+
+    def test_set_audit_table_ingestion_complete(self):
+        """Checks audit table correctly sets ingestion_complete to the requested value"""
+        test_file_key = "RSV_Vaccinations_v5_X26_20210730T12000000.csv"
+        test_message_id = "1234"
+        test_start_time = 1627647000
+        audit_table.set_audit_table_ingestion_complete(test_file_key, test_message_id, test_start_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_COMPLETE} = :{AuditTableKeys.INGESTION_COMPLETE}",
+            ExpressionAttributeNames={f"#{AuditTableKeys.INGESTION_COMPLETE}": AuditTableKeys.INGESTION_COMPLETE},
+            ExpressionAttributeValues={f":{AuditTableKeys.INGESTION_COMPLETE}": {"S": "20210730T12100000"}},
+            ConditionExpression="attribute_exists(message_id)",
+            ReturnValues="UPDATED_NEW",
+        )
+        self.mock_logger.info.assert_called_once()
+
+    def test_set_audit_table_ingestion_complete_throws_exception_with_invalid_id(self):
+        test_file_key = "RSV_Vaccinations_v5_X26_20210730T12000000.csv"
+        test_message_id = "1234"
+        test_start_time = 1627647000
+        self.mock_dynamodb_client.update_item.side_effect = Exception("fail!")
+        with self.assertRaises(UnhandledAuditTableError) as ctx:
+            audit_table.set_audit_table_ingestion_complete(test_file_key, test_message_id, test_start_time)
+        self.assertIn("fail!", str(ctx.exception))
+        self.mock_logger.error.assert_called_once()

lambdas/ack_backend/tests/test_splunk_logging.py

@@ -424,6 +424,7 @@ def test_splunk_update_ack_file_logged(self):
                 "update_ack_file.change_audit_table_status_to_processed"
             ) as mock_change_audit_table_status_to_processed,  # noqa: E999
             patch("update_ack_file.set_records_succeeded_count") as mock_set_records_succeeded_count,  # noqa: E999
+            patch("logging_decorators.set_audit_table_ingestion_complete"),  # noqa: E999
         ):  # noqa: E999
             result = lambda_handler(generate_event(messages, include_eof_message=True), context={})
 

lambdas/ack_backend/tests/test_update_ack_file_flow.py

@@ -41,10 +41,15 @@ def setUp(self):
         self.mock_get_record_count = self.get_record_count_patcher.start()
         self.set_records_succeeded_count_patcher = patch("update_ack_file.set_records_succeeded_count")
         self.mock_set_records_succeeded_count = self.set_records_succeeded_count_patcher.start()
+        self.set_audit_table_ingestion_complete_patcher = patch("logging_decorators.set_audit_table_ingestion_complete")
+        self.mock_set_audit_table_ingestion_complete = self.set_audit_table_ingestion_complete_patcher.start()
 
     def tearDown(self):
         self.logger_patcher.stop()
         self.change_audit_status_patcher.stop()
+        self.get_record_count_patcher.stop()
+        self.set_records_succeeded_count_patcher.stop()
+        self.set_audit_table_ingestion_complete_patcher.stop()
 
     def test_audit_table_updated_correctly_when_ack_process_complete(self):
         """VED-167 - Test that the audit table has been updated correctly"""