sync

Author: Lucas McDonald

DynamoDbEncryption/runtimes/python/src/aws_dbesdk_dynamodb/encrypted/paginator.py

@@ -0,0 +1,199 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""High-level helper class to provide an encrypting wrapper for boto3 DynamoDB paginators."""
+from collections.abc import Callable, Generator
+from copy import deepcopy
+from typing import Any
+
+from botocore.paginate import (
+    Paginator,
+)
+
+from aws_dbesdk_dynamodb.encrypted.boto3_interface import EncryptedBotoInterface
+from aws_dbesdk_dynamodb.internal.client_to_resource import ClientShapeToResourceShapeConverter
+from aws_dbesdk_dynamodb.internal.resource_to_client import ResourceShapeToClientShapeConverter
+from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb.models import (
+    DynamoDbTablesEncryptionConfig,
+)
+from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb_transforms.client import (
+    DynamoDbEncryptionTransforms,
+)
+from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb_transforms.models import (
+    QueryInputTransformInput,
+    QueryOutputTransformInput,
+    ScanInputTransformInput,
+    ScanOutputTransformInput,
+)
+
+
+class EncryptedPaginator(EncryptedBotoInterface):
+    """Wrapping class for the boto3 Paginator that decrypts returned items before returning them."""
+
+    def __init__(
+        self,
+        *,
+        paginator: Paginator,
+        encryption_config: DynamoDbTablesEncryptionConfig,
+        expect_standard_dictionaries: bool | None = False,
+    ):
+        """
+        Create an EncryptedPaginator.
+
+        Args:
+            paginator (Paginator): A boto3 Paginator object for DynamoDB operations.
+                This can be either a "query" or "scan" Paginator.
+            encryption_config (DynamoDbTablesEncryptionConfig): Encryption configuration object.
+            expect_standard_dictionaries (Optional[bool]): Does the underlying boto3 client expect items
+                to be standard Python dictionaries? This should only be set to True if you are using a
+                client obtained from a service resource or table resource (ex: ``table.meta.client``).
+                If this is True, EncryptedClient will expect item-like shapes to be
+                standard Python dictionaries (default: False).
+
+        """
+        self._paginator = paginator
+        self._encryption_config = encryption_config
+        self._transformer = DynamoDbEncryptionTransforms(config=encryption_config)
+        self._expect_standard_dictionaries = expect_standard_dictionaries
+        self._resource_to_client_shape_converter = ResourceShapeToClientShapeConverter()
+        self._client_to_resource_shape_converter = ClientShapeToResourceShapeConverter(delete_table_name=False)
+
+    def paginate(self, **kwargs) -> Generator[dict, None, None]:
+        """
+        Yield a generator that paginates through responses from DynamoDB, decrypting items.
+
+        Note:
+            Calling ``botocore.paginate.Paginator``'s ``paginate`` method for Query or Scan
+            returns a ``PageIterator`` object, but this implementation returns a Python generator.
+            However, you can use this generator to iterate exactly as described in the
+            boto3 documentation:
+
+            https://botocore.amazonaws.com/v1/documentation/api/latest/topics/paginators.html
+
+            Any other operations on this class will defer to the underlying boto3 Paginator's implementation.
+
+        Args:
+            **kwargs: Keyword arguments passed directly to the underlying DynamoDB paginator.
+
+                For a Scan operation, structure these arguments according to:
+
+                https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb/paginator/Scan.html
+
+                For a Query operation, structure these arguments according to:
+
+                https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb/paginator/Query.html
+
+        Returns:
+            Generator[dict, None, None]: A generator yielding pages as dictionaries.
+            The items in the pages will be decrypted locally after being read from DynamoDB.
+
+        """
+        if self._paginator._model.name == "Query":
+            yield from self._paginate_query(**kwargs)
+        elif self._paginator._model.name == "Scan":
+            yield from self._paginate_scan(**kwargs)
+        else:
+            yield from self._paginator.paginate(**kwargs)
+
+    def _paginate_query(self, **paginate_query_kwargs):
+        return self._paginate_request(
+            paginate_kwargs=paginate_query_kwargs,
+            input_item_to_ddb_transform_method=self._resource_to_client_shape_converter.query_request,
+            input_item_to_dict_transform_method=self._client_to_resource_shape_converter.query_request,
+            input_transform_method=self._transformer.query_input_transform,
+            input_transform_shape=QueryInputTransformInput,
+            output_item_to_ddb_transform_method=self._resource_to_client_shape_converter.query_response,
+            output_item_to_dict_transform_method=self._client_to_resource_shape_converter.query_response,
+            output_transform_method=self._transformer.query_output_transform,
+            output_transform_shape=QueryOutputTransformInput,
+        )
+
+    def _paginate_scan(self, **paginate_scan_kwargs):
+        return self._paginate_request(
+            paginate_kwargs=paginate_scan_kwargs,
+            input_item_to_ddb_transform_method=self._resource_to_client_shape_converter.scan_request,
+            input_item_to_dict_transform_method=self._client_to_resource_shape_converter.scan_request,
+            input_transform_method=self._transformer.scan_input_transform,
+            input_transform_shape=ScanInputTransformInput,
+            output_item_to_ddb_transform_method=self._resource_to_client_shape_converter.scan_response,
+            output_item_to_dict_transform_method=self._client_to_resource_shape_converter.scan_response,
+            output_transform_method=self._transformer.scan_output_transform,
+            output_transform_shape=ScanOutputTransformInput,
+        )
+
+    def _paginate_request(
+        self,
+        *,
+        paginate_kwargs: dict[str, Any],
+        input_item_to_ddb_transform_method: Callable,
+        input_item_to_dict_transform_method: Callable,
+        input_transform_method: Callable,
+        input_transform_shape: Any,
+        output_item_to_ddb_transform_method: Callable,
+        output_item_to_dict_transform_method: Callable,
+        output_transform_method: Callable,
+        output_transform_shape: Any,
+    ):
+        client_kwargs = deepcopy(paginate_kwargs)
+        try:
+            # Remove PaginationConfig from the request if it exists.
+            # The input_transform_method does not expect it.
+            # It is added back to the request sent to the SDK.
+            pagination_config = client_kwargs["PaginationConfig"]
+            del client_kwargs["PaginationConfig"]
+        except KeyError:
+            pagination_config = None
+
+        # If _expect_standard_dictionaries is true, input items are expected to be standard dictionaries,
+        # and need to be converted to DDB-JSON before encryption.
+        if self._expect_standard_dictionaries:
+            if "TableName" in client_kwargs:
+                self._resource_to_client_shape_converter.table_name = client_kwargs["TableName"]
+            client_kwargs = input_item_to_ddb_transform_method(client_kwargs)
+
+        # Apply DBESDK transformations to the input
+        transformed_request = input_transform_method(input_transform_shape(sdk_input=client_kwargs)).transformed_input
+
+        # If _expect_standard_dictionaries is true, the boto3 client expects items to be standard dictionaries,
+        # and need to be converted from DDB-JSON to a standard dictionary before being passed to the boto3 client.
+        if self._expect_standard_dictionaries:
+            transformed_request = input_item_to_dict_transform_method(transformed_request)
+
+        if pagination_config is not None:
+            transformed_request["PaginationConfig"] = pagination_config
+
+        sdk_page_response = self._paginator.paginate(**transformed_request)
+
+        for page in sdk_page_response:
+            # If _expect_standard_dictionaries is true, the boto3 client returns items as standard dictionaries,
+            # and needs to convert the standard dictionary to DDB-JSON before passing the response to the DBESDK.
+            if self._expect_standard_dictionaries:
+                page = output_item_to_ddb_transform_method(page)
+
+            # Apply DBESDK transformation to the boto3 output
+            dbesdk_response = output_transform_method(
+                output_transform_shape(
+                    original_input=client_kwargs,
+                    sdk_output=page,
+                )
+            ).transformed_output
+
+            # Copy any missing fields from the SDK output to the response (e.g. ConsumedCapacity)
+            dbesdk_response = self._copy_sdk_response_to_dbesdk_response(page, dbesdk_response)
+
+            # If _expect_standard_dictionaries is true, the boto3 client expects items to be standard dictionaries,
+            # and need to be converted from DDB-JSON to a standard dictionary before returning the response.
+            if self._expect_standard_dictionaries:
+                dbesdk_response = output_item_to_dict_transform_method(dbesdk_response)
+
+            yield dbesdk_response
+
+    @property
+    def _boto_client_attr_name(self) -> str:
+        """
+        Name of the attribute containing the underlying boto3 client.
+
+        Returns:
+            str: '_paginator'
+
+        """
+        return "_paginator"

DynamoDbEncryption/runtimes/python/test/integ/encrypted/test_paginator.py

@@ -0,0 +1,190 @@
+import boto3
+import pytest
+
+from aws_dbesdk_dynamodb.encrypted.client import EncryptedClient
+
+from ...constants import (
+    INTEG_TEST_DEFAULT_DYNAMODB_TABLE_NAME,
+    INTEG_TEST_DEFAULT_TABLE_CONFIGS,
+)
+from ...items import (
+    complex_item_ddb,
+    complex_item_dict,
+    complex_key_ddb,
+    complex_key_dict,
+    simple_item_ddb,
+    simple_item_dict,
+    simple_key_ddb,
+    simple_key_dict,
+)
+from ...requests import (
+    basic_put_item_request_ddb,
+    basic_put_item_request_dict,
+    basic_query_paginator_request,
+    basic_scan_paginator_request,
+)
+from . import sort_dynamodb_json_lists
+
+
+# Creates a matrix of tests for each value in param,
+# with a user-friendly string for test output:
+# expect_standard_dictionaries = True -> "standard_dicts"
+# expect_standard_dictionaries = False -> "ddb_json"
+@pytest.fixture(params=[True, False], ids=["standard_dicts", "ddb_json"])
+def expect_standard_dictionaries(request):
+    return request.param
+
+
+def encrypted_client(expect_standard_dictionaries):
+    return EncryptedClient(
+        client=plaintext_client(expect_standard_dictionaries),
+        encryption_config=INTEG_TEST_DEFAULT_TABLE_CONFIGS,
+        expect_standard_dictionaries=expect_standard_dictionaries,
+    )
+
+
+def plaintext_client(expect_standard_dictionaries):
+    if expect_standard_dictionaries:
+        client = boto3.resource("dynamodb").Table(INTEG_TEST_DEFAULT_DYNAMODB_TABLE_NAME).meta.client
+    else:
+        client = boto3.client("dynamodb")
+    return client
+
+
+# Creates a matrix of tests for each value in param,
+# with a user-friendly string for test output:
+# encrypted = True -> "encrypted"
+# encrypted = False -> "plaintext"
+@pytest.fixture(params=[True, False], ids=["encrypted", "plaintext"])
+def encrypted(request):
+    return request.param
+
+
+@pytest.fixture
+def client(encrypted, expect_standard_dictionaries):
+    if encrypted:
+        return encrypted_client(expect_standard_dictionaries)
+    else:
+        return plaintext_client(expect_standard_dictionaries)
+
+
+@pytest.fixture
+def query_paginator(client):
+    return client.get_paginator("query")
+
+
+@pytest.fixture
+def scan_paginator(client):
+    return client.get_paginator("scan")
+
+
+# Creates a matrix of tests for each value in param,
+# with a user-friendly string for test output:
+# use_complex_item = True -> "complex_item"
+# use_complex_item = False -> "simple_item"
+@pytest.fixture(params=[True, False], ids=["complex_item", "simple_item"])
+def use_complex_item(request):
+    return request.param
+
+
+@pytest.fixture
+def test_key(expect_standard_dictionaries, use_complex_item):
+    """Get a single test item in the appropriate format for the client."""
+    if expect_standard_dictionaries:
+        if use_complex_item:
+            return complex_key_dict
+        return simple_key_dict
+    if use_complex_item:
+        return complex_key_ddb
+    return simple_key_ddb
+
+
+@pytest.fixture
+def multiple_test_keys(expect_standard_dictionaries):
+    """Get two test keys in the appropriate format for the client."""
+    if expect_standard_dictionaries:
+        return [simple_key_dict, complex_key_dict]
+    return [simple_key_ddb, complex_key_ddb]
+
+
+@pytest.fixture
+def test_item(expect_standard_dictionaries, use_complex_item):
+    """Get a single test item in the appropriate format for the client."""
+    if expect_standard_dictionaries:
+        if use_complex_item:
+            return complex_item_dict
+        return simple_item_dict
+    if use_complex_item:
+        return complex_item_ddb
+    return simple_item_ddb
+
+
+@pytest.fixture
+def paginate_query_request(expect_standard_dictionaries, test_key):
+    if expect_standard_dictionaries:
+        return {**basic_query_paginator_request(test_key), "TableName": INTEG_TEST_DEFAULT_DYNAMODB_TABLE_NAME}
+    return basic_query_paginator_request(test_key)
+
+
+@pytest.fixture
+def put_item_request(expect_standard_dictionaries, test_item):
+    if expect_standard_dictionaries:
+        # Client requests with `expect_standard_dictionaries=True` use dict-formatted requests
+        # with an added "TableName" key.
+        return {**basic_put_item_request_dict(test_item), "TableName": INTEG_TEST_DEFAULT_DYNAMODB_TABLE_NAME}
+    return basic_put_item_request_ddb(test_item)
+
+
+def test_GIVEN_query_paginator_WHEN_paginate_THEN_returns_expected_items(
+    client, query_paginator, paginate_query_request, put_item_request, test_item
+):
+    # Given: item in table
+    client.put_item(**put_item_request)
+    # Given: Query paginator
+    # When: Paginate
+    response = query_paginator.paginate(**paginate_query_request)
+    # Then: Returns encrypted items
+    items = []
+    for page in response:
+        if "Items" in page:
+            for item in page["Items"]:
+                items.append(item)
+    assert len(items) == 1
+    # DynamoDB JSON uses lists to represent sets, so strict equality can fail.
+    # Sort lists to ensure consistent ordering when comparing expected and actual items.
+    expected_item = sort_dynamodb_json_lists(test_item)
+    actual_item = sort_dynamodb_json_lists(items[0])
+    # Then: Items are equal
+    assert expected_item == actual_item
+
+
+@pytest.fixture
+def paginate_scan_request(expect_standard_dictionaries, test_item):
+    if expect_standard_dictionaries:
+        request = {**basic_scan_paginator_request(test_item), "TableName": INTEG_TEST_DEFAULT_DYNAMODB_TABLE_NAME}
+    else:
+        request = basic_scan_paginator_request(test_item)
+    return request
+
+
+def test_GIVEN_scan_paginator_WHEN_paginate_THEN_returns_expected_items(
+    client, scan_paginator, paginate_scan_request, put_item_request, test_item
+):
+    # Given: item in table
+    client.put_item(**put_item_request)
+    # Given: Scan paginator
+    # When: Paginate
+    response = scan_paginator.paginate(**paginate_scan_request)
+    # Then: Returns encrypted items
+    items = []
+    for page in response:
+        if "Items" in page:
+            for item in page["Items"]:
+                items.append(item)
+    assert len(items) == 1
+    # DynamoDB JSON uses lists to represent sets, so strict equality can fail.
+    # Sort lists to ensure consistent ordering when comparing expected and actual items.
+    expected_item = sort_dynamodb_json_lists(test_item)
+    actual_item = sort_dynamodb_json_lists(items[0])
+    # Then: Items are equal
+    assert expected_item == actual_item