diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py b/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py index b05d8216b50..df356dd6ebb 100644 --- a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py +++ b/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py @@ -146,7 +146,6 @@ def _get_record(self, idempotency_key) -> DataRecord: def _put_record(self, data_record: DataRecord) -> None: item = { **self._get_key(data_record.idempotency_key), - self.key_attr: {"S": data_record.idempotency_key}, self.expiry_attr: {"N": str(data_record.expiry_timestamp)}, self.status_attr: {"S": data_record.status}, } diff --git a/tests/functional/idempotency/conftest.py b/tests/functional/idempotency/conftest.py index 7e5fa0e7c61..53813a8b49d 100644 --- a/tests/functional/idempotency/conftest.py +++ b/tests/functional/idempotency/conftest.py @@ -95,6 +95,16 @@ def expected_params_update_item(serialized_lambda_response, hashed_idempotency_k } +@pytest.fixture +def expected_params_update_item_compound_key_static_pk_value( + expected_params_update_item, hashed_idempotency_key, static_pk_value +): + return { + **expected_params_update_item, + "Key": {"id": {"S": static_pk_value}, "sk": {"S": hashed_idempotency_key}}, + } + + @pytest.fixture def expected_params_update_item_with_validation( serialized_lambda_response, hashed_idempotency_key, hashed_validation_key @@ -150,6 +160,22 @@ def expected_params_put_item(hashed_idempotency_key): } +@pytest.fixture +def expected_params_put_item_compound_key_static_pk_value( + expected_params_put_item, hashed_idempotency_key, static_pk_value +): + return { + **expected_params_put_item, + "Item": { + "expiration": {"N": stub.ANY}, + "in_progress_expiration": {"N": stub.ANY}, + "id": {"S": static_pk_value}, + "sk": {"S": hashed_idempotency_key}, + "status": {"S": "INPROGRESS"}, + }, + } + + @pytest.fixture def expected_params_put_item_with_validation(hashed_idempotency_key, hashed_validation_key): return { @@ -200,6 +226,11 @@ def hashed_idempotency_key_with_envelope(request, lambda_apigw_event): ) +@pytest.fixture +def static_pk_value(): + return "static-value" + + @pytest.fixture def hashed_validation_key(lambda_apigw_event): return hash_idempotency_key(lambda_apigw_event["requestContext"]) @@ -215,6 +246,13 @@ def persistence_store_compound(config): return DynamoDBPersistenceLayer(table_name=TABLE_NAME, boto_config=config, key_attr="id", sort_key_attr="sk") +@pytest.fixture +def persistence_store_compound_static_pk_value(config, static_pk_value): + return DynamoDBPersistenceLayer( + table_name=TABLE_NAME, boto_config=config, key_attr="id", sort_key_attr="sk", static_pk_value=static_pk_value + ) + + @pytest.fixture def idempotency_config(config, request, default_jmespath): return IdempotencyConfig( diff --git a/tests/functional/idempotency/test_idempotency.py b/tests/functional/idempotency/test_idempotency.py index dfc6b03b60c..68aeabeb50a 100644 --- a/tests/functional/idempotency/test_idempotency.py +++ b/tests/functional/idempotency/test_idempotency.py @@ -1504,3 +1504,34 @@ def lambda_handler(event, context): stubber.assert_no_pending_responses() stubber.deactivate() + + +@pytest.mark.parametrize("idempotency_config", [{"use_local_cache": False}], indirect=True) +def test_idempotent_lambda_compound_static_pk_value_has_correct_pk( + idempotency_config: IdempotencyConfig, + persistence_store_compound_static_pk_value: DynamoDBPersistenceLayer, + lambda_apigw_event, + expected_params_put_item_compound_key_static_pk_value, + expected_params_update_item_compound_key_static_pk_value, + lambda_response, + lambda_context, +): + """ + 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) + ddb_response = {} + + stubber.add_response("put_item", ddb_response, expected_params_put_item_compound_key_static_pk_value) + stubber.add_response("update_item", ddb_response, expected_params_update_item_compound_key_static_pk_value) + stubber.activate() + + @idempotent(config=idempotency_config, persistence_store=persistence_store_compound_static_pk_value) + def lambda_handler(event, context): + return lambda_response + + lambda_handler(lambda_apigw_event, lambda_context) + + stubber.assert_no_pending_responses() + stubber.deactivate()