VED-744 Batch Audit Table TTL (#790)

* init: ExpiresAt added to audit table

* review fixes

* review fixes II

Author: JamesW1-NHS

filenameprocessor/src/audit_table.py

@@ -8,6 +8,7 @@ def upsert_audit_table(
     message_id: str,
     file_key: str,
     created_at_formatted_str: str,
+    expiry_timestamp: int,
     queue_name: str,
     file_status: str
 ) -> None:
@@ -24,6 +25,7 @@ def upsert_audit_table(
                 AuditTableKeys.QUEUE_NAME: {"S": queue_name},
                 AuditTableKeys.STATUS: {"S": file_status},
                 AuditTableKeys.TIMESTAMP: {"S": created_at_formatted_str},
+                AuditTableKeys.EXPIRES_AT: {"N": str(expiry_timestamp)},
             },
             ConditionExpression="attribute_not_exists(message_id)",  # Prevents accidental overwrites
         )

filenameprocessor/src/constants.py

@@ -14,6 +14,7 @@
 
 SOURCE_BUCKET_NAME = os.getenv("SOURCE_BUCKET_NAME")
 AUDIT_TABLE_NAME = os.getenv("AUDIT_TABLE_NAME")
+AUDIT_TABLE_TTL_DAYS = os.getenv("AUDIT_TABLE_TTL_DAYS")
 VALID_VERSIONS = ["V5"]
 
 SUPPLIER_PERMISSIONS_HASH_KEY = "supplier_permissions"
@@ -48,3 +49,4 @@ class AuditTableKeys(StrEnum):
     QUEUE_NAME = "queue_name"
     STATUS = "status"
     TIMESTAMP = "timestamp"
+    EXPIRES_AT = "expires_at"

filenameprocessor/src/file_name_processor.py

@@ -8,7 +8,7 @@
 
 import argparse
 from uuid import uuid4
-from utils_for_filenameprocessor import get_created_at_formatted_string, move_file
+from utils_for_filenameprocessor import get_creation_and_expiry_times, move_file
 from file_key_validation import validate_file_key, is_file_in_directory_root
 from send_sqs_message import make_and_send_sqs_message
 from make_and_upload_ack_file import make_and_upload_the_ack_file
@@ -63,14 +63,14 @@ def handle_record(record) -> dict:
 
     try:
         message_id = str(uuid4())
-        created_at_formatted_string = get_created_at_formatted_string(bucket_name, file_key)
+        created_at_formatted_string, expiry_timestamp = get_creation_and_expiry_times(bucket_name, file_key)
 
         vaccine_type, supplier = validate_file_key(file_key)
         permissions = validate_vaccine_type_permissions(vaccine_type=vaccine_type, supplier=supplier)
 
         queue_name = f"{supplier}_{vaccine_type}"
         upsert_audit_table(
-            message_id, file_key, created_at_formatted_string, queue_name, FileStatus.QUEUED
+            message_id, file_key, created_at_formatted_string, expiry_timestamp, queue_name, FileStatus.QUEUED
         )
         make_and_send_sqs_message(
             file_key, message_id, permissions, vaccine_type, supplier, created_at_formatted_string
@@ -100,7 +100,7 @@ def handle_record(record) -> dict:
 
         queue_name = f"{supplier}_{vaccine_type}"
         upsert_audit_table(
-            message_id, file_key, created_at_formatted_string, queue_name, FileStatus.PROCESSED
+            message_id, file_key, created_at_formatted_string, expiry_timestamp, queue_name, FileStatus.PROCESSED
         )
 
         # Create ack file

filenameprocessor/src/utils_for_filenameprocessor.py

@@ -1,11 +1,16 @@
 """Utils for filenameprocessor lambda"""
+from datetime import timedelta
 from clients import s3_client, logger
+from constants import AUDIT_TABLE_TTL_DAYS
 
 
-def get_created_at_formatted_string(bucket_name: str, file_key: str) -> str:
-    """Get the created_at_formatted_string from the response"""
+def get_creation_and_expiry_times(bucket_name: str, file_key: str) -> (str, int):
+    """Get 'created_at_formatted_string' and 'expires_at' from the response"""
     response = s3_client.get_object(Bucket=bucket_name, Key=file_key)
-    return response["LastModified"].strftime("%Y%m%dT%H%M%S00")
+    creation_datetime = response["LastModified"]
+    expiry_datetime = creation_datetime + timedelta(days=int(AUDIT_TABLE_TTL_DAYS))
+    expiry_timestamp = int(expiry_datetime.timestamp())
+    return creation_datetime.strftime("%Y%m%dT%H%M%S00"), expiry_timestamp
 
 
 def move_file(bucket_name: str, source_file_key: str, destination_file_key: str) -> None:

filenameprocessor/tests/test_audit_table.py

@@ -54,7 +54,8 @@ def test_upsert_audit_table(self):
             file_key=ravs_rsv_test_file.file_key,
             created_at_formatted_str=ravs_rsv_test_file.created_at_formatted_string,
             queue_name=ravs_rsv_test_file.queue_name,
-            file_status=FileStatus.PROCESSED
+            file_status=FileStatus.PROCESSED,
+            expiry_timestamp=ravs_rsv_test_file.expires_at
         )
 
         assert_audit_table_entry(ravs_rsv_test_file, FileStatus.PROCESSED)
