feat(Idempotency): add feature for manipulating idempotent responses

aws_lambda_powertools/utilities/idempotency/__init__.py

@@ -2,7 +2,9 @@
 Utility for adding idempotency to lambda functions
 """
 
-from aws_lambda_powertools.utilities.idempotency.config import IdempotentHookData, IdempotentHookFunction
+from aws_lambda_powertools.utilities.idempotency.hook import (
+    IdempotentHookFunction,
+)
 from aws_lambda_powertools.utilities.idempotency.persistence.base import (
     BasePersistenceLayer,
 )
@@ -18,6 +20,5 @@
     "idempotent",
     "idempotent_function",
     "IdempotencyConfig",
-    "IdempotentHookData",
     "IdempotentHookFunction",
 )

aws_lambda_powertools/utilities/idempotency/base.py

@@ -5,7 +5,6 @@
 
 from aws_lambda_powertools.utilities.idempotency.config import (
     IdempotencyConfig,
-    IdempotentHookData,
 )
 from aws_lambda_powertools.utilities.idempotency.exceptions import (
     IdempotencyAlreadyInProgressError,
@@ -230,12 +229,15 @@ def _handle_for_status(self, data_record: DataRecord) -> Optional[Any]:
             )
         response_dict: Optional[dict] = data_record.response_json_as_dict()
         if response_dict is not None:
+            serialized_response = self.output_serializer.from_dict(response_dict)
             if self.config.response_hook is not None:
+                logger.debug("Response hook configured, invoking function")
                 return self.config.response_hook(
-                    self.output_serializer.from_dict(response_dict),
-                    IdempotentHookData(data_record),
+                    serialized_response,
+                    data_record,
                 )
-            return self.output_serializer.from_dict(response_dict)
+            return serialized_response
+
         return None
 
     def _get_function_response(self):

aws_lambda_powertools/utilities/idempotency/config.py

@@ -1,45 +1,9 @@
-from typing import Any, Dict, Optional
+from typing import Dict, Optional
 
-from aws_lambda_powertools.shared.types import Protocol
-from aws_lambda_powertools.utilities.idempotency.persistence.datarecord import DataRecord
+from aws_lambda_powertools.utilities.idempotency import IdempotentHookFunction
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 
-class IdempotentHookData:
-    """
-    Idempotent Hook Data
-
-    Contains data relevant to the current Idempotent record which matches the current request.
-    All IdempotentHook functions will be passed this data as well as the current Response.
-    """
-
-    def __init__(self, data_record: DataRecord) -> None:
-        self._idempotency_key = data_record.idempotency_key
-        self._status = data_record.status
-        self._expiry_timestamp = data_record.expiry_timestamp
-
-    @property
-    def idempotency_key(self) -> str:
-        return self._idempotency_key
-
-    @property
-    def status(self) -> str:
-        return self._status
-
-    @property
-    def expiry_timestamp(self) -> Optional[int]:
-        return self._expiry_timestamp
-
-
-class IdempotentHookFunction(Protocol):
-    """
-    The IdempotentHookFunction.
-    This class defines the calling signature for IdempotentHookFunction callbacks.
-    """
-
-    def __call__(self, response: Any, idempotent_data: IdempotentHookData) -> Any: ...
-
-
 class IdempotencyConfig:
     def __init__(
         self,

aws_lambda_powertools/utilities/idempotency/hook.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+from aws_lambda_powertools.shared.types import Protocol
+from aws_lambda_powertools.utilities.idempotency.persistence.datarecord import DataRecord
+
+
+class IdempotentHookFunction(Protocol):
+    """
+    The IdempotentHookFunction.
+    This class defines the calling signature for IdempotentHookFunction callbacks.
+    """
+
+    def __call__(self, response: Any, idempotent_data: DataRecord) -> Any: ...

docs/utilities/idempotency.md

@@ -454,6 +454,40 @@ sequenceDiagram
 <i>Idempotent successful request cached</i>
 </center>
 
+#### Successful request with response_hook configured
+
+<center>
+```mermaid
+sequenceDiagram
+    participant Client
+    participant Lambda
+    participant Response hook
+    participant Persistence Layer
+    alt initial request
+        Client->>Lambda: Invoke (event)
+        Lambda->>Persistence Layer: Get or set idempotency_key=hash(payload)
+        activate Persistence Layer
+        Note over Lambda,Persistence Layer: Set record status to INPROGRESS. <br> Prevents concurrent invocations <br> with the same payload
+        Lambda-->>Lambda: Call your function
+        Lambda->>Persistence Layer: Update record with result
+        deactivate Persistence Layer
+        Persistence Layer-->>Persistence Layer: Update record
+        Note over Lambda,Persistence Layer: Set record status to COMPLETE. <br> New invocations with the same payload <br> now return the same result
+        Lambda-->>Client: Response sent to client
+    else retried request
+        Client->>Lambda: Invoke (event)
+        Lambda->>Persistence Layer: Get or set idempotency_key=hash(payload)
+        activate Persistence Layer
+        Persistence Layer-->>Response hook: Already exists in persistence layer.
+        deactivate Persistence Layer
+        Note over Response hook,Persistence Layer: Record status is COMPLETE and not expired
+        Response hook->>Lambda: Response hook invoked
+        Lambda-->>Client: Same response sent to client
+    end
+```
+<i>Idempotent successful request with response hook</i>
+</center>
+
 #### Expired idempotency records
 
 <center>
@@ -699,16 +733,16 @@ For advanced configurations, such as setting up SSL certificates or customizing
 
 Idempotent decorator can be further configured with **`IdempotencyConfig`** as seen in the previous example. These are the available options for further configuration
 
-| Parameter                       | Default | Description                                                                                                                                                                                                                         |
-|---------------------------------|---------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| **event_key_jmespath**          | `""`    | JMESPath expression to extract the idempotency key from the event record using [built-in functions](./jmespath_functions.md#built-in-jmespath-functions){target="_blank"}                                                           |
-| **payload_validation_jmespath** | `""`    | JMESPath expression to validate whether certain parameters have changed in the event while the event payload                                                                                                                        |
-| **raise_on_no_idempotency_key** | `False` | Raise exception if no idempotency key was found in the request                                                                                                                                                                      |
-| **expires_after_seconds**       | 3600    | The number of seconds to wait before a record is expired                                                                                                                                                                            |
-| **use_local_cache**             | `False` | Whether to locally cache idempotency results                                                                                                                                                                                        |
-| **local_cache_max_items**       | 256     | Max number of items to store in local cache                                                                                                                                                                                         |
-| **hash_function**               | `md5`   | Function to use for calculating hashes, as provided by [hashlib](https://docs.python.org/3/library/hashlib.html){target="_blank" rel="nofollow"} in the standard library.                                                           |
-| **response_hook**               | `None`  | Function to use for processing the stored Idempotent response.  This function hook is called when an already returned response is found.  See [Modifying The Idempotent Response](idempotency.md#modifying-the-idempotent-repsonse) |
+| Parameter                       | Default | Description                                                                                                                                                                                                                                  |
+|---------------------------------|---------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| **event_key_jmespath**          | `""`    | JMESPath expression to extract the idempotency key from the event record using [built-in functions](./jmespath_functions.md#built-in-jmespath-functions){target="_blank"}                                                                    |
+| **payload_validation_jmespath** | `""`    | JMESPath expression to validate whether certain parameters have changed in the event while the event payload                                                                                                                                 |
+| **raise_on_no_idempotency_key** | `False` | Raise exception if no idempotency key was found in the request                                                                                                                                                                               |
+| **expires_after_seconds**       | 3600    | The number of seconds to wait before a record is expired                                                                                                                                                                                     |
+| **use_local_cache**             | `False` | Whether to locally cache idempotency results                                                                                                                                                                                                 |
+| **local_cache_max_items**       | 256     | Max number of items to store in local cache                                                                                                                                                                                                  |
+| **hash_function**               | `md5`   | Function to use for calculating hashes, as provided by [hashlib](https://docs.python.org/3/library/hashlib.html){target="_blank" rel="nofollow"} in the standard library.                                                                    |
+| **response_hook**               | `None`  | Function to use for processing the stored Idempotent response. This function hook is called when an existing idempotent response is found. See [Manipulating The Idempotent Response](idempotency.md#manipulating-the-idempotent-response) |
 
 ### Handling concurrent executions with the same payload
 
@@ -910,22 +944,38 @@ You can create your own persistent store from scratch by inheriting the `BasePer
 
     For example, the `_put_record` method needs to raise an exception if a non-expired record already exists in the data store with a matching key.
 
-### Modifying the Idempotent Repsonse
-
-The IdempotentConfig allows you to specify a _**response_hook**_ which is a function that will be called when an idempotent response is loaded from the PersistenceStore.
+### Manipulating the Idempotent Response
 
-You can provide the response_hook using _**IdempotentConfig**_.
+You can set up a `response_hook` in the `IdempotentConfig` class to access the returned data when an operation is idempotent. The hook function will be called with the current deserialized response object and the Idempotency record.
 
 === "Using an Idempotent Response Hook"
 
-```python hl_lines="10-15 19"
---8<-- "examples/idempotency/src/working_with_response_hook.py"
-```
+    ```python hl_lines="18 20 23 32"
+    --8<-- "examples/idempotency/src/working_with_response_hook.py"
+    ```
+
+=== "Sample event"
+
+    ```json
+    --8<-- "examples/idempotency/src/working_with_response_hook_payload.json"
+    ```
 
 ???+ info "Info: Using custom de-serialization?"
 
     The response_hook is called after the custom de-serialization so the payload you process will be the de-serialized version.
 
+#### Being a good citizen
+
+Using Response hooks can add subtle improvements to manipulating returned data from idempotent operations, but also add significant complexity if you're not careful.
+
+Keep the following in mind when authoring hooks for Idempotency utility:
+
+1. **Response hook works exclusively when operations are idempotent.** Carefully consider the logic within the `Response hook` and prevent any attempt to access the key from relying exclusively on idempotent operations.
+
+2. **Catch your own exceptions.** Catch and handle known exceptions to your logic.
+
+3. **Watch out when you are decorating the Lambda Handler and using the Response hook.** If you don't catch and handle exceptions in your `Response hook`, your function might not run properly.
+
 ## Compatibility with other utilities
 
 ### Batch

examples/idempotency/src/working_with_response_hook.py

@@ -1,27 +1,56 @@
+import datetime
+import uuid
 from typing import Dict
 
+from aws_lambda_powertools import Logger
 from aws_lambda_powertools.utilities.idempotency import (
     DynamoDBPersistenceLayer,
     IdempotencyConfig,
-    IdempotentHookData,
-    idempotent,
+    idempotent_function,
+)
+from aws_lambda_powertools.utilities.idempotency.persistence.base import (
+    DataRecord,
 )
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
+logger = Logger()
+
+
+def my_response_hook(response: Dict, idempotent_data: DataRecord) -> Dict:
+    # Return inserted Header data into the Idempotent Response
+    response["x-idempotent-key"] = idempotent_data.idempotency_key
 
-def my_response_hook(response: Dict, idempotent_data: IdempotentHookData) -> Dict:
-    # How to add a field to the response
-    response["is_idempotent_response"] = True
+    # expiry_timestamp could be None so include if set
+    expiry_timestamp = idempotent_data.expiry_timestamp
+    if expiry_timestamp:
+        expiry_time = datetime.datetime.fromtimestamp(int(expiry_timestamp))
+        response["x-idempotent-expiration"] = expiry_time.isoformat()
 
     # Must return the response here
     return response
 
 
-persistence_layer = DynamoDBPersistenceLayer(table_name="IdempotencyTable")
+dynamodb = DynamoDBPersistenceLayer(table_name="IdempotencyTable")
+config = IdempotencyConfig(response_hook=my_response_hook)
+
+
+@idempotent_function(data_keyword_argument="order", config=config, persistence_store=dynamodb)
+def process_order(order: dict) -> dict:
+    # create the order_id
+    order_id = str(uuid.uuid4())
+
+    # create your logic to save the order
+    # append the order_id created
+    order["order_id"] = order_id
 
-config = IdempotencyConfig(event_key_jmespath="body", response_hook=my_response_hook)
+    # return the order
+    return {"order": order}
 
 
-@idempotent(persistence_store=persistence_layer, config=config)
-def lambda_handler(event: dict, context: LambdaContext) -> dict:
-    return event
+def lambda_handler(event: dict, context: LambdaContext):
+    config.register_lambda_context(context)  # see Lambda timeouts section
+    try:
+        logger.info(f"Processing order id {event.get('order_id')}")
+        return process_order(order=event.get("order"))
+    except Exception as err:
+        return {"status_code": 400, "error": f"Erro processing {str(err)}"}

examples/idempotency/src/working_with_response_hook_payload.json

@@ -0,0 +1,8 @@
+{
+  "order" : {
+    "user_id": "xyz",
+    "product_id": "123456789",
+    "quantity": 2,
+    "value": 30
+  }
+}

tests/functional/idempotency/test_idempotency.py

@@ -27,7 +27,6 @@
     IdempotencyHandler,
     _prepare_data,
 )
-from aws_lambda_powertools.utilities.idempotency.config import IdempotentHookData
 from aws_lambda_powertools.utilities.idempotency.exceptions import (
     IdempotencyAlreadyInProgressError,
     IdempotencyInconsistentStateError,
@@ -2043,7 +2042,7 @@ def test_idempotent_lambda_already_completed_response_hook_is_called(
     Test idempotent decorator where event with matching event key has already been successfully processed
     """
 
-    def idempotent_response_hook(response: Any, idempotent_data: IdempotentHookData) -> Any:
+    def idempotent_response_hook(response: Any, idempotent_data: DataRecord) -> Any:
         """Modify the response provided by adding a new key"""
         response["idempotent_response"] = True