diff --git a/aws_lambda_powertools/utilities/idempotency/base.py b/aws_lambda_powertools/utilities/idempotency/base.py index eb4fe0eab5f..1ab1b9fa570 100644 --- a/aws_lambda_powertools/utilities/idempotency/base.py +++ b/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 diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/datarecord.py b/aws_lambda_powertools/utilities/idempotency/persistence/datarecord.py index 12e025b5e98..e9da1daf8eb 100644 --- a/aws_lambda_powertools/utilities/idempotency/persistence/datarecord.py +++ b/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,6 +45,8 @@ 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 @@ -51,6 +54,7 @@ def __init__( 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: diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py b/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py index 0b04ec135c7..18371a7d252 100644 --- a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py +++ b/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: diff --git a/tests/functional/idempotency/_boto3/test_idempotency.py b/tests/functional/idempotency/_boto3/test_idempotency.py index f92fb639350..17f14c2c182 100644 --- a/tests/functional/idempotency/_boto3/test_idempotency.py +++ b/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,