Merge pull request #968 from aws-powertools/feature/idempotency-hooks

feat: Add support for response hooks in idempotency utility

Author: hjgraca

docs/utilities/idempotency.md

@@ -846,6 +846,50 @@ Data would then be stored in DynamoDB like this:
 | idempotency#MyLambdaFunction | f091d2527ad1c78f05d54cc3f363be80 | 1636549585 | IN_PROGRESS |                                      |
 
 
+### 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 deserialized response object and the Idempotency `DataRecord`.
+
+#### Using Response Hooks
+
+The example below shows how to append HTTP headers to an `APIGatewayProxyResponse`:
+
+```csharp
+Idempotency.Config()
+    .WithConfig(IdempotencyOptions.Builder()
+        .WithEventKeyJmesPath("powertools_json(body).address")
+        .WithResponseHook((responseData, dataRecord) => {
+            if (responseData is APIGatewayProxyResponse proxyResponse)
+            {
+                proxyResponse.Headers ??= new Dictionary<string, string>();
+                proxyResponse.Headers["x-idempotency-response"] = "true";
+                proxyResponse.Headers["x-idempotency-expiration"] = dataRecord.ExpiryTimestamp.ToString();
+                return proxyResponse;
+            }
+            return responseData;
+        })
+        .Build())
+    .WithPersistenceStore(DynamoDBPersistenceStore.Builder()
+        .WithTableName(Environment.GetEnvironmentVariable("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 C# 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.
+
+
 ## AOT Support
 
 Native AOT trims your application code as part of the compilation to ensure that the binary is as small as possible. .NET 8 for Lambda provides improved trimming support compared to previous versions of .NET.

libraries/src/AWS.Lambda.Powertools.Idempotency/IdempotencyOptions.cs

@@ -13,6 +13,8 @@
  * permissions and limitations under the License.
  */
 
+using System;
+
 namespace AWS.Lambda.Powertools.Idempotency;
 
 /// <summary>
@@ -57,32 +59,24 @@ public class IdempotencyOptions
     /// as supported by <see cref="System.Security.Cryptography.HashAlgorithm"/> (eg. SHA1, SHA-256, ...)
     /// </summary>
     public string HashFunction { get; }
+    /// <summary>
+    /// Delegate for manipulating idempotent responses.
+    /// </summary>
+    public Func<object, Persistence.DataRecord, object> ResponseHook { get; }
 
     /// <summary>
     /// Constructor of <see cref="IdempotencyOptions"/>.
     /// </summary>
-    /// <param name="eventKeyJmesPath"></param>
-    /// <param name="payloadValidationJmesPath"></param>
-    /// <param name="throwOnNoIdempotencyKey"></param>
-    /// <param name="useLocalCache"></param>
-    /// <param name="localCacheMaxItems"></param>
-    /// <param name="expirationInSeconds"></param>
-    /// <param name="hashFunction"></param>
-    internal IdempotencyOptions(
-        string eventKeyJmesPath, 
-        string payloadValidationJmesPath, 
-        bool throwOnNoIdempotencyKey, 
-        bool useLocalCache, 
-        int localCacheMaxItems, 
-        long expirationInSeconds, 
-        string hashFunction)
+    /// <param name="builder">The builder containing the configuration values</param>
+    internal IdempotencyOptions(IdempotencyOptionsBuilder builder)
     {
-        EventKeyJmesPath = eventKeyJmesPath;
-        PayloadValidationJmesPath = payloadValidationJmesPath;
-        ThrowOnNoIdempotencyKey = throwOnNoIdempotencyKey;
-        UseLocalCache = useLocalCache;
-        LocalCacheMaxItems = localCacheMaxItems;
-        ExpirationInSeconds = expirationInSeconds;
-        HashFunction = hashFunction;
+        EventKeyJmesPath = builder.EventKeyJmesPath;
+        PayloadValidationJmesPath = builder.PayloadValidationJmesPath;
+        ThrowOnNoIdempotencyKey = builder.ThrowOnNoIdempotencyKey;
+        UseLocalCache = builder.UseLocalCache;
+        LocalCacheMaxItems = builder.LocalCacheMaxItems;
+        ExpirationInSeconds = builder.ExpirationInSeconds;
+        HashFunction = builder.HashFunction;
+        ResponseHook = builder.ResponseHook;
     }
 }

libraries/src/AWS.Lambda.Powertools.Idempotency/IdempotencyOptionsBuilder.cs

@@ -43,21 +43,59 @@ public class IdempotencyOptionsBuilder
     private string _hashFunction = "MD5";
 
     /// <summary>
-    /// Initialize and return an instance of IdempotencyConfig.
+    /// Response hook function
+    /// </summary>
+    private Func<object, AWS.Lambda.Powertools.Idempotency.Persistence.DataRecord, object> _responseHook;
+
+    /// <summary>
+    /// Gets the event key JMESPath expression.
+    /// </summary>
+    internal string EventKeyJmesPath => _eventKeyJmesPath;
+
+    /// <summary>
+    /// Gets the payload validation JMESPath expression.
+    /// </summary>
+    internal string PayloadValidationJmesPath => _payloadValidationJmesPath;
+
+    /// <summary>
+    /// Gets whether to throw exception if no idempotency key is found.
+    /// </summary>
+    internal bool ThrowOnNoIdempotencyKey => _throwOnNoIdempotencyKey;
+
+    /// <summary>
+    /// Gets whether local cache is enabled.
+    /// </summary>
+    internal bool UseLocalCache => _useLocalCache;
+
+    /// <summary>
+    /// Gets the maximum number of items in the local cache.
+    /// </summary>
+    internal int LocalCacheMaxItems => _localCacheMaxItems;
+
+    /// <summary>
+    /// Gets the expiration in seconds.
+    /// </summary>
+    internal long ExpirationInSeconds => _expirationInSeconds;
+
+    /// <summary>
+    /// Gets the hash function.
+    /// </summary>
+    internal string HashFunction => _hashFunction;
+
+    /// <summary>
+    /// Gets the response hook function.
+    /// </summary>
+    internal Func<object, AWS.Lambda.Powertools.Idempotency.Persistence.DataRecord, object> ResponseHook => _responseHook;
+
+    /// <summary>
+    /// Initialize and return an instance of IdempotencyOptions.
     /// Example:
-    /// IdempotencyConfig.Builder().WithUseLocalCache().Build();
-    /// This instance must then be passed to the Idempotency.Config:
-    /// Idempotency.Config().WithConfig(config).Configure();
+    /// new IdempotencyOptionsBuilder().WithUseLocalCache().Build();
+    /// This instance can then be passed to Idempotency.Configure:
+    /// Idempotency.Configure(builder => builder.WithOptions(options));
     /// </summary>
-    /// <returns>an instance of IdempotencyConfig</returns>
-    public IdempotencyOptions Build() =>
-        new(_eventKeyJmesPath,
-            _payloadValidationJmesPath,
-            _throwOnNoIdempotencyKey,
-            _useLocalCache,
-            _localCacheMaxItems,
-            _expirationInSeconds,
-            _hashFunction);
+    /// <returns>an instance of IdempotencyOptions</returns>
+    public IdempotencyOptions Build() => new(this);
 
     /// <summary>
     /// A JMESPath expression to extract the idempotency key from the event record.
@@ -133,4 +171,15 @@ public IdempotencyOptionsBuilder WithHashFunction(string hashFunction)
 #endif
         return this;
     }
+
+    /// <summary>
+    /// Set a response hook function, to be called with the response and the data record.
+    /// </summary>
+    /// <param name="hook">The response hook function</param>
+    /// <returns>the instance of the builder (to chain operations)</returns>
+    public IdempotencyOptionsBuilder WithResponseHook(Func<object, AWS.Lambda.Powertools.Idempotency.Persistence.DataRecord, object> hook)
+    {
+        _responseHook = hook;
+        return this;
+    }
 }

libraries/src/AWS.Lambda.Powertools.Idempotency/Internal/IdempotencyAspectHandler.cs

@@ -208,6 +208,24 @@ private Task<T> HandleForStatus(DataRecord record)
                     {
                         throw new IdempotencyPersistenceLayerException("Unable to cast function response as " + typeof(T).Name);
                     }
+                    // Response hook logic
+                    var responseHook = Idempotency.Instance.IdempotencyOptions?.ResponseHook;
+                    if (responseHook != null)
+                    {
+                        try
+                        {
+                            var hooked = responseHook(result, record);
+                            if (hooked is T typedHooked)
+                            {
+                                return Task.FromResult(typedHooked);
+                            }
+                            // If hook returns wrong type, fallback to original result
+                        }
+                        catch (Exception)
+                        {
+                            // If hook throws, fallback to original result
+                        }
+                    }
                     return Task.FromResult(result);
                 }
                 catch (Exception e)

libraries/src/AWS.Lambda.Powertools.Idempotency/Persistence/DynamoDBPersistenceStore.cs

@@ -370,9 +370,9 @@ public class DynamoDBPersistenceStoreBuilder
     private AmazonDynamoDBClient _dynamoDbClient;
 
     /// <summary>
-    /// Initialize and return a new instance of {@link DynamoDBPersistenceStore}.
+    /// Initialize and return a new instance of <see cref="DynamoDBPersistenceStore"/>.
     /// Example:
-    ///    DynamoDBPersistenceStore.builder().withTableName("idempotency_store").build();
+    ///    new DynamoDBPersistenceStoreBuilder().WithTableName("idempotency_store").Build();
     /// </summary>
     /// <returns></returns>
     /// <exception cref="ArgumentNullException"></exception>