fix(idempotency): use the record key for DynamoDB composite items

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),

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"

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

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()