[PRMT-611] Implement pagination for DynamoDB queries for the statistical reports (#780)

Author: NogaNHS

lambdas/services/base/dynamo_service.py

@@ -4,6 +4,7 @@
 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,
@@ -76,6 +77,34 @@ def query_table_by_index(
             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")
+
+            dynamodb_scan_result = results["Items"]
+
+            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
+
     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.

lambdas/services/statistical_report_service.py

@@ -6,6 +6,7 @@
 import polars as pl
 import polars.selectors as column_select
 from inflection import humanize
+
 from models.report.statistics import (
     ApplicationData,
     LoadedStatisticData,
@@ -22,6 +23,7 @@
 
 logger = LoggingService(__name__)
 
+
 class StatisticalReportService:
     def __init__(self):
         self.dynamo_service = DynamoDBService()
@@ -67,12 +69,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:
-            response = self.dynamo_service.query_all_fields(
+            dynamodb_items_for_date = self.dynamo_service.query_with_pagination(
                 table_name=self.statistic_table,
                 search_key="Date",
                 search_condition=date,
             )
-            dynamodb_items.extend(response["Items"])
+            dynamodb_items.extend(dynamodb_items_for_date)
 
         loaded_data = load_from_dynamodb_items(dynamodb_items)
 
@@ -159,7 +161,9 @@ def summarise_application_data(
                 pl.concat_list("active_user_ids_hashed")
                 .flatten()
                 .unique()
-                .map_elements(lambda col: str(col.sort().to_list()), return_dtype=pl.Utf8)
+                .map_elements(
+                    lambda col: str(col.sort().to_list()), return_dtype=pl.Utf8
+                )
                 .alias("unique_active_user_ids_hashed"),
                 pl.concat_list("active_user_ids_hashed")
                 .flatten()
@@ -179,7 +183,7 @@ def join_dataframes_by_ods_code(
 
         for other_dataframe in data_to_report[1:]:
             joined_dataframe = joined_dataframe.join(
-                other_dataframe, on="ods_code", how='full', coalesce=True
+                other_dataframe, on="ods_code", how="full", coalesce=True
             )
 
         return joined_dataframe

lambdas/tests/unit/helpers/data/statistic/mock_statistic_data.py

@@ -299,8 +299,8 @@
 )
 
 MOCK_DYNAMODB_QUERY_RESPONSE = [
-    {"Items": [item for item in MOCK_DYNAMODB_ITEMS if item["Date"] == "20240510"]},
-    {"Items": [item for item in MOCK_DYNAMODB_ITEMS if item["Date"] == "20240511"]},
+    [item for item in MOCK_DYNAMODB_ITEMS if item["Date"] == "20240510"],
+    [item for item in MOCK_DYNAMODB_ITEMS if item["Date"] == "20240511"],
 ]
 
 EXPECTED_WEEKLY_SUMMARY = pl.DataFrame(

lambdas/tests/unit/services/base/test_dynamo_service.py

@@ -1,9 +1,10 @@
-from typing import Optional
+import copy
 from unittest.mock import call
 
 import pytest
 from boto3.dynamodb.conditions import Attr, Key
 from botocore.exceptions import ClientError
+
 from enums.dynamo_filter import AttributeOperator
 from enums.metadata_field_names import DocumentReferenceMetadataFields
 from services.base.dynamo_service import DynamoDBService
@@ -25,6 +26,7 @@ def mock_service(set_env, mocker):
     mocker.patch("boto3.resource")
     service = DynamoDBService()
     yield service
+    DynamoDBService.instance = None
 
 
 @pytest.fixture
@@ -55,6 +57,18 @@ def mock_filter_expression():
     ).build()
     yield filter_expression
 
+def mock_scan_implementation(
+    **kwargs
+):
+    key = kwargs.get("ExclusiveStartKey")
+    if not key:
+        return copy.deepcopy(MOCK_PAGINATED_RESPONSE_1)
+    elif key.get("ID") == "id_token_for_page_2":
+        return copy.deepcopy(MOCK_PAGINATED_RESPONSE_2)
+    elif key.get("ID") == "id_token_for_page_3":
+        return copy.deepcopy(MOCK_PAGINATED_RESPONSE_3)
+    return None
+
 
 def test_query_with_requested_fields_returns_items_from_dynamo(
     mock_service, mock_table
@@ -492,16 +506,6 @@ def test_scan_whole_table_return_items_in_response(
 def test_scan_whole_table_handles_pagination(
     mock_service, mock_scan_method, mock_filter_expression
 ):
-    def mock_scan_implementation(
-        ExclusiveStartKey: Optional[dict[str, str]] = None, **_kwargs
-    ):
-        if not ExclusiveStartKey:
-            return MOCK_PAGINATED_RESPONSE_1
-        elif ExclusiveStartKey.get("ID") == "id_token_for_page_2":
-            return MOCK_PAGINATED_RESPONSE_2
-        elif ExclusiveStartKey.get("ID") == "id_token_for_page_3":
-            return MOCK_PAGINATED_RESPONSE_3
-
     mock_project_expression = "mock_project_expression"
     mock_scan_method.side_effect = mock_scan_implementation
 
@@ -577,3 +581,32 @@ def test_dynamo_service_singleton_instance(mocker):
     instance_2 = DynamoDBService()
 
     assert instance_1 is instance_2
+
+
+def test_query_with_pagination(mock_service, mock_table):
+    mock_table.return_value.query.side_effect = mock_scan_implementation
+    expected_result = EXPECTED_ITEMS_FOR_PAGINATED_RESULTS
+    search_key_obj = Key("NhsNumber").eq(TEST_NHS_NUMBER)
+
+    expected_calls = [
+        call(
+            KeyConditionExpression=search_key_obj,
+        ),
+        call(
+            KeyConditionExpression=search_key_obj,
+            ExclusiveStartKey={"ID": "id_token_for_page_2"},
+        ),
+        call(
+            KeyConditionExpression=search_key_obj,
+            ExclusiveStartKey={"ID": "id_token_for_page_3"},
+        ),
+    ]
+
+    actual = mock_service.query_with_pagination(
+        table_name=MOCK_TABLE_NAME,
+        search_key="NhsNumber",
+        search_condition=TEST_NHS_NUMBER,
+    )
+    assert expected_result == actual
+    mock_table.assert_called_with(MOCK_TABLE_NAME)
+    mock_table.return_value.query.assert_has_calls(expected_calls)

lambdas/tests/unit/services/test_statistical_report_service.py

@@ -103,7 +103,9 @@ def test_make_weekly_summary(set_env, mocker):
 
 def test_get_statistic_data(mock_dynamodb_service, mock_service):
     mock_service.dates_to_collect = ["20240510", "20240511"]
-    mock_dynamodb_service.query_all_fields.side_effect = MOCK_DYNAMODB_QUERY_RESPONSE
+    mock_dynamodb_service.query_with_pagination.side_effect = (
+        MOCK_DYNAMODB_QUERY_RESPONSE
+    )
 
     actual = mock_service.get_statistic_data()
     expected = ALL_MOCKED_STATISTIC_DATA
@@ -123,7 +125,7 @@ def test_get_statistic_data(mock_dynamodb_service, mock_service):
         ),
     ]
 
-    mock_dynamodb_service.query_all_fields.assert_has_calls(expected_calls)
+    mock_dynamodb_service.query_with_pagination.assert_has_calls(expected_calls)
 
 
 def test_get_statistic_data_raise_error_if_all_data_are_empty(