Skip to content

Commit 665749d

Browse files
author
Lucas McDonald
committed
sync
1 parent 0e4fd02 commit 665749d

File tree

6 files changed

+756
-0
lines changed

6 files changed

+756
-0
lines changed
Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,199 @@
1+
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""High-level helper class to provide an encrypting wrapper for boto3 DynamoDB paginators."""
4+
from collections.abc import Callable, Generator
5+
from copy import deepcopy
6+
from typing import Any
7+
8+
from botocore.paginate import (
9+
Paginator,
10+
)
11+
12+
from aws_dbesdk_dynamodb.encrypted.boto3_interface import EncryptedBotoInterface
13+
from aws_dbesdk_dynamodb.internal.client_to_resource import ClientShapeToResourceShapeConverter
14+
from aws_dbesdk_dynamodb.internal.resource_to_client import ResourceShapeToClientShapeConverter
15+
from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb.models import (
16+
DynamoDbTablesEncryptionConfig,
17+
)
18+
from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb_transforms.client import (
19+
DynamoDbEncryptionTransforms,
20+
)
21+
from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb_transforms.models import (
22+
QueryInputTransformInput,
23+
QueryOutputTransformInput,
24+
ScanInputTransformInput,
25+
ScanOutputTransformInput,
26+
)
27+
28+
29+
class EncryptedPaginator(EncryptedBotoInterface):
30+
"""Wrapping class for the boto3 Paginator that decrypts returned items before returning them."""
31+
32+
def __init__(
33+
self,
34+
*,
35+
paginator: Paginator,
36+
encryption_config: DynamoDbTablesEncryptionConfig,
37+
expect_standard_dictionaries: bool | None = False,
38+
):
39+
"""
40+
Create an EncryptedPaginator.
41+
42+
Args:
43+
paginator (Paginator): A boto3 Paginator object for DynamoDB operations.
44+
This can be either a "query" or "scan" Paginator.
45+
encryption_config (DynamoDbTablesEncryptionConfig): Encryption configuration object.
46+
expect_standard_dictionaries (Optional[bool]): Does the underlying boto3 client expect items
47+
to be standard Python dictionaries? This should only be set to True if you are using a
48+
client obtained from a service resource or table resource (ex: ``table.meta.client``).
49+
If this is True, EncryptedClient will expect item-like shapes to be
50+
standard Python dictionaries (default: False).
51+
52+
"""
53+
self._paginator = paginator
54+
self._encryption_config = encryption_config
55+
self._transformer = DynamoDbEncryptionTransforms(config=encryption_config)
56+
self._expect_standard_dictionaries = expect_standard_dictionaries
57+
self._resource_to_client_shape_converter = ResourceShapeToClientShapeConverter()
58+
self._client_to_resource_shape_converter = ClientShapeToResourceShapeConverter(delete_table_name=False)
59+
60+
def paginate(self, **kwargs) -> Generator[dict, None, None]:
61+
"""
62+
Yield a generator that paginates through responses from DynamoDB, decrypting items.
63+
64+
Note:
65+
Calling ``botocore.paginate.Paginator``'s ``paginate`` method for Query or Scan
66+
returns a ``PageIterator`` object, but this implementation returns a Python generator.
67+
However, you can use this generator to iterate exactly as described in the
68+
boto3 documentation:
69+
70+
https://botocore.amazonaws.com/v1/documentation/api/latest/topics/paginators.html
71+
72+
Any other operations on this class will defer to the underlying boto3 Paginator's implementation.
73+
74+
Args:
75+
**kwargs: Keyword arguments passed directly to the underlying DynamoDB paginator.
76+
77+
For a Scan operation, structure these arguments according to:
78+
79+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb/paginator/Scan.html
80+
81+
For a Query operation, structure these arguments according to:
82+
83+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb/paginator/Query.html
84+
85+
Returns:
86+
Generator[dict, None, None]: A generator yielding pages as dictionaries.
87+
The items in the pages will be decrypted locally after being read from DynamoDB.
88+
89+
"""
90+
if self._paginator._model.name == "Query":
91+
yield from self._paginate_query(**kwargs)
92+
elif self._paginator._model.name == "Scan":
93+
yield from self._paginate_scan(**kwargs)
94+
else:
95+
yield from self._paginator.paginate(**kwargs)
96+
97+
def _paginate_query(self, **paginate_query_kwargs):
98+
return self._paginate_request(
99+
paginate_kwargs=paginate_query_kwargs,
100+
input_item_to_ddb_transform_method=self._resource_to_client_shape_converter.query_request,
101+
input_item_to_dict_transform_method=self._client_to_resource_shape_converter.query_request,
102+
input_transform_method=self._transformer.query_input_transform,
103+
input_transform_shape=QueryInputTransformInput,
104+
output_item_to_ddb_transform_method=self._resource_to_client_shape_converter.query_response,
105+
output_item_to_dict_transform_method=self._client_to_resource_shape_converter.query_response,
106+
output_transform_method=self._transformer.query_output_transform,
107+
output_transform_shape=QueryOutputTransformInput,
108+
)
109+
110+
def _paginate_scan(self, **paginate_scan_kwargs):
111+
return self._paginate_request(
112+
paginate_kwargs=paginate_scan_kwargs,
113+
input_item_to_ddb_transform_method=self._resource_to_client_shape_converter.scan_request,
114+
input_item_to_dict_transform_method=self._client_to_resource_shape_converter.scan_request,
115+
input_transform_method=self._transformer.scan_input_transform,
116+
input_transform_shape=ScanInputTransformInput,
117+
output_item_to_ddb_transform_method=self._resource_to_client_shape_converter.scan_response,
118+
output_item_to_dict_transform_method=self._client_to_resource_shape_converter.scan_response,
119+
output_transform_method=self._transformer.scan_output_transform,
120+
output_transform_shape=ScanOutputTransformInput,
121+
)
122+
123+
def _paginate_request(
124+
self,
125+
*,
126+
paginate_kwargs: dict[str, Any],
127+
input_item_to_ddb_transform_method: Callable,
128+
input_item_to_dict_transform_method: Callable,
129+
input_transform_method: Callable,
130+
input_transform_shape: Any,
131+
output_item_to_ddb_transform_method: Callable,
132+
output_item_to_dict_transform_method: Callable,
133+
output_transform_method: Callable,
134+
output_transform_shape: Any,
135+
):
136+
client_kwargs = deepcopy(paginate_kwargs)
137+
try:
138+
# Remove PaginationConfig from the request if it exists.
139+
# The input_transform_method does not expect it.
140+
# It is added back to the request sent to the SDK.
141+
pagination_config = client_kwargs["PaginationConfig"]
142+
del client_kwargs["PaginationConfig"]
143+
except KeyError:
144+
pagination_config = None
145+
146+
# If _expect_standard_dictionaries is true, input items are expected to be standard dictionaries,
147+
# and need to be converted to DDB-JSON before encryption.
148+
if self._expect_standard_dictionaries:
149+
if "TableName" in client_kwargs:
150+
self._resource_to_client_shape_converter.table_name = client_kwargs["TableName"]
151+
client_kwargs = input_item_to_ddb_transform_method(client_kwargs)
152+
153+
# Apply DBESDK transformations to the input
154+
transformed_request = input_transform_method(input_transform_shape(sdk_input=client_kwargs)).transformed_input
155+
156+
# If _expect_standard_dictionaries is true, the boto3 client expects items to be standard dictionaries,
157+
# and need to be converted from DDB-JSON to a standard dictionary before being passed to the boto3 client.
158+
if self._expect_standard_dictionaries:
159+
transformed_request = input_item_to_dict_transform_method(transformed_request)
160+
161+
if pagination_config is not None:
162+
transformed_request["PaginationConfig"] = pagination_config
163+
164+
sdk_page_response = self._paginator.paginate(**transformed_request)
165+
166+
for page in sdk_page_response:
167+
# If _expect_standard_dictionaries is true, the boto3 client returns items as standard dictionaries,
168+
# and needs to convert the standard dictionary to DDB-JSON before passing the response to the DBESDK.
169+
if self._expect_standard_dictionaries:
170+
page = output_item_to_ddb_transform_method(page)
171+
172+
# Apply DBESDK transformation to the boto3 output
173+
dbesdk_response = output_transform_method(
174+
output_transform_shape(
175+
original_input=client_kwargs,
176+
sdk_output=page,
177+
)
178+
).transformed_output
179+
180+
# Copy any missing fields from the SDK output to the response (e.g. ConsumedCapacity)
181+
dbesdk_response = self._copy_sdk_response_to_dbesdk_response(page, dbesdk_response)
182+
183+
# If _expect_standard_dictionaries is true, the boto3 client expects items to be standard dictionaries,
184+
# and need to be converted from DDB-JSON to a standard dictionary before returning the response.
185+
if self._expect_standard_dictionaries:
186+
dbesdk_response = output_item_to_dict_transform_method(dbesdk_response)
187+
188+
yield dbesdk_response
189+
190+
@property
191+
def _boto_client_attr_name(self) -> str:
192+
"""
193+
Name of the attribute containing the underlying boto3 client.
194+
195+
Returns:
196+
str: '_paginator'
197+
198+
"""
199+
return "_paginator"
Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
Integration tests for encrypted interfaces.
2+
3+
These integration tests verify that encrypted boto3 interfaces behave as drop-in replacements for plaintext boto3 interfaces.
4+
5+
Each test runs with both a plaintext client and an encrypted client, using the same request parameters and expecting the same response.
6+
7+
This validates that encrypted clients expect the same input shapes as plaintext clients
8+
and encrypted clients return the same output shapes as plaintext clients.
9+
10+
This guarantees that users can substitute encrypted interfaces without modifying their application logic.

0 commit comments

Comments
 (0)