[PRMP-1079] Dynamodb query limit doesn't work as expected.

lambdas/handlers/generate_document_manifest_handler.py

@@ -1,4 +1,3 @@
-from boto3.dynamodb.types import TypeDeserializer
 from enums.lambda_error import LambdaError
 from enums.logging_app_interaction import LoggingAppInteraction
 from models.zip_trace import DocumentManifestZipTrace
@@ -9,6 +8,7 @@
 from utils.decorators.handle_lambda_exceptions import handle_lambda_exceptions
 from utils.decorators.override_error_check import override_error_check
 from utils.decorators.set_audit_arg import set_request_context_for_logging
+from utils.dynamo_utils import deserialize_dynamodb_object
 from utils.lambda_exceptions import GenerateManifestZipException
 from utils.lambda_response import ApiGatewayResponse
 from utils.request_context import request_context
@@ -71,9 +71,4 @@ def manifest_zip_handler(zip_trace_item):
 
 
 def prepare_zip_trace_data(new_zip_trace: dict) -> dict:
-    deserialize = TypeDeserializer().deserialize
-    parsed_dynamodb_items = {
-        key: deserialize(dynamodb_value)
-        for key, dynamodb_value in new_zip_trace.items()
-    }
-    return parsed_dynamodb_items
+    return deserialize_dynamodb_object(new_zip_trace)

lambdas/services/base/dynamo_service.py

@@ -11,6 +11,8 @@
     create_expression_attribute_values,
     create_expressions,
     create_update_expression,
+    deserialize_dynamodb_object,
+    serialize_dict_to_dynamodb_object,
 )
 from utils.exceptions import DynamoServiceException
 
@@ -29,6 +31,7 @@ def __new__(cls):
     def __init__(self):
         if not self.initialised:
             self.dynamodb = boto3.resource("dynamodb", region_name="eu-west-2")
+            self.client = boto3.client("dynamodb")
             self.initialised = True
 
     def get_table(self, table_name: str):
@@ -429,3 +432,57 @@ def build_update_transaction_item(
                 },
             }
         }
+
+    def query_table_with_paginator(
+        self,
+        table_name: str,
+        index_name: str,
+        key: str,
+        condition: str,
+        filter_expression: str | None = None,
+        expression_attribute_names: str | None = None,
+        expression_attribute_values: dict | None = None,
+        limit: int = 20,
+        page_size: int = 1,
+        start_key: str | None = None,
+    ) -> dict:
+
+        try:
+            query_params = {
+                "TableName": table_name,
+                "IndexName": index_name,
+                "KeyConditionExpression": f"{key}=:i",
+                "PaginationConfig": {
+                    "MaxItems": limit,
+                    "PageSize": page_size,
+                    "StartingToken": start_key,
+                },
+            }
+
+            if expression_attribute_values is None:
+                expression_attribute_values = {}
+
+            expression_attribute_values[":i"] = condition
+
+            if filter_expression:
+                query_params["FilterExpression"] = filter_expression
+
+            if expression_attribute_names:
+                query_params["ExpressionAttributeNames"] = expression_attribute_names
+
+            if expression_attribute_values:
+                query_params["ExpressionAttributeValues"] = (
+                    serialize_dict_to_dynamodb_object(expression_attribute_values)
+                )
+
+            paginator = self.client.get_paginator("query")
+            response = paginator.paginate(**query_params).build_full_result()
+
+            response["Items"] = [
+                deserialize_dynamodb_object(item) for item in response["Items"]
+            ]
+            return response
+
+        except Exception as e:
+            logger.error("Failed to query DynamoDB")
+            raise e

lambdas/services/document_service.py

