refactor filenameprocessor file path into two

Author: Akol125

lambdas/filenameprocessor/src/constants.py

@@ -11,7 +11,7 @@
 )
 
 SOURCE_BUCKET_NAME = os.getenv("SOURCE_BUCKET_NAME")
-DPS_DESTINATION_BUCKET_NAME = os.getenv("DPS_DESTINATION_BUCKET_NAME")
+DPS_DESTINATION_BUCKET_NAME = os.getenv("ACK_BUCKET_NAME")
 AUDIT_TABLE_NAME = os.getenv("AUDIT_TABLE_NAME")
 AUDIT_TABLE_TTL_DAYS = os.getenv("AUDIT_TABLE_TTL_DAYS")
 VALID_VERSIONS = ["V5"]

lambdas/filenameprocessor/src/file_name_processor.py

@@ -55,8 +55,6 @@ def handle_record(record) -> dict:
             "error": str(error),
         }
 
-    vaccine_type = "unknown"
-    supplier = "unknown"
     expiry_timestamp = "unknown"
 
     if bucket_name != SOURCE_BUCKET_NAME:
@@ -74,56 +72,128 @@ def handle_record(record) -> dict:
     message_id = "Message id was not created"
     created_at_formatted_string = "created_at_time not identified"
 
-    try:
-        message_id = str(uuid4())
-        s3_response = get_s3_client().get_object(Bucket=bucket_name, Key=file_key)
-        created_at_formatted_string, expiry_timestamp = get_creation_and_expiry_times(s3_response)
+    message_id = str(uuid4())
+    s3_response = get_s3_client().get_object(Bucket=bucket_name, Key=file_key)
+    created_at_formatted_string, expiry_timestamp = get_creation_and_expiry_times(s3_response)
 
-        if file_key.startswith(EXTENDED_ATTRIBUTES_PREFIXES):
-            queue_name = validate_extended_attributes_file_key(file_key)
-            move_file_outside_bucket(bucket_name, file_key, DPS_DESTINATION_BUCKET_NAME, f"archive/{file_key}")
+    if file_key.startswith(EXTENDED_ATTRIBUTES_PREFIXES):
+        return handle_extended_attributes_file(
+            file_key,
+            bucket_name,
+            message_id,
+            created_at_formatted_string,
+            expiry_timestamp,
+        )
+    else:
+        return handle_batch_file(
+            file_key,
+            bucket_name,
+            message_id,
+            created_at_formatted_string,
+            expiry_timestamp,
+        )
+
+
+def get_file_status_for_error(error: Exception) -> str:
+    """Creates a file status based on the type of error that was thrown"""
+    if isinstance(error, VaccineTypePermissionsError):
+        return f"{FileStatus.NOT_PROCESSED} - {FileNotProcessedReason.UNAUTHORISED}"
 
-            upsert_audit_table(
-                message_id,
+    return FileStatus.FAILED
+
+
+def handle_unexpected_bucket_name(bucket_name: str, file_key: str) -> dict:
+    """Handles scenario where Lambda was not invoked by the data-sources bucket. Should not occur due to terraform
+    config and overarching design"""
+    try:
+        if file_key.startswith(EXTENDED_ATTRIBUTES_PREFIXES):
+            extended_attribute_identifier = validate_extended_attributes_file_key(file_key)
+            logger.error(
+                "Unable to process file %s due to unexpected bucket name %s",
                 file_key,
-                created_at_formatted_string,
-                expiry_timestamp,
-                queue_name,
-                FileStatus.PROCESSING,
+                bucket_name,
             )
+            message = f"Failed to process file due to unexpected bucket name {bucket_name}"
+            return {
+                "statusCode": 500,
+                "message": message,
+                "file_key": file_key,
+                "vaccine_supplier_info": extended_attribute_identifier,
+            }
         else:
             vaccine_type, supplier = validate_batch_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,
-                expiry_timestamp,
-                queue_name,
-                FileStatus.QUEUED,
-            )
-            make_and_send_sqs_message(
+            logger.error(
+                "Unable to process file %s due to unexpected bucket name %s",
                 file_key,
-                message_id,
-                permissions,
-                vaccine_type,
-                supplier,
-                created_at_formatted_string,
+                bucket_name,
             )
+            message = f"Failed to process file due to unexpected bucket name {bucket_name}"
 
-            logger.info("Lambda invocation successful for file '%s'", file_key)
-
-            # Return details for logs
             return {
-                "statusCode": 200,
-                "message": "Successfully sent to SQS for further processing",
+                "statusCode": 500,
+                "message": message,
                 "file_key": file_key,
-                "message_id": message_id,
                 "vaccine_type": vaccine_type,
                 "supplier": supplier,
             }
 
+    except Exception as error:
+        logger.error(
+            "Unable to process file due to unexpected bucket name %s and file key %s",
+            bucket_name,
+            file_key,
+        )
+        message = f"Failed to process file due to unexpected bucket name {bucket_name} and file key {file_key}"
+
+        return {
+            "statusCode": 500,
+            "message": message,
+            "file_key": file_key,
+            "vaccine_type": "unknown",
+            "supplier": "unknown",
+            "error": str(error),
+        }
+
+
+def handle_batch_file(file_key, bucket_name, message_id, created_at_formatted_string, expiry_timestamp) -> dict:
+    """
+    Processes a single record for batch file.
+    Returns a dictionary containing information to be included in the logs.
+    """
+    vaccine_type = "unknown"
+    supplier = "unknown"
+    try:
+        vaccine_type, supplier = validate_batch_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,
+            expiry_timestamp,
+            queue_name,
+            FileStatus.QUEUED,
+        )
+        make_and_send_sqs_message(
+            file_key,
+            message_id,
+            permissions,
+            vaccine_type,
+            supplier,
+            created_at_formatted_string,
+        )
+        logger.info("Lambda invocation successful for file '%s'", file_key)
+
+        return {
+            "statusCode": 200,
+            "message": "Batch file successfully processed",
+            "file_key": file_key,
+            "message_id": message_id,
+            "vaccine_type": vaccine_type,
+            "supplier": supplier,
+            "queue_name": queue_name,
+        }
     except (  # pylint: disable=broad-exception-caught
         VaccineTypePermissionsError,
         InvalidFileKeyError,
@@ -133,8 +203,8 @@ def handle_record(record) -> dict:
     ) as error:
         logger.error("Error processing file '%s': %s", file_key, str(error))
 
