fix(idempotency): include sk in error msgs when using composite key (#6325)

Improve IdempotencyAlreadyInProgressError message

Author: leandrodamascena

aws_lambda_powertools/utilities/idempotency/base.py

@@ -250,10 +250,14 @@ def _handle_for_status(self, data_record: DataRecord) -> Any | None:
                     "item should have been expired in-progress because it already time-outed.",
                 )
 
-            raise IdempotencyAlreadyInProgressError(
+            inprogress_error_message = (
                 f"Execution already in progress with idempotency key: "
-                f"{self.persistence_store.event_key_jmespath}={data_record.idempotency_key}",
+                f"{self.persistence_store.event_key_jmespath}={data_record.idempotency_key}"
             )
+            if data_record.sort_key is not None:
+                inprogress_error_message += f" and sort key: {data_record.sort_key}"
+
+            raise IdempotencyAlreadyInProgressError(inprogress_error_message)
 
         response_dict = data_record.response_json_as_dict()
         serialized_response = self.output_serializer.from_dict(response_dict) if response_dict else None

aws_lambda_powertools/utilities/idempotency/persistence/datarecord.py

@@ -27,6 +27,7 @@ def __init__(
         in_progress_expiry_timestamp: int | None = None,
         response_data: str = "",
         payload_hash: str = "",
+        sort_key: str | None = None,
     ) -> None:
         """
 
@@ -44,13 +45,16 @@ def __init__(
             hashed representation of payload
         response_data: str, optional
             response data from previous executions using the record
+        sort_key: str, optional
+            sort key when using composite key
         """
         self.idempotency_key = idempotency_key
         self.payload_hash = payload_hash
         self.expiry_timestamp = expiry_timestamp
         self.in_progress_expiry_timestamp = in_progress_expiry_timestamp
         self._status = status
         self.response_data = response_data
+        self.sort_key = sort_key
 
     @property
     def is_expired(self) -> bool:

aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py

@@ -168,6 +168,7 @@ def _item_to_data_record(self, item: dict[str, Any]) -> DataRecord:
             in_progress_expiry_timestamp=data.get(self.in_progress_expiry_attr),
             response_data=data.get(self.data_attr),
             payload_hash=data.get(self.validation_key_attr),
+            sort_key=data[self.sort_key_attr] if self.sort_key_attr is not None else None,
         )
 
     def _get_record(self, idempotency_key) -> DataRecord:

tests/functional/idempotency/_boto3/test_idempotency.py

@@ -770,6 +770,54 @@ def lambda_handler(event, context):
     stubber.deactivate()
 
 
+@pytest.mark.parametrize("idempotency_config", [{"use_local_cache": False}, {"use_local_cache": True}], indirect=True)
+def test_idempotent_lambda_expires_in_progress_before_expire_with_sort_key(
+    idempotency_config: IdempotencyConfig,
+    persistence_store_compound_static_pk_value: DynamoDBPersistenceLayer,
+    lambda_apigw_event,
+    timestamp_future,
+    lambda_response,
+    hashed_idempotency_key,
+    lambda_context,
+):
+    stubber = stub.Stubber(persistence_store_compound_static_pk_value.client)
+
+    stubber.add_client_error("put_item", "ConditionalCheckFailedException")
+
+    now = datetime.datetime.now()
+    period = datetime.timedelta(seconds=5)
+    timestamp_expires_in_progress = int((now + period).timestamp() * 1000)
+
+    expected_params_get_item = {
+        "TableName": TABLE_NAME,
+        "Key": {"id": {"S": "static-value"}, "sk": {"S": hashed_idempotency_key}},
+        "ConsistentRead": True,
+    }
+    ddb_response_get_item = {
+        "Item": {
+            "id": {"S": "static-value"},
+            "expiration": {"N": timestamp_future},
+            "in_progress_expiration": {"N": str(timestamp_expires_in_progress)},
+            "data": {"S": '{"message": "test", "statusCode": 200'},
+            "status": {"S": "INPROGRESS"},
+            "sk": {"S": hashed_idempotency_key},
+        },
+    }
+    stubber.add_response("get_item", ddb_response_get_item, expected_params_get_item)
+
+    stubber.activate()
+
+    @idempotent(config=idempotency_config, persistence_store=persistence_store_compound_static_pk_value)
+    def lambda_handler(event, context):
+        return lambda_response
+
+    with pytest.raises(IdempotencyAlreadyInProgressError, match="and sort key"):
+        lambda_handler(lambda_apigw_event, lambda_context)
+
+    stubber.assert_no_pending_responses()
+    stubber.deactivate()
+
+
 @pytest.mark.parametrize("idempotency_config", [{"use_local_cache": False}, {"use_local_cache": True}], indirect=True)
 def test_idempotent_lambda_expires_in_progress_after_expire(
     idempotency_config: IdempotencyConfig,