Refactor FhirController to handle larger search responses and improve error handling for oversized queries

Author: Thomas-Boyle

lambdas/backend/src/controller/fhir_controller.py

@@ -160,7 +160,7 @@ def _get_immunization_by_identifier(self, search_params: dict[str, list[str]], s
         identifier = Identifier.construct(system=identifier_components[0], value=identifier_components[1])
 
         search_bundle = self.fhir_service.get_immunization_by_identifier(identifier, supplier_system, element)
-        prepared_search_bundle = self._prepare_search_bundle(search_bundle)
+        prepared_search_bundle = self._prepare_search_bundle(search_bundle.json())
 
         return create_response(200, prepared_search_bundle)
 
@@ -177,12 +177,7 @@ def _search_immunizations(self, search_params: dict[str, list[str]], supplier_sy
             result.invalid_immunization_targets,
         )
 
-        if self._has_too_many_search_results(search_bundle):
-            raise TooManyResultsError("Search returned too many results. Please narrow down the search")
-
-        prepared_search_bundle = self._prepare_search_bundle(search_bundle)
-
-        return create_response(200, prepared_search_bundle)
+        return self._create_search_response(search_bundle)
 
     def _search_immunizations_by_target_disease(self, search_params: dict[str, list[str]], supplier_system: str) -> dict:
         result = validate_and_retrieve_search_params_by_disease(search_params)
@@ -208,7 +203,15 @@ def _search_immunizations_by_target_disease(self, search_params: dict[str, list[
                 invalid_target_diseases=result.invalid_target_diseases,
             )
 
-        prepared_search_bundle = self._prepare_search_bundle(search_bundle)
+        return self._create_search_response(search_bundle)
+
+    def _create_search_response(self, search_bundle: Bundle) -> dict:
+        search_response_json = search_bundle.json(use_decimal=True)
+
+        if len(search_response_json) > MAX_RESPONSE_SIZE_BYTES:
+            raise TooManyResultsError("Search returned too many results. Please narrow down the search")
+
+        prepared_search_bundle = self._prepare_search_bundle(search_response_json)
         return create_response(200, prepared_search_bundle)
 
     @staticmethod
@@ -220,11 +223,11 @@ def _is_valid_immunization_id(self, immunization_id: str) -> bool:
         return False if not re.match(self._IMMUNIZATION_ID_PATTERN, immunization_id) else True
 
     @staticmethod
-    def _prepare_search_bundle(search_response: Bundle) -> dict:
+    def _prepare_search_bundle(search_response_json: str) -> dict:
         """Workaround for fhir.resources dict() or json() removing the empty "entry" list. Team also specified that
         total should be the final key in the object. Should investigate if this can be resolved with later version of
         the library."""
-        search_response_dict = json.loads(search_response.json())
+        search_response_dict = json.loads(search_response_json)
 
         if "entry" not in search_response_dict:
             search_response_dict["entry"] = []
@@ -244,13 +247,6 @@ def _is_identifier_search(search_params: dict[str, list[str]]) -> bool:
             or IdentifierSearchParameterName.ELEMENTS in search_params
         )
 
-    @staticmethod
-    def _has_too_many_search_results(search_response: Bundle) -> bool:
-        """Checks whether the response is too large - 6MB Lambda limit. Note: this condition should never happen as it
-        would require a very large number of vaccs for a single patient. It is also very rudimentary and all it does is
-        ensure we can raise and return a nice looking error. Consider using pagination as a more robust approach."""
-        return len(search_response.json(use_decimal=True)) > MAX_RESPONSE_SIZE_BYTES
-
     @staticmethod
     def _get_search_params_from_request(
         aws_event: APIGatewayProxyEventV1, is_post_endpoint_req: bool

lambdas/backend/tests/controller/test_fhir_controller_target_disease.py

@@ -287,3 +287,30 @@ def test_search_by_target_disease_with_mixed_valid_and_invalid_returns_200_with_
         self.assertEqual(len(call_args[1]["invalid_target_diseases"]), 1)
         self.assertIn("invalid-no-pipe", call_args[1]["invalid_target_diseases"][0])
         self.assertIn("Invalid format", call_args[1]["invalid_target_diseases"][0])
+
+    @patch("controller.fhir_controller.MAX_RESPONSE_SIZE_BYTES", 5)
+    def test_search_by_target_disease_returns_400_when_response_too_large(self):
+        """it should return the same narrow-the-search error for oversized target-disease searches"""
+        self.mock_redis.hget.side_effect = self._hget_target_disease_codes_and_mmr
+        self.service.search_immunizations.return_value = Bundle.construct(
+            entry=[BundleEntry.construct(resource=Immunization.construct(**{"id": "imms-1"}))],
+            link=[BundleLink.construct(relation="self", url="search-url")],
+            type="searchset",
+            total=1,
+        )
+        lambda_event = {
+            "headers": {"SupplierSystem": "test"},
+            "multiValueQueryStringParameters": {
+                self.patient_identifier_key: [self.patient_identifier_valid_value],
+                self.target_disease_key: [f"{self.snomed_system}|{self.measles_code}"],
+            },
+        }
+
+        response = self.controller.search_immunizations(lambda_event)
+        response_body = json.loads(response["body"])
+
+        self.assertEqual(response["statusCode"], 400)
+        self.assertEqual(
+            response_body["issue"][0]["diagnostics"],
+            "Search returned too many results. Please narrow down the search",
+        )

lambdas/backend/tests/repository/test_fhir_repository.py

@@ -1,7 +1,7 @@
 import time
 import unittest
 import uuid
-from unittest.mock import ANY, MagicMock, Mock, patch
+from unittest.mock import ANY, MagicMock, Mock, call, patch
 
 import botocore.exceptions
 import simplejson as json
@@ -606,6 +606,54 @@ def test_find_immunizations_returns_empty_list_when_not_found(self):
         )
         self.assertEqual(result, [])
 
+    def test_find_immunizations_queries_all_dynamodb_pages(self):
+        """it should continue querying DynamoDB until LastEvaluatedKey is exhausted"""
+        nhs_number = "a-patient-id"
+        imms1 = {"id": 1, "meta": {"versionId": 1}}
+        imms2 = {"id": 2, "meta": {"versionId": 2}}
+        first_page_last_evaluated_key = {"PK": "Immunization#1", "PatientPK": _make_patient_pk(nhs_number)}
+        self.table.query = MagicMock(
+            side_effect=[
+                {
+                    "ResponseMetadata": {"HTTPStatusCode": 200},
+                    "Items": [
+                        {
+                            "Resource": json.dumps(imms1),
+                            "PatientSK": "COVID#some_other_text",
+                            "Version": "1",
+                        }
+                    ],
+                    "LastEvaluatedKey": first_page_last_evaluated_key,
+                },
+                {
+                    "ResponseMetadata": {"HTTPStatusCode": 200},
+                    "Items": [
+                        {
+                            "Resource": json.dumps(imms2),
+                            "PatientSK": "COVID#some_other_text",
+                            "Version": "2",
+                        }
+                    ],
+                },
+            ]
+        )
+
+        results = self.repository.find_immunizations(nhs_number, {"COVID"})
+        expected_query_kwargs = {
+            "IndexName": "PatientGSI",
+            "KeyConditionExpression": Key("PatientPK").eq(_make_patient_pk(nhs_number)),
+            "FilterExpression": Attr("DeletedAt").not_exists() | Attr("DeletedAt").eq("reinstated"),
+        }
+
+        self.assertEqual(
+            self.table.query.call_args_list,
+            [
+                call(**expected_query_kwargs),
+                call(**{**expected_query_kwargs, "ExclusiveStartKey": first_page_last_evaluated_key}),
+            ],
+        )
+        self.assertListEqual(results, [imms1, imms2])
+
     def test_find_immunizations_returns_resources_including_meta(self):
         """it should map Resource list into a list of Immunizations"""
         imms1 = {"id": 1, "meta": {"versionId": 1}}