From 28fc89aa14b6e17709ee1ff4655ff409772312b1 Mon Sep 17 00:00:00 2001 From: Sumedh Aerram Date: Sat, 31 Oct 2026 16:19:46 -0700 Subject: [PATCH] fix(idempotency): use the record key for DynamoDB composite items Read the idempotency key from the sort key for composite tables. Replayed records then use the actual operation key for response hooks, local cache lookups and in-progress diagnostics. Cover both conditional-write response and get-item paths through the public decorator, retaining simple-key controls. --- .../idempotency/persistence/dynamodb.py | 2 +- docs/utilities/idempotency.md | 2 + .../src/working_with_response_hook.py | 2 +- .../_boto3/test_composite_key_records.py | 134 ++++++++++++++++++ 4 files changed, 138 insertions(+), 2 deletions(-) create mode 100644 tests/functional/idempotency/_boto3/test_composite_key_records.py 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()