[PRMP-777] Ruff Fixes (#867)

Author: robg-test

lambdas/handlers/get_fhir_document_reference_handler.py

@@ -1,5 +1,4 @@
 from enums.lambda_error import LambdaError
-from enums.mtls import MtlsCommonNames
 from oauthlib.oauth2 import WebApplicationClient
 from services.base.ssm_service import SSMService
 from services.dynamic_configuration_service import DynamicConfigurationService
@@ -9,14 +8,13 @@
 from utils.audit_logging_setup import LoggingService
 from utils.decorators.ensure_env_var import ensure_environment_variables
 from utils.decorators.handle_lambda_exceptions import handle_lambda_exceptions_fhir
-from utils.lambda_handler_utils import extract_bearer_token
-from utils.lambda_header_utils import validate_common_name_in_mtls
 from utils.decorators.set_audit_arg import set_request_context_for_logging
 from utils.exceptions import AuthorisationException, OidcApiException
 from utils.lambda_exceptions import (
     GetFhirDocumentReferenceException,
     SearchPatientException,
 )
+from utils.lambda_handler_utils import extract_bearer_token
 from utils.lambda_response import ApiGatewayResponse
 
 logger = LoggingService(__name__)

lambdas/services/base/dynamo_service.py

@@ -1,10 +1,9 @@
 import time
-from typing import Iterator, Optional, Sequence
+from typing import Optional, Sequence
 
 import boto3
 from boto3.dynamodb.conditions import Attr, ConditionBase, Key
 from botocore.exceptions import ClientError
-
 from utils.audit_logging_setup import LoggingService
 from utils.dynamo_utils import (
     create_expression_attribute_values,
@@ -120,7 +119,7 @@ def update_item(
 
         if condition_expression:
             update_item_args["ConditionExpression"] = condition_expression
-            
+
         return table.update_item(**update_item_args)
 
     def delete_item(self, table_name: str, key: dict):
@@ -157,10 +156,10 @@ def scan_table(
             raise e
 
     def scan_whole_table(
-            self,
-            table_name: str,
-            project_expression: Optional[str] = None,
-            filter_expression: Optional[str] = None,
+        self,
+        table_name: str,
+        project_expression: Optional[str] = None,
+        filter_expression: Optional[str] = None,
     ) -> list[dict]:
         try:
             table = self.get_table(table_name)
@@ -223,7 +222,7 @@ def batch_get_items(self, table_name: str, key_list: list[str]):
                     )
                     request_items = unprocessed_keys
                     retries += 1
-                    time.sleep((2 ** retries) * 0.1)
+                    time.sleep((2**retries) * 0.1)
                 else:
                     break
 
@@ -243,15 +242,13 @@ def get_item(self, table_name: str, key: dict):
             )
             raise e
 
-    def transact_write_items(
-        self, transact_items: Sequence[dict]
-    ):
+    def transact_write_items(self, transact_items: Sequence[dict]):
         """
         Execute a transactional write operation.
-        
+
         Args:
             transact_items: List of transaction items (Put, Update, Delete, ConditionCheck)
-            
+
         Raises:
             ClientError: If the transaction fails (e.g., TransactionCanceledException)
         """
@@ -263,55 +260,57 @@ def transact_write_items(
             logger.info("Transaction completed successfully")
             return response
         except ClientError as e:
-            error_code = e.response.get('Error', {}).get('Code', '')
-            if error_code == 'TransactionCanceledException':
+            error_code = e.response.get("Error", {}).get("Code", "")
+            if error_code == "TransactionCanceledException":
                 logger.error(f"Transaction cancelled: {str(e)}")
-                cancellation_reasons = e.response.get('CancellationReasons', [])
+                cancellation_reasons = e.response.get("CancellationReasons", [])
                 logger.error(f"Cancellation reasons: {cancellation_reasons}")
             else:
                 logger.error(f"Transaction failed with error: {str(e)}")
             raise e
 
     def build_update_transaction_item(
-        self, 
+        self,
         table_name: str,
-        document_key: dict, 
-        update_fields: dict, 
-        condition_fields: dict
+        document_key: dict,
+        update_fields: dict,
+        condition_fields: dict,
     ) -> dict:
         """
         Build a DynamoDB transaction update item with a conditional expression.
-        
+
         Args:
             table_name: The name of the DynamoDB table
             document_key: The key of the table to update
             update_fields: Dictionary of fields to update (already in DynamoDB format/aliases)
             condition_fields: Dictionary of field names and their expected values for the condition to pass
                              e.g., {"DocStatus": "final", "Version": 1}
-            
+
         Returns:
             A transaction item dict ready for transact_write_items
         """
         field_names = list(update_fields.keys())
         update_expression = create_update_expression(field_names)
         _, expression_attribute_names = create_expressions(field_names)
         expression_attribute_values = create_expression_attribute_values(update_fields)
-        
+
         # Build condition expression for multiple fields
         condition_expressions = []
         condition_attribute_names = {}
         condition_attribute_values = {}
-        
+
         for field_name, field_value in condition_fields.items():
             condition_placeholder = f"#{field_name}_attr"
             condition_value_placeholder = f":{field_name}_condition_val"
-            condition_expressions.append(f"{condition_placeholder} = {condition_value_placeholder}")
+            condition_expressions.append(
+                f"{condition_placeholder} = {condition_value_placeholder}"
+            )
             condition_attribute_names[condition_placeholder] = field_name
             condition_attribute_values[condition_value_placeholder] = field_value
-        
+
         # Join multiple conditions with AND
         condition_expression = " AND ".join(condition_expressions)
-        
+
         return {
             "Update": {
                 "TableName": table_name,
@@ -320,11 +319,11 @@ def build_update_transaction_item(
                 "ConditionExpression": condition_expression,
                 "ExpressionAttributeNames": {
                     **expression_attribute_names,
-                    **condition_attribute_names
+                    **condition_attribute_names,
                 },
                 "ExpressionAttributeValues": {
                     **expression_attribute_values,
-                    **condition_attribute_values
-                }
+                    **condition_attribute_values,
+                },
             }
         }

lambdas/services/put_fhir_document_reference_service.py

@@ -1,34 +1,34 @@
-import os
-
-from utils.audit_logging_setup import LoggingService
+from enums.lambda_error import LambdaError
+from enums.snomed_codes import SnomedCode
+from models.document_reference import DocumentReference
 from models.fhir.R4.fhir_document_reference import (
     DocumentReference as FhirDocumentReference,
 )
-from models.document_reference import DocumentReference
-from services.fhir_document_reference_service_base import FhirDocumentReferenceServiceBase
-from enums.snomed_codes import SnomedCode
-from enums.lambda_error import LambdaError
-from utils.lambda_exceptions import UpdateFhirDocumentReferenceException
+from pydantic import ValidationError
+from services.fhir_document_reference_service_base import (
+    FhirDocumentReferenceServiceBase,
+)
+from utils.audit_logging_setup import LoggingService
 from utils.dynamo_utils import DocTypeTableRouter
+from utils.exceptions import FhirDocumentReferenceException, InvalidNhsNumberException
 from utils.lambda_exceptions import (
     InvalidDocTypeException,
-)
-from pydantic import ValidationError
-from utils.exceptions import (
-    InvalidNhsNumberException,
-    FhirDocumentReferenceException,
+    UpdateFhirDocumentReferenceException,
 )
 
 logger = LoggingService(__name__)
 
+
 class PutFhirDocumentReferenceService(FhirDocumentReferenceServiceBase):
     def __init__(self):
         super().__init__()
         self.doc_router = DocTypeTableRouter()
 
     def process_fhir_document_reference(self, fhir_document: str) -> str:
         try:
-            validated_fhir_doc = FhirDocumentReference.model_validate_json(fhir_document)
+            validated_fhir_doc = FhirDocumentReference.model_validate_json(
+                fhir_document
+            )
         except ValidationError as e:
             logger.error(f"FHIR document validation error: {str(e)}")
             raise UpdateFhirDocumentReferenceException(400, LambdaError.DocRefNoParse)
@@ -38,38 +38,41 @@ def process_fhir_document_reference(self, fhir_document: str) -> str:
             doc_type = self._determine_document_type(validated_fhir_doc)
         except FhirDocumentReferenceException:
             logger.error("Could not determine document type")
-            raise UpdateFhirDocumentReferenceException(400, LambdaError.DocRefInvalidType)
+            raise UpdateFhirDocumentReferenceException(
+                400, LambdaError.DocRefInvalidType
+            )
 
         # Determine which DynamoDB table to use based on the document type
         dynamo_table = self._get_dynamo_table_for_doc_type(doc_type)
 
         # Get the current document from the database
         current_doc = self._get_current_doc(validated_fhir_doc, dynamo_table)
-        
+
         # Check that the NHS number in the request matches the one in the stored document
         request_nhs_number = self._validate_nhs_number(validated_fhir_doc, current_doc)
-        
+
         # Check that the version number matches the stored document
         self._validate_version_number(validated_fhir_doc, current_doc)
-        
+
         # Create a new document reference with the new version number
         new_doc_reference = self._create_new_document_reference(
-            validated_fhir_doc,
-            current_doc,
-            request_nhs_number,
-            doc_type
+            validated_fhir_doc, current_doc, request_nhs_number, doc_type
         )
 
         # Handle binary content if present, otherwise create a pre-signed URL
-        presigned_url = self._handle_document_save(new_doc_reference, validated_fhir_doc, dynamo_table)
-        
+        presigned_url = self._handle_document_save(
+            new_doc_reference, validated_fhir_doc, dynamo_table
+        )
+
         try:
             return self._create_fhir_response(new_doc_reference, presigned_url)
         except (ValidationError, InvalidNhsNumberException) as e:
             logger.error(f"FHIR document validation error: {str(e)}")
             raise UpdateFhirDocumentReferenceException(400, LambdaError.DocRefNoParse)
-        
-    def _get_current_doc(self, fhir_doc: FhirDocumentReference, dynamo_table: str) -> DocumentReference:
+
+    def _get_current_doc(
+        self, fhir_doc: FhirDocumentReference, dynamo_table: str
+    ) -> DocumentReference:
         try:
             current_doc = self._get_document_reference(fhir_doc.id, dynamo_table)
         except FhirDocumentReferenceException:
@@ -83,10 +86,12 @@ def _get_current_doc(self, fhir_doc: FhirDocumentReference, dynamo_table: str) -
             raise UpdateFhirDocumentReferenceException(
                 400, LambdaError.UpdateDocNotLatestVersion
             )
-        
+
         return current_doc
-    
-    def _validate_nhs_number(self, fhir_doc: FhirDocumentReference, current_doc: DocumentReference):
+
+    def _validate_nhs_number(
+        self, fhir_doc: FhirDocumentReference, current_doc: DocumentReference
+    ):
         try:
             request_nhs_number = fhir_doc.extract_nhs_number_from_fhir()
         except FhirDocumentReferenceException:
@@ -98,30 +103,38 @@ def _validate_nhs_number(self, fhir_doc: FhirDocumentReference, current_doc: Doc
             raise UpdateFhirDocumentReferenceException(
                 400, LambdaError.UpdateDocNHSNumberMismatch
             )
-        
+
         return request_nhs_number
-    
-    def _validate_version_number(self, fhir_doc: FhirDocumentReference, current_doc: DocumentReference):
+
+    def _validate_version_number(
+        self, fhir_doc: FhirDocumentReference, current_doc: DocumentReference
+    ):
         if fhir_doc.meta is None:
             logger.error("Missing version number")
-            raise UpdateFhirDocumentReferenceException(400, LambdaError.DocumentReferenceMissingParameters)
-        
+            raise UpdateFhirDocumentReferenceException(
+                400, LambdaError.DocumentReferenceMissingParameters
+            )
+
         if current_doc.version != fhir_doc.meta.versionId:
             logger.error("Version does not match current version.")
-            raise UpdateFhirDocumentReferenceException(400, LambdaError.UpdateDocVersionMismatch)
-    
+            raise UpdateFhirDocumentReferenceException(
+                400, LambdaError.UpdateDocVersionMismatch
+            )
+
     def _create_new_document_reference(
-            self, 
-            fhir_doc: FhirDocumentReference,
-            current_doc: DocumentReference,
-            nhs_number: str,
-            doc_type: SnomedCode,
-            ) -> DocumentReference:
+        self,
+        fhir_doc: FhirDocumentReference,
+        current_doc: DocumentReference,
+        nhs_number: str,
+        doc_type: SnomedCode,
+    ) -> DocumentReference:
         try:
             patient_details = self._check_nhs_number_with_pds(nhs_number)
         except FhirDocumentReferenceException:
-            raise UpdateFhirDocumentReferenceException(400, LambdaError.DocRefPatientSearchInvalid)
-        
+            raise UpdateFhirDocumentReferenceException(
+                400, LambdaError.DocRefPatientSearchInvalid
+            )
+
         new_doc_version = int(current_doc.version) + 1
 
         # Create a document reference model
@@ -135,13 +148,13 @@ def _create_new_document_reference(
         )
 
         return document_reference
-    
+
     def _handle_document_save(
-            self,
-            document_reference: DocumentReference,
-            fhir_doc: FhirDocumentReference,
-            dynamo_table: str,
-            ) -> str:
+        self,
+        document_reference: DocumentReference,
+        fhir_doc: FhirDocumentReference,
+        dynamo_table: str,
+    ) -> str:
         binary_content = fhir_doc.content[0].attachment.data
 
         presigned_url = None
@@ -150,21 +163,27 @@ def _handle_document_save(
             try:
                 self._store_binary_in_s3(document_reference, binary_content)
             except FhirDocumentReferenceException:
-                raise UpdateFhirDocumentReferenceException(500, LambdaError.DocRefNoParse)
+                raise UpdateFhirDocumentReferenceException(
+                    500, LambdaError.DocRefNoParse
+                )
         else:
             # Create a pre-signed URL for uploading
             try:
                 presigned_url = self._create_s3_presigned_url(document_reference)
             except FhirDocumentReferenceException:
-                raise UpdateFhirDocumentReferenceException(500, LambdaError.InternalServerError)
+                raise UpdateFhirDocumentReferenceException(
+                    500, LambdaError.InternalServerError
+                )
         try:
             # Save document reference to DynamoDB
             self._save_document_reference_to_dynamo(dynamo_table, document_reference)
         except FhirDocumentReferenceException:
-            raise UpdateFhirDocumentReferenceException(500, LambdaError.DocRefUploadInternalError)
-        
+            raise UpdateFhirDocumentReferenceException(
+                500, LambdaError.DocRefUploadInternalError
+            )
+
         return presigned_url
-    
+
     def _get_dynamo_table_for_doc_type(self, doc_type: SnomedCode) -> str:
         try:
             return self.doc_router.resolve(doc_type)