feat(idempotency): allow custom sdk clients in DynamoDBPersistenceLayer (#2087)

Author: heitorlessa

aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -41,8 +41,8 @@ def __init__(
         status: str = "",
         expiry_timestamp: Optional[int] = None,
         in_progress_expiry_timestamp: Optional[int] = None,
-        response_data: Optional[str] = "",
-        payload_hash: Optional[str] = None,
+        response_data: str = "",
+        payload_hash: str = "",
     ) -> None:
         """
 
@@ -117,15 +117,15 @@ def __init__(self):
         """Initialize the defaults"""
         self.function_name = ""
         self.configured = False
-        self.event_key_jmespath: Optional[str] = None
+        self.event_key_jmespath: str = ""
         self.event_key_compiled_jmespath = None
         self.jmespath_options: Optional[dict] = None
         self.payload_validation_enabled = False
         self.validation_key_jmespath = None
         self.raise_on_no_idempotency_key = False
         self.expires_after_seconds: int = 60 * 60  # 1 hour default
         self.use_local_cache = False
-        self.hash_function = None
+        self.hash_function = hashlib.md5
 
     def configure(self, config: IdempotencyConfig, function_name: Optional[str] = None) -> None:
         """

aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py

@@ -1,7 +1,9 @@
+from __future__ import annotations
+
 import datetime
 import logging
 import os
-from typing import Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 import boto3
 from boto3.dynamodb.types import TypeDeserializer
@@ -19,6 +21,10 @@
     DataRecord,
 )
 
+if TYPE_CHECKING:
+    from mypy_boto3_dynamodb import DynamoDBClient
+    from mypy_boto3_dynamodb.type_defs import AttributeValueTypeDef
+
 logger = logging.getLogger(__name__)
 
 
@@ -36,6 +42,7 @@ def __init__(
         validation_key_attr: str = "validation",
         boto_config: Optional[Config] = None,
         boto3_session: Optional[boto3.session.Session] = None,
+        boto3_client: "DynamoDBClient" | None = None,
     ):
         """
         Initialize the DynamoDB client
@@ -61,8 +68,10 @@ def __init__(
             DynamoDB attribute name for response data, by default "data"
         boto_config: botocore.config.Config, optional
             Botocore configuration to pass during client initialization
-        boto3_session : boto3.session.Session, optional
+        boto3_session : boto3.Session, optional
             Boto3 session to use for AWS API communication
+        boto3_client : DynamoDBClient, optional
+            Boto3 DynamoDB Client to use, boto3_session and boto_config will be ignored if both are provided
 
         Examples
         --------
@@ -78,10 +87,12 @@ def __init__(
             >>> def handler(event, context):
             >>>     return {"StatusCode": 200}
         """
-
-        self._boto_config = boto_config or Config()
-        self._boto3_session = boto3_session or boto3.session.Session()
-        self._client = self._boto3_session.client("dynamodb", config=self._boto_config)
+        if boto3_client is None:
+            self._boto_config = boto_config or Config()
+            self._boto3_session: boto3.Session = boto3_session or boto3.session.Session()
+            self.client: "DynamoDBClient" = self._boto3_session.client("dynamodb", config=self._boto_config)
+        else:
+            self.client = boto3_client
 
         if sort_key_attr == key_attr:
             raise ValueError(f"key_attr [{key_attr}] and sort_key_attr [{sort_key_attr}] cannot be the same!")
@@ -149,7 +160,7 @@ def _item_to_data_record(self, item: Dict[str, Any]) -> DataRecord:
         )
 
     def _get_record(self, idempotency_key) -> DataRecord:
-        response = self._client.get_item(
+        response = self.client.get_item(
             TableName=self.table_name, Key=self._get_key(idempotency_key), ConsistentRead=True
         )
         try:
@@ -204,7 +215,7 @@ def _put_record(self, data_record: DataRecord) -> None:
             condition_expression = (
                 f"{idempotency_key_not_exist} OR {idempotency_expiry_expired} OR ({inprogress_expiry_expired})"
             )
-            self._client.put_item(
+            self.client.put_item(
                 TableName=self.table_name,
                 Item=item,
                 ConditionExpression=condition_expression,
@@ -233,7 +244,7 @@ def _put_record(self, data_record: DataRecord) -> None:
     def _update_record(self, data_record: DataRecord):
         logger.debug(f"Updating record for idempotency key: {data_record.idempotency_key}")
         update_expression = "SET #response_data = :response_data, #expiry = :expiry, " "#status = :status"
-        expression_attr_values = {
+        expression_attr_values: Dict[str, "AttributeValueTypeDef"] = {
             ":expiry": {"N": str(data_record.expiry_timestamp)},
             ":response_data": {"S": data_record.response_data},
             ":status": {"S": data_record.status},
@@ -249,15 +260,14 @@ def _update_record(self, data_record: DataRecord):
             expression_attr_values[":validation_key"] = {"S": data_record.payload_hash}
             expression_attr_names["#validation_key"] = self.validation_key_attr
 
-        kwargs = {
-            "Key": self._get_key(data_record.idempotency_key),
-            "UpdateExpression": update_expression,
-            "ExpressionAttributeValues": expression_attr_values,
-            "ExpressionAttributeNames": expression_attr_names,
-        }
-
-        self._client.update_item(TableName=self.table_name, **kwargs)
+        self.client.update_item(
+            TableName=self.table_name,
+            Key=self._get_key(data_record.idempotency_key),
+            UpdateExpression=update_expression,
+            ExpressionAttributeNames=expression_attr_names,
+            ExpressionAttributeValues=expression_attr_values,
+        )
 
     def _delete_record(self, data_record: DataRecord) -> None:
         logger.debug(f"Deleting record for idempotency key: {data_record.idempotency_key}")
-        self._client.delete_item(TableName=self.table_name, Key={**self._get_key(data_record.idempotency_key)})
+        self.client.delete_item(TableName=self.table_name, Key={**self._get_key(data_record.idempotency_key)})

tests/functional/idempotency/test_idempotency.py

@@ -76,7 +76,7 @@ def test_idempotent_lambda_already_completed(
     Test idempotent decorator where event with matching event key has already been successfully processed
     """
 
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
     ddb_response = {
         "Item": {
             "id": {"S": hashed_idempotency_key},
@@ -120,7 +120,7 @@ def test_idempotent_lambda_in_progress(
     Test idempotent decorator where lambda_handler is already processing an event with matching event key
     """
 
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
 
     expected_params = {
         "TableName": TABLE_NAME,
@@ -172,7 +172,7 @@ def test_idempotent_lambda_in_progress_with_cache(
     """
     save_to_cache_spy = mocker.spy(persistence_store, "_save_to_cache")
     retrieve_from_cache_spy = mocker.spy(persistence_store, "_retrieve_from_cache")
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
 
     expected_params = {
         "TableName": TABLE_NAME,
@@ -234,7 +234,7 @@ def test_idempotent_lambda_first_execution(
     Test idempotent decorator when lambda is executed with an event with a previously unknown event key
     """
 
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
     ddb_response = {}
 
     stubber.add_response("put_item", ddb_response, expected_params_put_item)
@@ -269,7 +269,7 @@ def test_idempotent_lambda_first_execution_cached(
     """
     save_to_cache_spy = mocker.spy(persistence_store, "_save_to_cache")
     retrieve_from_cache_spy = mocker.spy(persistence_store, "_retrieve_from_cache")
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
     ddb_response = {}
 
     stubber.add_response("put_item", ddb_response, expected_params_put_item)
@@ -310,7 +310,7 @@ def test_idempotent_lambda_first_execution_event_mutation(
     Ensures we're passing data by value, not reference.
     """
     event = copy.deepcopy(lambda_apigw_event)
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
     ddb_response = {}
     stubber.add_response(
         "put_item",
@@ -350,7 +350,7 @@ def test_idempotent_lambda_expired(
     expiry window
     """
 
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
 
     ddb_response = {}
 
@@ -385,7 +385,7 @@ def test_idempotent_lambda_exception(
     # Create a new provider
 
     # Stub the boto3 client
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
 
     ddb_response = {}
     expected_params_delete_item = {"TableName": TABLE_NAME, "Key": {"id": {"S": hashed_idempotency_key}}}
@@ -427,7 +427,7 @@ def test_idempotent_lambda_already_completed_with_validation_bad_payload(
     Test idempotent decorator where event with matching event key has already been successfully processed
     """
 
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
     ddb_response = {
         "Item": {
             "id": {"S": hashed_idempotency_key},
@@ -471,7 +471,7 @@ def test_idempotent_lambda_expired_during_request(
     returns inconsistent/rapidly changing result between put_item and get_item calls.
     """
 
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
 
     ddb_response_get_item = {
         "Item": {
@@ -524,7 +524,7 @@ def test_idempotent_persistence_exception_deleting(
     Test idempotent decorator when lambda is executed with an event with a previously unknown event key, but
     lambda_handler raises an exception which is retryable.
     """
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
 
     ddb_response = {}
 
@@ -556,7 +556,7 @@ def test_idempotent_persistence_exception_updating(
     Test idempotent decorator when lambda is executed with an event with a previously unknown event key, but
     lambda_handler raises an exception which is retryable.
     """
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
 
     ddb_response = {}
 
@@ -587,7 +587,7 @@ def test_idempotent_persistence_exception_getting(
     Test idempotent decorator when lambda is executed with an event with a previously unknown event key, but
     lambda_handler raises an exception which is retryable.
     """
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
 
     stubber.add_client_error("put_item", "ConditionalCheckFailedException")
     stubber.add_client_error("get_item", "UnexpectedException")
@@ -625,7 +625,7 @@ def test_idempotent_lambda_first_execution_with_validation(
     """
     Test idempotent decorator when lambda is executed with an event with a previously unknown event key
     """
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
     ddb_response = {}
 
     stubber.add_response("put_item", ddb_response, expected_params_put_item_with_validation)
@@ -661,7 +661,7 @@ def test_idempotent_lambda_with_validator_util(
     validator utility to unwrap the event
     """
 
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
     ddb_response = {
         "Item": {
             "id": {"S": hashed_idempotency_key_with_envelope},
@@ -704,7 +704,7 @@ def test_idempotent_lambda_expires_in_progress_before_expire(
     hashed_idempotency_key,
     lambda_context,
 ):
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
 
     stubber.add_client_error("put_item", "ConditionalCheckFailedException")
 
@@ -751,7 +751,7 @@ def test_idempotent_lambda_expires_in_progress_after_expire(
     hashed_idempotency_key,
     lambda_context,
 ):
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
 
     for _ in range(MAX_RETRIES + 1):
         stubber.add_client_error("put_item", "ConditionalCheckFailedException")
@@ -1070,7 +1070,7 @@ def test_custom_jmespath_function_overrides_builtin_functions(
 def test_idempotent_lambda_save_inprogress_error(persistence_store: DynamoDBPersistenceLayer, lambda_context):
     # GIVEN a miss configured persistence layer
     # like no table was created for the idempotency persistence layer
-    stubber = stub.Stubber(persistence_store._client)
+    stubber = stub.Stubber(persistence_store.client)
     service_error_code = "ResourceNotFoundException"
     service_message = "Custom message"
 
@@ -1327,7 +1327,7 @@ def test_idempotency_disabled_envvar(monkeypatch, lambda_context, persistence_st
     # Scenario to validate no requests sent to dynamodb table when 'POWERTOOLS_IDEMPOTENCY_DISABLED' is set
     mock_event = {"data": "value"}
 
-    persistence_store._client = MagicMock()
+    persistence_store.client = MagicMock()
 
     monkeypatch.setenv("POWERTOOLS_IDEMPOTENCY_DISABLED", "1")
 
@@ -1342,7 +1342,7 @@ def dummy_handler(event, context):
     dummy(data=mock_event)
     dummy_handler(mock_event, lambda_context)
 
-    assert len(persistence_store._client.method_calls) == 0
+    assert len(persistence_store.client.method_calls) == 0
 
 
 @pytest.mark.parametrize("idempotency_config", [{"use_local_cache": True}], indirect=True)
@@ -1351,7 +1351,7 @@ def test_idempotent_function_duplicates(
 ):
     # Scenario to validate the both methods are called
     mock_event = {"data": "value"}
-    persistence_store._client = MagicMock()
+    persistence_store.client = MagicMock()
 
     @idempotent_function(data_keyword_argument="data", persistence_store=persistence_store, config=idempotency_config)
     def one(data):
@@ -1363,7 +1363,7 @@ def two(data):
 
     assert one(data=mock_event) == "one"
     assert two(data=mock_event) == "two"
-    assert len(persistence_store._client.method_calls) == 4
+    assert len(persistence_store.client.method_calls) == 4
 
 
 def test_invalid_dynamodb_persistence_layer():
@@ -1475,7 +1475,7 @@ def test_idempotent_lambda_compound_already_completed(
     Test idempotent decorator having a DynamoDBPersistenceLayer with a compound key
     """
 
-    stubber = stub.Stubber(persistence_store_compound._client)
+    stubber = stub.Stubber(persistence_store_compound.client)
     stubber.add_client_error("put_item", "ConditionalCheckFailedException")
     ddb_response = {
         "Item": {
@@ -1520,7 +1520,7 @@ def test_idempotent_lambda_compound_static_pk_value_has_correct_pk(
     Test idempotent decorator having a DynamoDBPersistenceLayer with a compound key and a static PK value
     """
 
-    stubber = stub.Stubber(persistence_store_compound_static_pk_value._client)
+    stubber = stub.Stubber(persistence_store_compound_static_pk_value.client)
     ddb_response = {}
 
     stubber.add_response("put_item", ddb_response, expected_params_put_item_compound_key_static_pk_value)

tests/unit/idempotency/__init__.py

@@ -0,0 +1,21 @@
+from dataclasses import dataclass
+
+from aws_lambda_powertools.utilities.idempotency import DynamoDBPersistenceLayer
+from tests.e2e.utils.data_builder.common import build_random_value
+
+
+def test_custom_sdk_client_injection():
+    # GIVEN
+    @dataclass
+    class DummyClient:
+        table_name: str
+
+    table_name = build_random_value()
+    fake_client = DummyClient(table_name)
+
+    # WHEN
+    persistence_layer = DynamoDBPersistenceLayer(table_name, boto3_client=fake_client)
+
+    # THEN
+    assert persistence_layer.table_name == table_name
+    assert persistence_layer.client == fake_client