@@ -68,7 +69,8 @@ def test_upsert_audit_table_with_duplicate_message_id_raises_exception(self):
             file_key=ravs_rsv_test_file.file_key,
             created_at_formatted_str=ravs_rsv_test_file.created_at_formatted_string,
             queue_name=ravs_rsv_test_file.queue_name,
-            file_status=FileStatus.PROCESSED
+            file_status=FileStatus.PROCESSED,
+            expiry_timestamp=ravs_rsv_test_file.expires_at
         )
 
         assert_audit_table_entry(ravs_rsv_test_file, FileStatus.PROCESSED)
@@ -80,4 +82,5 @@ def test_upsert_audit_table_with_duplicate_message_id_raises_exception(self):
                 created_at_formatted_str=ravs_rsv_test_file.created_at_formatted_string,
                 queue_name=ravs_rsv_test_file.queue_name,
                 file_status=FileStatus.PROCESSED,
+                expiry_timestamp=ravs_rsv_test_file.expires_at
             )

filenameprocessor/tests/test_lambda_handler.py

@@ -18,7 +18,7 @@
     MOCK_ODS_CODE_TO_SUPPLIER
 )
 from tests.utils_for_tests.mock_environment_variables import MOCK_ENVIRONMENT_DICT, BucketNames, Sqs
-from tests.utils_for_tests.values_for_tests import MOCK_CREATED_AT_FORMATTED_STRING, MockFileDetails
+from tests.utils_for_tests.values_for_tests import MOCK_CREATED_AT_FORMATTED_STRING, MOCK_EXPIRES_AT, MockFileDetails
 
 # Ensure environment variables are mocked before importing from src files
 with patch.dict("os.environ", MOCK_ENVIRONMENT_DICT):
@@ -62,10 +62,11 @@ def run(self, result=None):
 
         # Set up common patches to be applied to all tests in the class (these can be overridden in individual tests.)
         common_patches = [
-            # Patch get_created_at_formatted_string, so that the ack file key can be deduced  (it is already unittested
+            # Patch get_creation_and_expiry_times, so that the ack file key can be deduced  (it is already unittested
             # separately). Note that files numbered '1', which are predominantly used in these tests, use the
             # MOCK_CREATED_AT_FORMATTED_STRING.
-            patch("file_name_processor.get_created_at_formatted_string", return_value=MOCK_CREATED_AT_FORMATTED_STRING),
+            patch("file_name_processor.get_creation_and_expiry_times",
+                  return_value=(MOCK_CREATED_AT_FORMATTED_STRING, MOCK_EXPIRES_AT)),
             # Patch redis_client to use a fake redis client.
             patch("elasticache.redis_client", new=fakeredis.FakeStrictRedis()),
             # Patch the permissions config to allow all suppliers full permissions for all vaccine types.
@@ -242,6 +243,7 @@ def test_lambda_invalid_file_key_no_other_files_in_queue(self):
                 "queue_name": {"S": "unknown_unknown"},
                 "status": {"S": "Processed"},
                 "timestamp": {"S": file_details.created_at_formatted_string},
+                "expires_at": {"N": str(file_details.expires_at)},
             }
         ]
         self.assertEqual(self.get_audit_table_items(), expected_table_items)

filenameprocessor/tests/test_utils_for_filenameprocessor.py

@@ -2,7 +2,7 @@
 
 from unittest import TestCase
 from unittest.mock import patch
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from moto import mock_s3
 from boto3 import client as boto3_client
 
