diff --git a/aws_lambda_powertools/utilities/idempotency/idempotency.py b/aws_lambda_powertools/utilities/idempotency/idempotency.py index b77c3013cbb..6f73a842af4 100644 --- a/aws_lambda_powertools/utilities/idempotency/idempotency.py +++ b/aws_lambda_powertools/utilities/idempotency/idempotency.py @@ -11,6 +11,7 @@ IdempotencyInconsistentStateError, IdempotencyItemAlreadyExistsError, IdempotencyItemNotFoundError, + IdempotencyKeyError, IdempotencyPersistenceLayerError, IdempotencyValidationError, ) @@ -132,6 +133,8 @@ def handle(self) -> Any: # We call save_inprogress first as an optimization for the most common case where no idempotent record # already exists. If it succeeds, there's no need to call get_record. self.persistence_store.save_inprogress(event=self.event, context=self.context) + except IdempotencyKeyError: + raise except IdempotencyItemAlreadyExistsError: # Now we know the item already exists, we can retrieve it record = self._get_idempotency_record() diff --git a/tests/functional/idempotency/test_idempotency.py b/tests/functional/idempotency/test_idempotency.py index 503ec7d6183..25f76af48be 100644 --- a/tests/functional/idempotency/test_idempotency.py +++ b/tests/functional/idempotency/test_idempotency.py @@ -828,3 +828,20 @@ def lambda_handler(event, context): stubber.assert_no_pending_responses() stubber.deactivate() assert "Failed to save in progress record to idempotency store" == e.value.args[0] + + +def test_handler_raise_idempotency_key_error(persistence_store: DynamoDBPersistenceLayer, lambda_context): + # GIVEN raise_on_no_idempotency_key is True + idempotency_config = IdempotencyConfig(event_key_jmespath="idemKey", raise_on_no_idempotency_key=True) + + # WHEN handling the idempotent call + # AND save_inprogress raises a IdempotencyKeyError + @idempotent(persistence_store=persistence_store, config=idempotency_config) + def handler(event, context): + raise ValueError("Should not be raised") + + # THEN idempotent should re-raise the IdempotencyKeyError + with pytest.raises(IdempotencyKeyError) as e: + handler({}, lambda_context) + + assert "No data found to create a hashed idempotency_key" == e.value.args[0]