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

Author: phipag

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 our 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;
+        }
     }
 }

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

@@ -14,12 +14,21 @@
 
 package software.amazon.lambda.powertools.idempotency.internal;
 
-import com.amazonaws.services.lambda.runtime.Context;
-import com.fasterxml.jackson.databind.JsonNode;
+import static software.amazon.lambda.powertools.idempotency.persistence.DataRecord.Status.EXPIRED;
+import static software.amazon.lambda.powertools.idempotency.persistence.DataRecord.Status.INPROGRESS;
+
+import java.time.Instant;
+import java.util.OptionalInt;
+import java.util.function.BiFunction;
+
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.reflect.MethodSignature;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.fasterxml.jackson.databind.JsonNode;
+
 import software.amazon.lambda.powertools.idempotency.Idempotency;
 import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyAlreadyInProgressException;
 import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyInconsistentStateException;
@@ -32,14 +41,9 @@
 import software.amazon.lambda.powertools.idempotency.persistence.DataRecord;
 import software.amazon.lambda.powertools.utilities.JsonConfig;
 
-import java.time.Instant;
-import java.util.OptionalInt;
-
-import static software.amazon.lambda.powertools.idempotency.persistence.DataRecord.Status.EXPIRED;
-import static software.amazon.lambda.powertools.idempotency.persistence.DataRecord.Status.INPROGRESS;
-
 /**
- * Internal class that will handle the Idempotency, and use the {@link software.amazon.lambda.powertools.idempotency.persistence.PersistenceStore}
+ * Internal class that will handle the Idempotency, and use the
+ * {@link software.amazon.lambda.powertools.idempotency.persistence.PersistenceStore}
  * to store the result of previous calls.
  */
 public class IdempotencyHandler {
@@ -106,7 +110,8 @@ private Object processIdempotency() throws Throwable {
 
     /**
      * Tries to determine the remaining time available for the current lambda invocation.
-     * Currently, it only works if the idempotent handler decorator is used or using {@link Idempotency#registerLambdaContext(Context)}
+     * Currently, it only works if the idempotent handler decorator is used or using
+     * {@link Idempotency#registerLambdaContext(Context)}
      *
      * @return the remaining time in milliseconds or empty if the context was not provided/found
      */
@@ -144,7 +149,8 @@ private DataRecord getIdempotencyRecord() {
     /**
      * Take appropriate action based on data_record's status
      *
-     * @param record DataRecord
+     * @param record
+     *            DataRecord
      * @return Function's response previously used for this idempotency key, if it has successfully executed already.
      */
     private Object handleForStatus(DataRecord record) {
@@ -167,10 +173,26 @@ private Object handleForStatus(DataRecord record) {
         try {
             LOG.debug("Response for key '{}' retrieved from idempotency store, skipping the function",
                     record.getIdempotencyKey());
+
+            final BiFunction<Object, DataRecord, Object> responseHook = Idempotency.getInstance().getConfig()
+                    .getResponseHook();
+            final Object responseData;
+
             if (returnType.equals(String.class)) {
-                return record.getResponseData();
+                // Primitive String data will be returned raw and not de-serialized from JSON.
+                responseData = record.getResponseData();
+            } else {
+                responseData = JsonConfig.get().getObjectMapper().reader().readValue(record.getResponseData(),
+                        returnType);
             }
-            return JsonConfig.get().getObjectMapper().reader().readValue(record.getResponseData(), returnType);
+
+            if (responseHook != null) {
+                LOG.debug(
+                        "Found user-defined idempotent response hook. Calling with response data and persistence store data record.");
+                return responseHook.apply(responseData, record);
+            }
+
+            return responseData;
         } catch (Exception e) {
             throw new IdempotencyPersistenceLayerException(
                     "Unable to get function response as " + returnType.getSimpleName(), e);