diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py b/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py index f652fb0f0a2..5fb4aed159c 100644 --- a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py +++ b/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py @@ -162,7 +162,7 @@ def _item_to_data_record(self, item: dict[str, Any]) -> DataRecord: """ data = self._deserializer.deserialize({"M": item}) return DataRecord( - idempotency_key=data[self.key_attr], + idempotency_key=data[self.sort_key_attr] if self.sort_key_attr else data[self.key_attr], status=data[self.status_attr], expiry_timestamp=data[self.expiry_attr], in_progress_expiry_timestamp=data.get(self.in_progress_expiry_attr), diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md index da7e70ad89c..3c6e64e1247 100644 --- a/docs/utilities/idempotency.md +++ b/docs/utilities/idempotency.md @@ -983,6 +983,8 @@ You can create your own persistent store from scratch by inheriting the `BasePer ### Manipulating the Idempotent Response +With a DynamoDB composite key, `DataRecord.idempotency_key` contains the sort key value identifying the operation, rather than the shared static partition key. Replayed records use this same key for local caching. + You can set up a `response_hook` in the `IdempotentConfig` class to manipulate the returned data when an operation is idempotent. The hook function will be called with the current deserialized response object and the Idempotency record. === "Using an Idempotent Response Hook" diff --git a/examples/idempotency/src/working_with_response_hook.py b/examples/idempotency/src/working_with_response_hook.py index 68768a12548..52248357ae5 100644 --- a/examples/idempotency/src/working_with_response_hook.py +++ b/examples/idempotency/src/working_with_response_hook.py @@ -16,7 +16,7 @@ def my_response_hook(response: dict, idempotent_data: DataRecord) -> dict: - # Return inserted Header data into the Idempotent Response + # With a composite key, this is the operation's sort key, not the shared partition key response["x-idempotent-key"] = idempotent_data.idempotency_key # expiry_timestamp can be None so include if set diff --git a/tests/functional/idempotency/_boto3/test_composite_key_records.py b/tests/functional/idempotency/_boto3/test_composite_key_records.py new file mode 100644 index 00000000000..e01942630d2 --- /dev/null +++ b/tests/functional/idempotency/_boto3/test_composite_key_records.py @@ -0,0 +1,134 @@ +import pytest +from botocore.stub import Stubber + +from aws_lambda_powertools.utilities.idempotency import DynamoDBPersistenceLayer, IdempotencyConfig, idempotent +from aws_lambda_powertools.utilities.idempotency.exceptions import IdempotencyAlreadyInProgressError + + +@pytest.fixture(params=[None, "sk", "record_key"]) +def record_store(request, config): + return DynamoDBPersistenceLayer( + table_name="TEST_TABLE", + boto_config=config, + sort_key_attr=request.param, + static_pk_value="shared-partition", + ) + + +def stored_item(store, key, expiry, status="COMPLETED"): + item = { + "id": {"S": "shared-partition" if store.sort_key_attr else key}, + "expiration": {"N": expiry}, + "status": {"S": status}, + "data": {"S": '{"message": "stored"}'}, + } + if store.sort_key_attr: + item[store.sort_key_attr] = {"S": key} + return item + + +def stub_existing_record(stubber, store, item, source): + if source == "condition": + stubber.add_client_error("put_item", "ConditionalCheckFailedException", modeled_fields={"Item": item}) + else: + stubber.add_client_error("put_item", "ConditionalCheckFailedException") + key = {"id": item["id"]} + if store.sort_key_attr: + key[store.sort_key_attr] = item[store.sort_key_attr] + stubber.add_response( + "get_item", + {"Item": item}, + { + "TableName": "TEST_TABLE", + "Key": key, + "ConsistentRead": True, + }, + ) + + +@pytest.mark.parametrize("source", ["condition", "get_item"]) +@pytest.mark.parametrize("use_local_cache", [False, True]) +def test_response_hook_receives_record_key( + record_store, + source, + use_local_cache, + lambda_apigw_event, + lambda_context, + timestamp_future, + hashed_idempotency_key, + default_jmespath, +): + keys = [] + + def response_hook(response, record): + keys.append((record.idempotency_key, record.sort_key)) + return response + + config = IdempotencyConfig( + event_key_jmespath=default_jmespath, + use_local_cache=use_local_cache, + response_hook=response_hook, + ) + + @idempotent(config=config, persistence_store=record_store) + def lambda_handler(event, context): + pytest.fail("A stored response must bypass the handler") + + item = stored_item(record_store, hashed_idempotency_key, timestamp_future) + with Stubber(record_store.client) as stubber: + stub_existing_record(stubber, record_store, item, source) + assert lambda_handler(lambda_apigw_event, lambda_context) == {"message": "stored"} + stubber.assert_no_pending_responses() + + expected_sort_key = hashed_idempotency_key if record_store.sort_key_attr else None + assert keys == [(hashed_idempotency_key, expected_sort_key)] + + +@pytest.mark.parametrize("source", ["condition", "get_item"]) +def test_stored_response_is_replayed_from_local_cache( + record_store, + source, + lambda_apigw_event, + lambda_context, + timestamp_future, + hashed_idempotency_key, + default_jmespath, +): + config = IdempotencyConfig(event_key_jmespath=default_jmespath, use_local_cache=True) + + @idempotent(config=config, persistence_store=record_store) + def lambda_handler(event, context): + pytest.fail("A stored response must bypass the handler") + + item = stored_item(record_store, hashed_idempotency_key, timestamp_future) + with Stubber(record_store.client) as stubber: + stub_existing_record(stubber, record_store, item, source) + assert lambda_handler(lambda_apigw_event, lambda_context) == {"message": "stored"} + # No more DynamoDB requests are stubbed: the second invocation must use the cache. + assert lambda_handler(lambda_apigw_event, lambda_context) == {"message": "stored"} + stubber.assert_no_pending_responses() + + +@pytest.mark.parametrize("source", ["condition", "get_item"]) +def test_in_progress_error_identifies_record_key( + record_store, + source, + lambda_apigw_event, + lambda_context, + timestamp_future, + hashed_idempotency_key, + default_jmespath, +): + config = IdempotencyConfig(event_key_jmespath=default_jmespath) + + @idempotent(config=config, persistence_store=record_store) + def lambda_handler(event, context): + pytest.fail("An in-progress invocation must bypass the handler") + + item = stored_item(record_store, hashed_idempotency_key, timestamp_future, status="INPROGRESS") + with Stubber(record_store.client) as stubber: + stub_existing_record(stubber, record_store, item, source) + with pytest.raises(IdempotencyAlreadyInProgressError) as exc: + lambda_handler(lambda_apigw_event, lambda_context) + assert f"{default_jmespath}={hashed_idempotency_key}" in str(exc.value) + stubber.assert_no_pending_responses()