-        queue_name = f"{supplier}_{vaccine_type}"
         file_status = get_file_status_for_error(error)
+        queue_name = f"{supplier}_{vaccine_type}"
 
         upsert_audit_table(
             message_id,
@@ -165,64 +235,61 @@ def handle_record(record) -> dict:
         }
 
 
-def get_file_status_for_error(error: Exception) -> str:
-    """Creates a file status based on the type of error that was thrown"""
-    if isinstance(error, VaccineTypePermissionsError):
-        return f"{FileStatus.NOT_PROCESSED} - {FileNotProcessedReason.UNAUTHORISED}"
-
-    return FileStatus.FAILED
-
-
-def handle_unexpected_bucket_name(bucket_name: str, file_key: str) -> dict:
-    """Handles scenario where Lambda was not invoked by the data-sources bucket. Should not occur due to terraform
-    config and overarching design"""
+def handle_extended_attributes_file(
+    file_key, bucket_name, message_id, created_at_formatted_string, expiry_timestamp
+) -> dict:
+    """
+    Processes a single record for extended attributes file.
+    Returns a dictionary containing information to be included in the logs.
+    """
     try:
-        if file_key.startswith(EXTENDED_ATTRIBUTES_PREFIXES):
-            validate_extended_attributes_file_key(file_key)
-            logger.error(
-                "Unable to process file %s due to unexpected bucket name %s",
-                file_key,
-                bucket_name,
-            )
-            message = f"Failed to process file due to unexpected bucket name {bucket_name}"
-            return {
-                "statusCode": 500,
-                "message": message,
-                "file_key": file_key,
-                "vaccine_type": "extended_attributes",
-                "supplier": "unknown",
-            }
-        else:
-            vaccine_type, supplier = validate_batch_file_key(file_key)
-            logger.error(
-                "Unable to process file %s due to unexpected bucket name %s",
-                file_key,
-                bucket_name,
-            )
-            message = f"Failed to process file due to unexpected bucket name {bucket_name}"
+        extended_attribute_identifier = validate_extended_attributes_file_key(file_key)
+        move_file_outside_bucket(bucket_name, file_key, DPS_DESTINATION_BUCKET_NAME, f"archive/{file_key}")
+        queue_name = extended_attribute_identifier
 
-            return {
-                "statusCode": 500,
-                "message": message,
-                "file_key": file_key,
-                "vaccine_type": vaccine_type,
-                "supplier": supplier,
-            }
+        upsert_audit_table(
+            message_id,
+            file_key,
+            created_at_formatted_string,
+            expiry_timestamp,
+            queue_name,
+            FileStatus.PROCESSING,
+        )
+        return {
+            "statusCode": 200,
+            "message": "Extended Attributes file successfully processed",
+            "file_key": file_key,
+            "message_id": message_id,
+            "queue_name": queue_name,
+        }
+    except (  # pylint: disable=broad-exception-caught
+        VaccineTypePermissionsError,
+        InvalidFileKeyError,
+        UnhandledAuditTableError,
+        UnhandledSqsError,
+        Exception,
+    ) as error:
+        logger.error("Error processing file '%s': %s", file_key, str(error))
 
