[PRMP-175] Implementing consistent pagination when querying DynamoDB (#779)

Co-authored-by: Robert Gaskin <[email protected]>

Author: NogaNHS

lambdas/services/base/dynamo_service.py

@@ -36,16 +36,15 @@ def get_table(self, table_name):
             logger.error(str(e), {"Result": "Unable to connect to DB"})
             raise e
 
-    def query_table_by_index(
-            self,
-            table_name,
-            index_name,
-            search_key,
-            search_condition: str,
-            requested_fields: list[str] = None,
-            query_filter: Attr | ConditionBase = None,
-            exclusive_start_key: dict = None,
-    ):
+    def query_table(
+        self,
+        table_name,
+        search_key,
+        search_condition: str,
+        index_name: str = None,
+        requested_fields: list[str] = None,
+        query_filter: Attr | ConditionBase = None,
+    ) -> list[dict]:
         try:
             table = self.get_table(table_name)
 
@@ -62,61 +61,21 @@ def query_table_by_index(
 
             if query_filter:
                 query_params["FilterExpression"] = query_filter
-            if exclusive_start_key:
-                query_params["ExclusiveStartKey"] = exclusive_start_key
-
-            results = table.query(**query_params)
-
-            if results is None or "Items" not in results:
-                logger.error(f"Unusable results in DynamoDB: {results!r}")
-                raise DynamoServiceException("Unrecognised response from DynamoDB")
-
-            return results
-        except ClientError as e:
-            logger.error(str(e), {"Result": f"Unable to query table: {table_name}"})
-            raise e
-
-    def query_with_pagination(
-            self, table_name: str, search_key: str, search_condition: str
-    ):
-
-        try:
-            table = self.get_table(table_name)
-            results = table.query(
-                KeyConditionExpression=Key(search_key).eq(search_condition)
-            )
-            if results is None or "Items" not in results:
-                logger.error(f"Unusable results in DynamoDB: {results!r}")
-                raise DynamoServiceException("Unrecognised response from DynamoDB")
+            items = []
+            while True:
+                results = table.query(**query_params)
 
-            dynamodb_scan_result = results["Items"]
+                if results is None or "Items" not in results:
+                    logger.error(f"Unusable results in DynamoDB: {results!r}")
+                    raise DynamoServiceException("Unrecognised response from DynamoDB")
 
-            while "LastEvaluatedKey" in results:
-                start_key_for_next_page = results["LastEvaluatedKey"]
-                results = table.query(
-                    KeyConditionExpression=Key(search_key).eq(search_condition),
-                    ExclusiveStartKey=start_key_for_next_page,
-                )
-                dynamodb_scan_result.extend(results["Items"])
-            return dynamodb_scan_result
-
-        except ClientError as e:
-            logger.error(str(e), {"Result": f"Unable to query table: {table_name}"})
-            raise e
+                items += results["Items"]
 
-    def query_all_fields(self, table_name: str, search_key: str, search_condition: str):
-        """
-        Allow querying dynamodb table without explicitly defining the fields to retrieve.
-        """
-        try:
-            table = self.get_table(table_name)
-            results = table.query(
-                KeyConditionExpression=Key(search_key).eq(search_condition)
-            )
-            if results is None or "Items" not in results:
-                logger.error(f"Unusable results in DynamoDB: {results!r}")
-                raise DynamoServiceException("Unrecognised response from DynamoDB")
-            return results
+                if "LastEvaluatedKey" in results:
+                    query_params["ExclusiveStartKey"] = results["LastEvaluatedKey"]
+                else:
+                    break
+            return items
         except ClientError as e:
             logger.error(str(e), {"Result": f"Unable to query table: {table_name}"})
             raise e

lambdas/services/document_manifest_job_service.py

@@ -170,7 +170,7 @@ def query_zip_trace(self, job_id: str, nhs_number: str) -> DocumentManifestZipTr
             attr_operator=AttributeOperator.EQUAL,
             filter_value=nhs_number,
         ).build()
-        response = self.dynamo_service.query_table_by_index(
+        response = self.dynamo_service.query_table(
             table_name=self.zip_trace_table,
             index_name="JobIdIndex",
             search_key="JobId",
@@ -180,7 +180,7 @@ def query_zip_trace(self, job_id: str, nhs_number: str) -> DocumentManifestZipTr
         )
 
         try:
-            zip_trace = DocumentManifestZipTrace.model_validate(response["Items"][0])
+            zip_trace = DocumentManifestZipTrace.model_validate(response[0])
             return zip_trace
         except (KeyError, IndexError, ValidationError) as e:
             logger.error(

lambdas/services/document_service.py

@@ -59,30 +59,22 @@ def fetch_documents_from_table(
         query_filter: Attr | ConditionBase = None,
     ) -> list[DocumentReference]:
         documents = []
-        exclusive_start_key = None
-
-        while True:
-            response = self.dynamo_service.query_table_by_index(
-                table_name=table,
-                index_name=index_name,
-                search_key=search_key,
-                search_condition=search_condition,
-                query_filter=query_filter,
-                exclusive_start_key=exclusive_start_key,
-            )
 
-            for item in response["Items"]:
-                try:
-                    document = DocumentReference.model_validate(item)
-                    documents.append(document)
-                except ValidationError as e:
-                    logger.error(f"Validation error on document: {item}")
-                    logger.error(f"{e}")
-                    continue
-            if "LastEvaluatedKey" in response:
-                exclusive_start_key = response["LastEvaluatedKey"]
-            else:
-                break
+        response = self.dynamo_service.query_table(
+            table_name=table,
+            index_name=index_name,
+            search_key=search_key,
+            search_condition=search_condition,
+            query_filter=query_filter,
+        )
+        for item in response:
+            try:
+                document = DocumentReference.model_validate(item)
+                documents.append(document)
+            except ValidationError as e:
+                logger.error(f"Validation error on document: {item}")
+                logger.error(f"{e}")
+                continue
         return documents
 
     def get_nhs_numbers_based_on_ods_code(self, ods_code: str) -> list[str]:

lambdas/services/im_alerting_service.py

@@ -235,7 +235,7 @@ def get_alarm_history(self, alarm_name: str):
         logger.info(f"Checking if {alarm_name} already exists on alarm table")
 
         try:
-            results = self.dynamo_service.query_all_fields(
+            results = self.dynamo_service.query_table(
                 table_name=self.table_name,
                 search_key="AlarmNameMetric",
                 search_condition=alarm_name,

lambdas/services/lloyd_george_stitch_job_service.py

@@ -115,9 +115,10 @@ def process_stitch_trace_response(self, stitch_trace: StitchTrace):
                     totalFileSizeInBytes=stitch_trace.total_file_size_in_bytes,
                 )
 
-    def validate_stitch_trace(self, response: dict) -> list[StitchTrace] | None:
+    def validate_stitch_trace(
+        self, stitch_trace_dynamo_response: list
+    ) -> list[StitchTrace] | None:
         try:
-            stitch_trace_dynamo_response = response.get("Items", [])
             if not stitch_trace_dynamo_response:
                 return None
             return [
@@ -149,7 +150,7 @@ def query_stitch_trace_with_nhs_number(self, nhs_number: str) -> list[StitchTrac
             attr_operator=AttributeOperator.EQUAL,
             filter_value=False,
         ).build()
-        response = self.dynamo_service.query_table_by_index(
+        response = self.dynamo_service.query_table(
             table_name=self.stitch_trace_table,
             index_name="NhsNumberIndex",
             search_key="NhsNumber",

lambdas/services/login_service.py

@@ -140,11 +140,11 @@ def generate_session(self, state, auth_code) -> dict:
     def have_matching_state_value_in_record(self, state: str) -> bool:
         state_table_name = os.environ["AUTH_STATE_TABLE_NAME"]
 
-        query_response = self.db_service.query_all_fields(
+        query_response = self.db_service.query_table(
             table_name=state_table_name, search_key="State", search_condition=state
         )
 
-        state_match = "Count" in query_response and query_response["Count"] == 1
+        state_match = len(query_response) == 1
 
         if state_match:
             try:

lambdas/services/manage_user_session_access.py

@@ -22,19 +22,18 @@ def find_login_session(self, ndr_session_id: str):
         logger.info(
             f"Retrieving session for session ID ending in: {redact_id_to_last_4_chars(ndr_session_id)}"
         )
-        query_response = self.db_service.query_all_fields(
+        query_response = self.db_service.query_table(
             table_name=self.session_table_name,
             search_key="NDRSessionId",
             search_condition=ndr_session_id,
         )
 
-        try:
-            current_session = query_response["Items"][0]
+        if query_response and len(query_response) > 0:
+            current_session = query_response[0]
             return current_session
-        except (KeyError, IndexError):
-            raise AuthorisationException(
-                f"Unable to find session for session ID ending in: {redact_id_to_last_4_chars(ndr_session_id)}"
-            )
+        raise AuthorisationException(
+            f"Unable to find session for session ID ending in: {redact_id_to_last_4_chars(ndr_session_id)}"
+        )
 
     def update_auth_session_with_permitted_search(
         self,

lambdas/services/ods_report_service.py

@@ -101,25 +101,14 @@ def query_table_by_index(self, ods_code: str):
             ]
         results = []
         for ods_code in ods_codes:
-            response = self.dynamo_service.query_table_by_index(
+            response = self.dynamo_service.query_table(
                 table_name=self.table_name,
                 index_name="OdsCodeIndex",
                 search_key=DocumentReferenceMetadataFields.CURRENT_GP_ODS.value,
                 search_condition=ods_code,
                 query_filter=NotDeleted,
             )
-            results += response["Items"]
-
-            while "LastEvaluatedKey" in response:
-                response = self.dynamo_service.query_table_by_index(
-                    table_name=self.table_name,
-                    index_name="OdsCodeIndex",
-                    exclusive_start_key=response["LastEvaluatedKey"],
-                    search_key=DocumentReferenceMetadataFields.CURRENT_GP_ODS.value,
-                    search_condition=ods_code,
-                    query_filter=NotDeleted,
-                )
-                results += response["Items"]
+            results += response
 
         if not results:
             logger.info("No records found for ODS code {}".format(ods_code))

lambdas/services/statistical_report_service.py

@@ -68,12 +68,12 @@ def get_statistic_data(self) -> LoadedStatisticData:
         logger.info(f"The period to report: {self.dates_to_collect}")
         dynamodb_items = []
         for date in self.dates_to_collect:
-            dynamodb_items_for_date = self.dynamo_service.query_with_pagination(
+            response = self.dynamo_service.query_table(
                 table_name=self.statistic_table,
                 search_key="Date",
                 search_condition=date,
             )
-            dynamodb_items.extend(dynamodb_items_for_date)
+            dynamodb_items.extend(response)
 
         loaded_data = load_from_dynamodb_items(dynamodb_items)