diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/base.py b/aws_lambda_powertools/utilities/idempotency/persistence/base.py index 28b284b8e5e..cad60cbe2b7 100644 --- a/aws_lambda_powertools/utilities/idempotency/persistence/base.py +++ b/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,7 +117,7 @@ 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 @@ -125,7 +125,7 @@ def __init__(self): 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: """ diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py b/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py index ce3aac2425e..654f8ca99d4 100644 --- a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py +++ b/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)}) diff --git a/tests/functional/idempotency/test_idempotency.py b/tests/functional/idempotency/test_idempotency.py index 68aeabeb50a..23d0537e533 100644 --- a/tests/functional/idempotency/test_idempotency.py +++ b/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) diff --git a/tests/unit/idempotency/__init__.py b/tests/unit/idempotency/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/idempotency/test_dynamodb_persistence.py b/tests/unit/idempotency/test_dynamodb_persistence.py new file mode 100644 index 00000000000..9455c41ad8d --- /dev/null +++ b/tests/unit/idempotency/test_dynamodb_persistence.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