fix(idempotency): revert dict mutation that impacted static_pk_value feature (#1970)

Co-authored-by:  Adam Tankanow <[email protected]>

Author: heitorlessa

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

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",
+    }

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