@@ -11,9 +11,10 @@
 
 # Ensure environment variables are mocked before importing from src files
 with patch.dict("os.environ", MOCK_ENVIRONMENT_DICT):
+    from constants import AUDIT_TABLE_TTL_DAYS
     from clients import REGION_NAME
     from utils_for_filenameprocessor import (
-        get_created_at_formatted_string,
+        get_creation_and_expiry_times,
         move_file
     )
 
@@ -32,20 +33,24 @@ def tearDown(self):
         """Tear down the s3 buckets"""
         GenericTearDown(s3_client)
 
-    def test_get_created_at_formatted_string(self):
-        """Test that get_created_at_formatted_string can correctly get the created_at_formatted_string"""
+    def test_get_creation_and_expiry_times(self):
+        """Test that get_creation_and_expiry_times can correctly get the created_at_formatted_string"""
         bucket_name = BucketNames.SOURCE
         file_key = "test_file_key"
 
         s3_client.put_object(Bucket=bucket_name, Key=file_key)
 
-        mock_last_modified = {"LastModified": datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)}
-        expected_result = "20240101T12000000"
+        mock_last_modified_created_at = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
+        mock_last_modified = {"LastModified": mock_last_modified_created_at}
+        expected_result_created_at = "20240101T12000000"
+        expected_expiry_datetime = mock_last_modified_created_at + timedelta(days=int(AUDIT_TABLE_TTL_DAYS))
+        expected_result_expires_at = int(expected_expiry_datetime.timestamp())
 
         with patch("utils_for_filenameprocessor.s3_client.get_object", return_value=mock_last_modified):
-            created_at_formatted_string = get_created_at_formatted_string(bucket_name, file_key)
+            created_at_formatted_string, expires_at = get_creation_and_expiry_times(bucket_name, file_key)
 
-        self.assertEqual(created_at_formatted_string, expected_result)
+        self.assertEqual(created_at_formatted_string, expected_result_created_at)
+        self.assertEqual(expires_at, expected_result_expires_at)
 
     def test_move_file(self):
         """Tests that move_file correctly moves a file from one location to another within a single S3 bucket"""

filenameprocessor/tests/utils_for_tests/mock_environment_variables.py

@@ -40,5 +40,6 @@ class Sqs:
     "REDIS_HOST": "localhost",
     "REDIS_PORT": "6379",
     "SPLUNK_FIREHOSE_NAME": Firehose.STREAM_NAME,
-    "AUDIT_TABLE_NAME": "immunisation-batch-internal-dev-audit-table"
+    "AUDIT_TABLE_NAME": "immunisation-batch-internal-dev-audit-table",
+    "AUDIT_TABLE_TTL_DAYS": "14",
 }

filenameprocessor/tests/utils_for_tests/values_for_tests.py

@@ -16,6 +16,7 @@
 # NOTE That that file details for files numbered anything other than one will have a different
 # created_at_formatted_string (the final digit of the year will be the file number, rather than '1')
 MOCK_CREATED_AT_FORMATTED_STRING = "20010101T00000000"
+MOCK_EXPIRES_AT = 947808000
 
 
 class FileDetails:
@@ -33,6 +34,7 @@ def __init__(self, supplier: str, vaccine_type: str, ods_code: str, file_number:
         self.queue_name = f"{self.supplier}_{self.vaccine_type}"
 
         self.created_at_formatted_string = f"200{file_number}0101T00000000"
+        self.expires_at = MOCK_EXPIRES_AT
         self.message_id = f"{self.supplier}_{self.vaccine_type}_test_id_{file_number}"
         self.name = f"{self.vaccine_type}/ {self.supplier} file"
 
@@ -61,6 +63,7 @@ def __init__(self, supplier: str, vaccine_type: str, ods_code: str, file_number:
             AuditTableKeys.FILENAME: {"S": self.file_key},
             AuditTableKeys.QUEUE_NAME: {"S": self.queue_name},
             AuditTableKeys.TIMESTAMP: {"S": self.created_at_formatted_string},
+            AuditTableKeys.EXPIRES_AT: {"N": str(self.expires_at)},
         }
 
 

terraform/dynamodb.tf

@@ -24,6 +24,11 @@ resource "aws_dynamodb_table" "audit-table" {
     type = "S"
   }
 
+  ttl {
+    attribute_name = "expires_at"
+    enabled        = true
+  }
+
   global_secondary_index {
     name            = "filename_index"
     hash_key        = "filename"