[NDR-310] Fix bug introduced by the S3 file key index. (#874)

jameslinnell · web-flow · commit 60b18acbcc6f · 2025-11-10T12:18:07.000Z
diff --git a/lambdas/models/document_reference.py b/lambdas/models/document_reference.py
@@ -46,10 +46,7 @@ class Subject(BaseModel):
 
 
 class UpdateBody(BaseModel):
-    model_config = ConfigDict(
-        validate_by_alias=True,
-        alias_generator=to_camel
-    )
+    model_config = ConfigDict(validate_by_alias=True, alias_generator=to_camel)
     resource_type: str
     subject: Subject
     content: list[SingleAttachment]
@@ -61,7 +58,7 @@ class GenericEventModel(BaseModel):
         validate_by_alias=True,
         validate_by_name=True,
         populate_by_name=True,
-        alias_generator=to_camel
+        alias_generator=to_camel,
     )
     http_method: str
     query_string_parameters: dict
@@ -190,7 +187,7 @@ def _build_s3_key(data: dict) -> str:
         s3_key = "/".join(key_parts)
 
         return s3_key
-    
+
     @staticmethod
     def _build_final_s3_key(data: dict) -> str:
         """Build the S3 key from document data."""
@@ -244,4 +241,3 @@ def infer_doc_status(self) -> str | None:
         if self.uploading:
             return "preliminary"
         return None
-
diff --git a/lambdas/services/base/dynamo_service.py b/lambdas/services/base/dynamo_service.py
@@ -41,7 +41,7 @@ def query_table(
         table_name,
         search_key,
         search_condition: str,
-        index_name: str = None,
+        index_name: str | None = None,
         requested_fields: list[str] = None,
         query_filter: Attr | ConditionBase = None,
     ) -> list[dict]:
diff --git a/lambdas/services/document_service.py b/lambdas/services/document_service.py
@@ -55,7 +55,7 @@ def fetch_documents_from_table(
         table: str,
         search_condition: str,
         search_key: str,
-        index_name: str = None,
+        index_name: str | None = None,
         query_filter: Attr | ConditionBase = None,
     ) -> list[DocumentReference]:
         documents = []
diff --git a/lambdas/services/upload_document_reference_service.py b/lambdas/services/upload_document_reference_service.py
@@ -33,6 +33,7 @@ def __init__(self):
         self.staging_s3_bucket_name = os.environ["STAGING_STORE_BUCKET_NAME"]
         self.table_name = os.environ["LLOYD_GEORGE_DYNAMODB_NAME"]
         self.destination_bucket_name = os.environ["LLOYD_GEORGE_BUCKET_NAME"]
+        self.doc_type = SnomedCodes.LLOYD_GEORGE.value
         self.document_service = DocumentService()
         self.dynamo_service = DynamoDBService()
         self.virus_scan_service = get_virus_scan_service()
