feat(idempotency): Add response hook feature (#1814)

* Add responseHook definition to IdempotencyConfig and call from IdempotencyHandler. Update idempotency example with an example use-case modifying API GW headers.

* Make debug log message more concise in IdempotencyHandler.java.

* Add unit test for idempotency response hook.

* Document idempotency response hook feature.

* Update examples/powertools-examples-idempotency/src/main/java/helloworld/App.java

Co-authored-by: Leandro Damascena <[email protected]>

---------

Co-authored-by: Leandro Damascena <[email protected]>

Author: phipag

docs/utilities/idempotency.md

@@ -55,7 +55,7 @@ times with the same parameters**. This makes idempotent operations safe to retry
                      <aspectLibraries>
                          <aspectLibrary>
                              <groupId>software.amazon.lambda</groupId>
-                             <artifactId>powertools-idempotency-dynamodb</artifactId>
+                             <artifactId>powertools-idempotency-core</artifactId>
                          </aspectLibrary>
                      </aspectLibraries>
                  </configuration>
@@ -584,20 +584,22 @@ IdempotencyConfig.builder()
                 .withUseLocalCache(true)
                 .withLocalCacheMaxItems(432)
                 .withHashFunction("SHA-256")
+                .withResponseHook((responseData, dataRecord) -> responseData)
                 .build()
 ```
 
 These are the available options for further configuration:
 
 | Parameter                                         | Default | Description                                                                                                                      |
 |---------------------------------------------------|---------|----------------------------------------------------------------------------------------------------------------------------------|
-| **EventKeyJMESPath**                              | `""`    | JMESPath expression to extract the idempotency key from the event record. See available [built-in functions](serialization) |
+| **EventKeyJMESPath**                              | `""`    | JMESPath expression to extract the idempotency key from the event record. See available [built-in functions](serialization)      |
 | **PayloadValidationJMESPath**                     | `""`    | JMESPath expression to validate whether certain parameters have changed in the event                                             |
-| **ThrowOnNoIdempotencyKey**                       | `False` | Throw exception if no idempotency key was found in the request                                                                   |
+| **ThrowOnNoIdempotencyKey**                       | `false` | Throw exception if no idempotency key was found in the request                                                                   |
 | **ExpirationInSeconds**                           | 3600    | The number of seconds to wait before a record is expired                                                                         |
 | **UseLocalCache**                                 | `false` | Whether to locally cache idempotency results (LRU cache)                                                                         |
 | **LocalCacheMaxItems**                            | 256     | Max number of items to store in local cache                                                                                      |
 | **HashFunction**                                  | `MD5`   | Algorithm to use for calculating hashes, as supported by `java.security.MessageDigest` (eg. SHA-1, SHA-256, ...)                 |
+| **ResponseHook**                                  | `null`  | Response hook to apply modifications to idempotent responses                                                                     |
 
 These features are detailed below.
 
@@ -855,6 +857,58 @@ You can extend the `BasePersistenceStore` class and implement the abstract metho
 
     For example, the `putRecord` method needs to throw an exception if a non-expired record already exists in the data store with a matching key.
 
+### Manipulating the Idempotent Response
+
+You can set up a response hook in the Idempotency configuration to manipulate the returned data when an operation is idempotent. The hook function will be called with the current de-serialized response `Object` and the Idempotency `DataRecord`.
+
+The example below shows how to append an HTTP header to an `APIGatewayProxyResponseEvent`.
+
+=== "Using an Idempotent Response Hook"
+
+    ```java hl_lines="3-20"
+    Idempotency.config().withConfig(
+            IdempotencyConfig.builder()
+                    .withResponseHook((responseData, dataRecord) -> {
+                        if (responseData instanceof APIGatewayProxyResponseEvent) {
+                            APIGatewayProxyResponseEvent proxyResponse = 
+                                (APIGatewayProxyResponseEvent) responseData;
+                            final Map<String, String> headers = new HashMap<>();
+                            headers.putAll(proxyResponse.getHeaders());
+                            // Append idempotency headers
+                            headers.put("x-idempotency-response", "true");
+                            headers.put("x-idempotency-expiration",
+                                    String.valueOf(dataRecord.getExpiryTimestamp()));
+
+                            proxyResponse.setHeaders(headers);
+
+                            return proxyResponse;
+                        }
+
+                        return responseData;
+                    })
+                    .build())
+            .withPersistenceStore(
+                    DynamoDBPersistenceStore.builder()
+                            .withTableName(System.getenv("IDEMPOTENCY_TABLE"))
+                            .build())
+            .configure();
+    ```
+
+???+ info "Info: Using custom de-serialization?"
+
+    The response hook is called after de-serialization so the payload you process will be the de-serialized Java object.
+
+#### Being a good citizen
+
+When using response hooks to manipulate returned data from idempotent operations, it's important to follow best practices to avoid introducing complexity or issues. Keep these guidelines in mind:
+
+1. **Response hook works exclusively when operations are idempotent.** The hook will not be called when an operation is not idempotent, or when the idempotent logic fails.
+
+2. **Catch and Handle Exceptions.** Your response hook code should catch and handle any exceptions that may arise from your logic. Unhandled exceptions will cause the Lambda function to fail unexpectedly.
+
+3. **Keep Hook Logic Simple** Response hooks should consist of minimal and straightforward logic for manipulating response data. Avoid complex conditional branching and aim for hooks that are easy to reason about.
+
+
 ## Compatibility with other utilities
 
 ### Validation utility

examples/powertools-examples-idempotency/src/main/java/helloworld/App.java

@@ -45,20 +45,36 @@ public App() {
 
     public App(DynamoDbClient client) {
         Idempotency.config().withConfig(
-                        IdempotencyConfig.builder()
-                                .withEventKeyJMESPath(
-                                        "powertools_json(body).address") // will retrieve the address field in the body which is a string transformed to json with `powertools_json`
-                                .build())
+                IdempotencyConfig.builder()
+                        .withEventKeyJMESPath("powertools_json(body).address")
+                        .withResponseHook((responseData, dataRecord) -> {
+                            if (responseData instanceof APIGatewayProxyResponseEvent) {
+                                APIGatewayProxyResponseEvent proxyResponse = (APIGatewayProxyResponseEvent) responseData;
+                                final Map<String, String> headers = new HashMap<>();
+                                headers.putAll(proxyResponse.getHeaders());
+                                headers.put("x-idempotency-response", "true");
+                                headers.put("x-idempotency-expiration",
+                                        String.valueOf(dataRecord.getExpiryTimestamp()));
+
+                                proxyResponse.setHeaders(headers);
+
+                                return proxyResponse;
+                            }
+
+                            return responseData;
+                        })
+                        .build())
                 .withPersistenceStore(
                         DynamoDBPersistenceStore.builder()
                                 .withDynamoDbClient(client)
                                 .withTableName(System.getenv("IDEMPOTENCY_TABLE"))
-                                .build()
-                ).configure();
+                                .build())
+                .configure();
     }
 
     /**
-     * This is our Lambda event handler. It accepts HTTP POST  requests from API gateway and returns the contents of the given URL. Requests are made idempotent
+     * This is your Lambda event handler. It accepts HTTP POST requests from API gateway and returns the contents of the
+     * given URL. Requests are made idempotent
      * by the idempotency library, and results are cached for the default 1h expiry time.
      * <p>
      * You can test the endpoint like this:
@@ -67,8 +83,10 @@ public App(DynamoDbClient client) {
      *     curl -X POST https://[REST-API-ID].execute-api.[REGION].amazonaws.com/Prod/helloidem/ -H "Content-Type: application/json" -d '{"address": "https://checkip.amazonaws.com"}'
      * </pre>
      * <ul>
-     *     <li>First call will execute the handleRequest normally, and store the response in the idempotency table (Look into DynamoDB)</li>
-     *     <li>Second call (and next ones) will retrieve from the cache (if cache is enabled, which is by default) or from the store, the handler won't be called. Until the expiration happens (by default 1 hour).</li>
+     * <li>First call will execute the handleRequest normally, and store the response in the idempotency table (Look
+     * into DynamoDB)</li>
+     * <li>Second call (and next ones) will retrieve from the cache (if cache is enabled, which is by default) or from
+     * the store, the handler won't be called. Until the expiration happens (by default 1 hour).</li>
      * </ul>
      */
     @Idempotent // The magic is here!
@@ -101,14 +119,14 @@ public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEv
         }
     }
 
-
     /**
      * Helper to retrieve the contents of the given URL and return them as a string.
      * <p>
      * We could also put the @Idempotent annotation here if we only wanted this sub-operation to be idempotent. Putting
      * it on the handler, however, reduces total execution time and saves us time!
      *
-     * @param address The URL to fetch
+     * @param address
+     *            The URL to fetch
      * @return The contents of the given URL
      * @throws IOException
      */
@@ -118,4 +136,4 @@ private String getPageContents(String address) throws IOException {
             return br.lines().collect(Collectors.joining(System.lineSeparator()));
         }
     }
-}
+}

powertools-idempotency/powertools-idempotency-core/src/main/java/software/amazon/lambda/powertools/idempotency/IdempotencyConfig.java

@@ -14,10 +14,13 @@
 
 package software.amazon.lambda.powertools.idempotency;
 
+import java.time.Duration;
+import java.util.function.BiFunction;
+
 import com.amazonaws.services.lambda.runtime.Context;
-import software.amazon.lambda.powertools.idempotency.internal.cache.LRUCache;
 
-import java.time.Duration;
+import software.amazon.lambda.powertools.idempotency.internal.cache.LRUCache;
+import software.amazon.lambda.powertools.idempotency.persistence.DataRecord;
 
 /**
  * Configuration of the idempotency feature. Use the {@link Builder} to create an instance.
@@ -30,18 +33,20 @@ public class IdempotencyConfig {
     private final String payloadValidationJMESPath;
     private final boolean throwOnNoIdempotencyKey;
     private final String hashFunction;
+    private final BiFunction<Object, DataRecord, Object> responseHook;
     private Context lambdaContext;
 
     private IdempotencyConfig(String eventKeyJMESPath, String payloadValidationJMESPath,
-                              boolean throwOnNoIdempotencyKey, boolean useLocalCache, int localCacheMaxItems,
-                              long expirationInSeconds, String hashFunction) {
+            boolean throwOnNoIdempotencyKey, boolean useLocalCache, int localCacheMaxItems,
+            long expirationInSeconds, String hashFunction, BiFunction<Object, DataRecord, Object> responseHook) {
         this.localCacheMaxItems = localCacheMaxItems;
         this.useLocalCache = useLocalCache;
         this.expirationInSeconds = expirationInSeconds;
         this.eventKeyJMESPath = eventKeyJMESPath;
         this.payloadValidationJMESPath = payloadValidationJMESPath;
         this.throwOnNoIdempotencyKey = throwOnNoIdempotencyKey;
         this.hashFunction = hashFunction;
+        this.responseHook = responseHook;
     }
 
     /**
@@ -89,6 +94,10 @@ public void setLambdaContext(Context lambdaContext) {
         this.lambdaContext = lambdaContext;
     }
 
+    public BiFunction<Object, DataRecord, Object> getResponseHook() {
+        return responseHook;
+    }
+
     public static class Builder {
 
         private int localCacheMaxItems = 256;
@@ -98,14 +107,18 @@ public static class Builder {
         private String payloadValidationJMESPath;
         private boolean throwOnNoIdempotencyKey = false;
         private String hashFunction = "MD5";
+        private BiFunction<Object, DataRecord, Object> responseHook;
 
         /**
          * Initialize and return an instance of {@link IdempotencyConfig}.<br>
          * Example:<br>
+         * 
          * <pre>
          * IdempotencyConfig.builder().withUseLocalCache().build();
          * </pre>
+         * 
          * This instance must then be passed to the {@link Idempotency.Config}:
+         * 
          * <pre>
          * Idempotency.config().withConfig(config).configure();
          * </pre>
@@ -120,13 +133,15 @@ public IdempotencyConfig build() {
                     useLocalCache,
                     localCacheMaxItems,
                     expirationInSeconds,
-                    hashFunction);
+                    hashFunction,
+                    responseHook);
         }
 
         /**
          * A JMESPath expression to extract the idempotency key from the event record. <br>
          * See <a href="https://jmespath.org/">https://jmespath.org/</a> for more details.<br>
-         * Common paths are: <ul>
+         * Common paths are:
+         * <ul>
          * <li><code>powertools_json(body)</code> for APIGatewayProxyRequestEvent and APIGatewayV2HTTPEvent</li>
          * <li><code>Records[*].powertools_json(body)</code> for SQSEvent</li>
          * <li><code>Records[0].Sns.Message | powertools_json(@)</code> for SNSEvent</li>
@@ -136,7 +151,8 @@ public IdempotencyConfig build() {
          * <li>...</li>
          * </ul>
          *
-         * @param eventKeyJMESPath path of the key in the Lambda event
+         * @param eventKeyJMESPath
+         *            path of the key in the Lambda event
          * @return the instance of the builder (to chain operations)
          */
         public Builder withEventKeyJMESPath(String eventKeyJMESPath) {
@@ -147,7 +163,8 @@ public Builder withEventKeyJMESPath(String eventKeyJMESPath) {
         /**
          * Set the maximum number of items to store in local cache, by default 256
          *
-         * @param localCacheMaxItems maximum number of items to store in local cache
+         * @param localCacheMaxItems
+         *            maximum number of items to store in local cache
          * @return the instance of the builder (to chain operations)
          */
         public Builder withLocalCacheMaxItems(int localCacheMaxItems) {
@@ -158,8 +175,9 @@ public Builder withLocalCacheMaxItems(int localCacheMaxItems) {
         /**
          * Whether to locally cache idempotency results, by default false
          *
-         * @param useLocalCache boolean that indicate if a local cache must be used in addition to the persistence store.
-         *                      If set to true, will use the {@link LRUCache}
+         * @param useLocalCache
+         *            boolean that indicate if a local cache must be used in addition to the persistence store.
+         *            If set to true, will use the {@link LRUCache}
          * @return the instance of the builder (to chain operations)
          */
         public Builder withUseLocalCache(boolean useLocalCache) {
@@ -170,7 +188,8 @@ public Builder withUseLocalCache(boolean useLocalCache) {
         /**
          * The number of seconds to wait before a record is expired
          *
-         * @param expiration expiration of the record in the store
+         * @param expiration
+         *            expiration of the record in the store
          * @return the instance of the builder (to chain operations)
          */
         public Builder withExpiration(Duration expiration) {
@@ -182,7 +201,8 @@ public Builder withExpiration(Duration expiration) {
          * A JMESPath expression to extract the payload to be validated from the event record. <br/>
          * See <a href="https://jmespath.org/">https://jmespath.org/</a> for more details.
          *
-         * @param payloadValidationJMESPath JMES Path of a part of the payload to be used for validation
+         * @param payloadValidationJMESPath
+         *            JMES Path of a part of the payload to be used for validation
          * @return the instance of the builder (to chain operations)
          */
         public Builder withPayloadValidationJMESPath(String payloadValidationJMESPath) {
@@ -193,8 +213,9 @@ public Builder withPayloadValidationJMESPath(String payloadValidationJMESPath) {
         /**
          * Whether to throw an exception if no idempotency key was found in the request, by default false
          *
-         * @param throwOnNoIdempotencyKey boolean to indicate if we must throw an Exception when
-         *                                idempotency key could not be found in the payload.
+         * @param throwOnNoIdempotencyKey
+         *            boolean to indicate if we must throw an Exception when idempotency key could not be found in the
+         *            payload.
          * @return the instance of the builder (to chain operations)
          */
         public Builder withThrowOnNoIdempotencyKey(boolean throwOnNoIdempotencyKey) {
@@ -215,15 +236,43 @@ public Builder withThrowOnNoIdempotencyKey() {
         /**
          * Function to use for calculating hashes, by default MD5.
          *
-         * @param hashFunction Can be any algorithm supported by {@link java.security.MessageDigest}, most commons are<ul>
-         *                     <li>MD5</li>
-         *                     <li>SHA-1</li>
-         *                     <li>SHA-256</li></ul>
+         * @param hashFunction
+         *            Can be any algorithm supported by {@link java.security.MessageDigest}, most commons are
+         *            <ul>
+         *            <li>MD5</li>
+         *            <li>SHA-1</li>
+         *            <li>SHA-256</li>
+         *            </ul>
          * @return the instance of the builder (to chain operations)
          */
         public Builder withHashFunction(String hashFunction) {
             this.hashFunction = hashFunction;
             return this;
         }
+
+        /**
+         * Response hook that will be called for each idempotent response. This hook will receive the de-serialized
+         * response data from the persistence store as first argument and the original DataRecord from the persistence
+         * store as second argument.
+         * 
+         * Usage:
+         * 
+         * <pre>
+         * IdempotencyConfig.builder().withResponseHook((responseData, dataRecord) -> {
+         *     // do something with the response data, for example:
+         *     if(responseData instanceof APIGatewayProxyRequestEvent) {
+         *         ((APIGatewayProxyRequestEvent) responseData).setHeaders(Map.of("x-idempotency-response", "true")
+         *     }
+         *     return responseData;
+         * })
+         * </pre>
+         * 
+         * @param responseHook
+         * @return
+         */
+        public Builder withResponseHook(BiFunction<Object, DataRecord, Object> responseHook) {
+            this.responseHook = responseHook;
+            return this;
+        }
     }
 }