-    except Exception as error:
-        logger.error(
-            "Unable to process file due to unexpected bucket name %s and file key %s",
-            bucket_name,
+        file_status = get_file_status_for_error(error)
+        extended_attribute_identifier = validate_extended_attributes_file_key(file_key)
+        queue_name = extended_attribute_identifier
+
+        upsert_audit_table(
+            message_id,
             file_key,
+            created_at_formatted_string,
+            expiry_timestamp,
+            extended_attribute_identifier,
+            file_status,
+            error_details=str(error),
         )
-        message = f"Failed to process file due to unexpected bucket name {bucket_name} and file key {file_key}"
 
         return {
             "statusCode": 500,
-            "message": message,
+            "message": f"Failed to process extended attributes file {file_key} from bucket {bucket_name}",
             "file_key": file_key,
-            "vaccine_type": "unknown",
-            "supplier": "unknown",
+            "message_id": message_id,
             "error": str(error),
         }
 

lambdas/filenameprocessor/tests/test_lambda_handler.py

@@ -28,6 +28,7 @@
     MOCK_BATCH_FILE_CONTENT,
     MOCK_CREATED_AT_FORMATTED_STRING,
     MOCK_EXPIRES_AT,
+    MOCK_EXTENDED_ATTRIBUTES_FILE_CONTENT,
     MockFileDetails,
 )
 
@@ -239,6 +240,67 @@ def test_lambda_handler_non_root_file(self):
         self.assert_no_sqs_message()
         self.assert_no_ack_file(file_details)
 
+    def test_lambda_handler_extended_attributes_success(self):
+        """
+        Tests that for an extended attributes file (prefix starts with 'Vaccination_Extended_Attributes'):
+        * The file is added to the audit table with a status of 'Processing'
+        * The queue_name stored is the extended attribute identifier
+        * The file is moved to the destination bucket under archive/
+        * No SQS message is sent
+        * No ack file is created
+        """
+
+        # Build an extended attributes file.
+        # FileDetails supports this when vaccine_type starts with 'Vaccination_Extended_Attributes'.
+        test_cases = [MockFileDetails.extended_attributes_file]
+
+        # Put file in source bucket
+        s3_client.put_object(
+            Bucket=BucketNames.SOURCE,
+            Key=test_cases[0].file_key,
+            Body=MOCK_EXTENDED_ATTRIBUTES_FILE_CONTENT,
+        )
+
+        # Patch uuid4 (message id), the identifier extraction, and prevent external copy issues by simulating move
+        with (
+            patch("file_name_processor.uuid4", return_value=test_cases[0].message_id),
+            patch(
+                "file_name_processor.validate_extended_attributes_file_key",
+                return_value=test_cases[0].ods_code + "_COVID",
+            ),
+            patch(
+                "file_name_processor.move_file_outside_bucket",
+                side_effect=lambda src_bucket, key, dst_bucket, dst_key: (
+                    s3_client.put_object(
+                        Bucket=BucketNames.DESTINATION,
+                        Key=dst_key,
+                        Body=s3_client.get_object(Bucket=src_bucket, Key=key)["Body"].read(),
+                    ),
+                    s3_client.delete_object(Bucket=src_bucket, Key=key),
+                ),
+            ),
+        ):
+            lambda_handler(self.make_event([self.make_record(test_cases[0].file_key)]), None)
+
+        # Assert audit table entry captured with Processing and queue_name set to the identifier
+        table_items = self.get_audit_table_items()
+        self.assertEqual(len(table_items), 1)
+        item = table_items[0]
+        self.assertEqual(item[AuditTableKeys.MESSAGE_ID]["S"], test_cases[0].message_id)
+        self.assertEqual(item[AuditTableKeys.FILENAME]["S"], test_cases[0].file_key)
+        self.assertEqual(item[AuditTableKeys.QUEUE_NAME]["S"], test_cases[0].ods_code + "_COVID")
+        self.assertEqual(item[AuditTableKeys.TIMESTAMP]["S"], test_cases[0].created_at_formatted_string)
+        self.assertEqual(item[AuditTableKeys.EXPIRES_AT]["N"], str(test_cases[0].expires_at))
+        # File should be moved to destination under archive/
+        dest_key = f"archive/{test_cases[0].file_key}"
+        print(f" destination file is at {s3_client.list_objects(Bucket=BucketNames.DESTINATION)}")
+        retrieved = s3_client.get_object(Bucket=BucketNames.DESTINATION, Key=dest_key)
+        self.assertIsNotNone(retrieved)
+
+        # No SQS and no ack file
+        self.assert_no_sqs_message()
+        self.assert_no_ack_file(test_cases[0])
+
     @patch("elasticache.get_redis_client")
     def test_lambda_invalid_file_key_no_other_files_in_queue(self, mock_get_redis_client):
         """