@@ -67,7 +68,7 @@ def handle_upload_document_reference_request(
             logger.error(f"Unexpected error processing document reference: {str(e)}")
             logger.error(f"Failed to process document reference: {object_key}")
             return
-        
+
     def _get_infrastructure_for_document_key(self, object_parts: list[str]) -> None:
         doc_type = None
         if object_parts[0] != "fhir_upload" or not (
@@ -76,6 +77,7 @@ def _get_infrastructure_for_document_key(self, object_parts: list[str]) -> None:
             return
 
         try:
+            self.doc_type = doc_type
             self.table_name = self.table_router.resolve(doc_type)
             self.destination_bucket_name = self.bucket_router.resolve(doc_type)
         except KeyError:
@@ -126,7 +128,9 @@ def _process_preliminary_document_reference(
     ):
         """Process the preliminary (uploading) document reference with virus scanning and file operations"""
         try:
-            virus_scan_result = self._perform_virus_scan(preliminary_document_reference, object_key)
+            virus_scan_result = self._perform_virus_scan(
+                preliminary_document_reference, object_key
+            )
             preliminary_document_reference.virus_scanner_result = virus_scan_result
 
             if virus_scan_result == VirusScanResult.CLEAN:
@@ -143,22 +147,38 @@ def _process_preliminary_document_reference(
             preliminary_document_reference.uploaded = True
             preliminary_document_reference.uploading = False
 
-            updated_doc_status = None
-            if virus_scan_result != VirusScanResult.CLEAN:
-                updated_doc_status = "cancelled"
+            if self.doc_type.code != SnomedCodes.PATIENT_DATA.value.code:
+                updated_doc_status = None
+                if virus_scan_result != VirusScanResult.CLEAN:
+                    updated_doc_status = "cancelled"
 
-                preliminary_document_reference.doc_status = updated_doc_status
-                self._update_dynamo_table(preliminary_document_reference)
-            else:
-                updated_doc_status = "final"
-                preliminary_document_reference.doc_status = updated_doc_status
+                    preliminary_document_reference.doc_status = updated_doc_status
+                    self._update_dynamo_table(preliminary_document_reference)
+                else:
+                    updated_doc_status = "final"
+                    preliminary_document_reference.doc_status = updated_doc_status
 
-                self._finalize_and_supersede_with_transaction(
-                    preliminary_document_reference
-                )
+                    self._finalize_and_supersede_with_transaction(
+                        preliminary_document_reference
+                    )
 
-                # Update NRL Pointer
-                # TODO: PRMP-390
+                    # Update NRL Pointer
+                    # TODO: PRMP-390
+                    #
+            else:
+                try:
+                    preliminary_document_reference.doc_status = (
+                        "cancelled"
+                        if virus_scan_result != VirusScanResult.CLEAN
+                        else "final"
+                    )
+
+                    self._update_dynamo_table(preliminary_document_reference)
+                except Exception as e:
+                    logger.error(
+                        f"Error processing document reference {preliminary_document_reference.id}: {str(e)}"
+                    )
+                    raise
 
         except TransactionConflictException as e:
             logger.error(
@@ -300,7 +320,9 @@ def _finalize_and_supersede_with_transaction(self, new_document: DocumentReferen
             new_document.uploading = False
             new_document.file_size = None
             self._update_dynamo_table(new_document)
-            self.delete_file_from_bucket(new_document.file_location, new_document.s3_version_id)
+            self.delete_file_from_bucket(
+                new_document.file_location, new_document.s3_version_id
+            )
             raise
         except Exception as e:
             logger.error(
@@ -350,22 +372,23 @@ def copy_files_from_staging_bucket(
         """Copy files from staging bucket to destination bucket"""
         try:
             logger.info("Copying files from staging bucket")
-
-            dest_file_key = document_reference.s3_file_key
+            dest_file_key = f"{document_reference.nhs_number}/{document_reference.id}"
+            if self.doc_type.code != SnomedCodes.PATIENT_DATA.value.code:
+                dest_file_key = document_reference.s3_file_key
 
             copy_result = self.s3_service.copy_across_bucket(
                 source_bucket=self.staging_s3_bucket_name,
                 source_file_key=source_file_key,
                 dest_bucket=self.destination_bucket_name,
                 dest_file_key=dest_file_key,
             )
-
+            if self.doc_type.code == SnomedCodes.PATIENT_DATA.value.code:
+                document_reference.s3_file_key = dest_file_key
             document_reference.s3_bucket_name = self.destination_bucket_name
             document_reference.file_location = document_reference._build_s3_location(
                 self.destination_bucket_name, dest_file_key
             )
             document_reference.s3_version_id = copy_result.get("VersionId")
-
             return copy_result
 
         except ClientError as e:
@@ -386,7 +409,9 @@ def delete_file_from_staging_bucket(self, source_file_key: str):
     def delete_file_from_bucket(self, file_location: str, version_id: str):
         """Delete file from bucket"""
         try:
-            s3_bucket_name, source_file_key = DocumentReference._parse_s3_location(file_location)
+            s3_bucket_name, source_file_key = DocumentReference._parse_s3_location(
+                file_location
+            )
             logger.info(
                 f"Deleting file from bucket: {s3_bucket_name}/{source_file_key}"
             )
@@ -412,6 +437,8 @@ def _update_dynamo_table(
                 "uploaded",
                 "uploading",
             }
+            if self.doc_type.code == SnomedCodes.PATIENT_DATA.value.code:
+                update_fields.add("s3_file_key")
 
             self.document_service.update_document(
                 table_name=self.table_name,
diff --git a/lambdas/tests/unit/services/test_update_document_reference_service.py b/lambdas/tests/unit/services/test_update_document_reference_service.py
diff --git a/lambdas/tests/unit/services/test_upload_document_reference_service.py b/lambdas/tests/unit/services/test_upload_document_reference_service.py