@@ -298,3 +298,43 @@ def create_dynamo_entry(
         except (ValidationError, ClientError) as e:
             logger.error(e)
             raise e
+
+    def query_table_with_paginator(
+        self,
+        index_name: str,
+        search_key: str,
+        search_condition: str,
+        table_name: str | None = None,
+        filter_expression: str | None = None,
+        expression_attribute_names: dict | None = None,
+        expression_attribute_values: dict | None = None,
+        limit: int | None = None,
+        page_size: int = 1,
+        start_key: str | None = None,
+        model_class: BaseModel | None = None,
+    ) -> tuple[list[BaseModel], str | None]:
+
+        try:
+            table_name = table_name or self.table_name
+            model_class = model_class or self.model_class
+
+            response = self.dynamo_service.query_table_with_paginator(
+                table_name=table_name,
+                index_name=index_name,
+                key=search_key,
+                condition=search_condition,
+                filter_expression=filter_expression,
+                expression_attribute_names=expression_attribute_names,
+                expression_attribute_values=expression_attribute_values,
+                limit=limit,
+                page_size=page_size,
+                start_key=start_key
+            )
+
+            references=[model_class.model_validate(item) for item in response["Items"]]
+
+            return references, response.get("NextToken")
+
+        except (ValidationError, ClientError) as e:
+            logger.error(e)
+            raise e

lambdas/services/document_upload_review_service.py

@@ -12,7 +12,10 @@
 from utils.audit_logging_setup import LoggingService
 from utils.aws_transient_error_check import is_transient_error
 from utils.dynamo_query_filter_builder import DynamoQueryFilterBuilder
-from utils.dynamo_utils import build_transaction_item
+from utils.dynamo_utils import (
+    build_mixed_condition_expression,
+    build_transaction_item,
+)
 from utils.exceptions import DocumentReviewException
 
 logger = LoggingService(__name__)
@@ -40,40 +43,32 @@ def model_class(self) -> type:
     def s3_bucket(self) -> str:
         return self._s3_bucket
 
-    def query_docs_pending_review_by_custodian_with_limit(
+    def query_docs_pending_review_with_paginator(
         self,
         ods_code: str,
         limit: int = DEFAULT_QUERY_LIMIT,
-        start_key: dict | None = None,
+        start_key: str | None = None,
         nhs_number: str | None = None,
         uploader: str | None = None,
-    ) -> tuple[list[DocumentUploadReviewReference], dict | None]:
+    ) -> tuple[list[DocumentUploadReviewReference], str | None]:
+
         logger.info(f"Getting review document references for custodian: {ods_code}")
 
-        filter_expression = self.build_review_dynamo_filter(
-            nhs_number=nhs_number, uploader=uploader
+        filter_expression, condition_attribute_names, condition_attribute_values = (
+            self.build_paginator_query_filter(nhs_number=nhs_number, uploader=uploader)
+        )
+        references, last_evaluated_key = self.query_table_with_paginator(
+            index_name="CustodianIndex",
+            search_key="Custodian",
+            search_condition=ods_code,
+            filter_expression=filter_expression,
+            expression_attribute_names=condition_attribute_names,
+            expression_attribute_values=condition_attribute_values,
+            limit=limit,
+            start_key=start_key,
         )
 
-        try:
-            response = self.dynamo_service.query_table_single(
-                table_name=self.table_name,
-                search_key="Custodian",
-                search_condition=ods_code,
-                index_name="CustodianIndex",
-                limit=limit,
-                start_key=start_key,
-                query_filter=filter_expression,
-            )
-
-            references = self._validate_review_references(response["Items"])
-
-            last_evaluated_key = response.get("LastEvaluatedKey", None)
-
-            return references, last_evaluated_key
-
-        except ClientError as e:
-            logger.error(e)
-            raise DocumentReviewException("Error querying document review references")
+        return references, last_evaluated_key
 
     def _validate_review_references(
         self, items: list[dict]
@@ -88,6 +83,36 @@ def _validate_review_references(
             logger.error(e)
             raise DocumentReviewException("Error validating document review references")
 
+    def build_paginator_query_filter(
+        self, nhs_number: str | None = None, uploader: str | None = None
+    ):
+        conditions = [
+            {
+                "field": "ReviewStatus",
+                "operator": "=",
+                "value": DocumentReviewStatus.PENDING_REVIEW.value,
+            }
+        ]
+        if nhs_number:
+            conditions.append(
+                {
+                    "field": "NhsNumber",
+                    "operator": "=",
+                    "value": nhs_number,
+                }
+            )
+
+        if uploader:
+            conditions.append(
+                {
+                    "field": "Author",
+                    "operator": "=",
+                    "value": uploader,
+                }
+            )
+
+        return build_mixed_condition_expression(conditions)
+
     def get_document(
         self, document_id: str, version: int | None
     ) -> DocumentUploadReviewReference | None:

lambdas/services/search_document_review_service.py

@@ -1,7 +1,3 @@
-import base64
-import decimal
-import json
-
 from enums.lambda_error import LambdaError
 from pydantic import ValidationError
 from services.document_upload_review_service import DocumentUploadReviewService
@@ -20,14 +16,13 @@ def process_request(
         self, ods_code: str, params: dict
     ) -> tuple[list[str], str | None]:
         try:
-
-            decoded_start_key = self.decode_start_key(params.get("nextPageToken", None))
+            start_key = params.get("nextPageToken", None)
 
             str_limit = params.get("limit", self.document_service.DEFAULT_QUERY_LIMIT)
             limit = int(str_limit)
 
             references, last_evaluated_key = self.get_review_document_references(
-                start_key=decoded_start_key,
+                start_key=start_key,
                 ods_code=ods_code,
                 limit=limit,
                 nhs_number=params.get("nhsNumber", None),
@@ -50,9 +45,7 @@ def process_request(
                 for reference in references
             ]
 
-            encoded_exclusive_start_key = self.encode_start_key(last_evaluated_key)
-
-            return output_refs, encoded_exclusive_start_key
+            return output_refs, last_evaluated_key
 
         except ValidationError as e:
             logger.error(e)
@@ -69,33 +62,14 @@ def get_review_document_references(
         self,
         ods_code: str,
         limit: int | None = None,
-        start_key: dict | None = None,
+        start_key: str | None = None,
         nhs_number: str | None = None,
         uploader: str | None = None,
     ):
-        return self.document_service.query_docs_pending_review_by_custodian_with_limit(
+        return self.document_service.query_docs_pending_review_with_paginator(
             ods_code=ods_code,
             limit=limit,
             start_key=start_key,
             nhs_number=nhs_number,
             uploader=uploader,
         )
-
-    def decode_start_key(self, encoded_start_key: str | None) -> dict[str, str] | None:
-        return (
-            json.loads(
-                base64.b64decode(encoded_start_key.encode("ascii")).decode("utf-8")
-            )
-            if encoded_start_key
-            else None
-        )
-
-    def encode_start_key(self, start_key: dict) -> str | None:
-        if start_key:
-            for key, value in start_key.items():
-                if isinstance(value, decimal.Decimal):
-                    start_key[key] = int(value)
-            return base64.b64encode(json.dumps(start_key).encode("ascii")).decode(
-                "utf-8"
-            )
-        return None