diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py b/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py index b05d8216b50..ce3aac2425e 100644 --- a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py +++ b/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py @@ -104,6 +104,21 @@ def __init__( super(DynamoDBPersistenceLayer, self).__init__() def _get_key(self, idempotency_key: str) -> dict: + """Build primary key attribute simple or composite based on params. + + When sort_key_attr is set, we must return a composite key with static_pk_value, + otherwise we use the idempotency key given. + + Parameters + ---------- + idempotency_key : str + idempotency key to use for simple primary key + + Returns + ------- + dict + simple or composite key for DynamoDB primary key + """ if self.sort_key_attr: return {self.key_attr: {"S": self.static_pk_value}, self.sort_key_attr: {"S": idempotency_key}} return {self.key_attr: {"S": idempotency_key}} @@ -145,8 +160,8 @@ def _get_record(self, idempotency_key) -> DataRecord: def _put_record(self, data_record: DataRecord) -> None: item = { + # get simple or composite primary key **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..75834f76688 100644 --- a/tests/functional/idempotency/conftest.py +++ b/tests/functional/idempotency/conftest.py @@ -215,6 +215,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( @@ -246,3 +253,37 @@ def _func_echo_decoder(self, value): @pytest.fixture def mock_function(): return mock.MagicMock() + + +@pytest.fixture +def static_pk_value(): + return "static-value" + + +@pytest.fixture +def expected_params_update_item_compound_key_static_pk_value( + expected_params_update_item, hashed_idempotency_key, static_pk_value +): + return { + # same as in any update_item transaction except the `Key` due to composite key value + **expected_params_update_item, + "Key": {"id": {"S": static_pk_value}, "sk": {"S": 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 { + # same as in any put_item transaction except the `Item` due to composite key value + **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"}, + }, + "TableName": "TEST_TABLE", + } 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()