Merge pull request #745 from hjgraca/feature/idempotency-key-prefix

chore: Add support for custom key prefixes in Idempotency utility

Author: hjgraca

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

@@ -63,6 +63,11 @@ namespace AWS.Lambda.Powertools.Idempotency;
 [Injection(typeof(UniversalWrapperAspect), Inherited = true)]
 public class IdempotentAttribute : UniversalWrapperAttribute
 {
+    /// <summary>
+    /// Custom prefix for idempotency key: key_prefix#hash
+    /// </summary>
+    public string KeyPrefix { get; set; }
+    
     /// <summary>
     ///     Wraps as a synchronous operation, simply throws IdempotencyConfigurationException
     /// </summary>
@@ -90,7 +95,7 @@ protected internal sealed override T WrapSync<T>(Func<object[], T> target, objec
 
         Task<T> ResultDelegate() => Task.FromResult(target(args));
 
-        var idempotencyHandler = new IdempotencyAspectHandler<T>(ResultDelegate, eventArgs.Method.Name, payload,GetContext(eventArgs));
+        var idempotencyHandler = new IdempotencyAspectHandler<T>(ResultDelegate, eventArgs.Method.Name, KeyPrefix, payload,GetContext(eventArgs));
         if (idempotencyHandler == null)
         {
             throw new Exception("Failed to create an instance of IdempotencyAspectHandler");
@@ -128,7 +133,7 @@ protected internal sealed override async Task<T> WrapAsync<T>(
 
         Task<T> ResultDelegate() => target(args);
 
-        var idempotencyHandler = new IdempotencyAspectHandler<T>(ResultDelegate, eventArgs.Method.Name, payload, GetContext(eventArgs));
+        var idempotencyHandler = new IdempotencyAspectHandler<T>(ResultDelegate, eventArgs.Method.Name, KeyPrefix, payload, GetContext(eventArgs));
         if (idempotencyHandler == null)
         {
             throw new Exception("Failed to create an instance of IdempotencyAspectHandler");

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

@@ -50,19 +50,21 @@ internal class IdempotencyAspectHandler<T>
     /// </summary>
     /// <param name="target"></param>
     /// <param name="functionName"></param>
+    /// <param name="keyPrefix"></param>
     /// <param name="payload"></param>
     /// <param name="lambdaContext"></param>
     public IdempotencyAspectHandler(
         Func<Task<T>> target,
         string functionName,
+        string keyPrefix,
         JsonDocument payload,
         ILambdaContext lambdaContext)
     {
         _target = target;
         _data = payload;
         _lambdaContext = lambdaContext;
         _persistenceStore = Idempotency.Instance.PersistenceStore;
-        _persistenceStore.Configure(Idempotency.Instance.IdempotencyOptions, functionName);
+        _persistenceStore.Configure(Idempotency.Instance.IdempotencyOptions, functionName, keyPrefix);
     }
 
     /// <summary>

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

@@ -1,12 +1,12 @@
 /*
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
- * 
+ *
  * Licensed under the Apache License, Version 2.0 (the "License").
  * You may not use this file except in compliance with the License.
  * A copy of the License is located at
- * 
+ *
  *  http://aws.amazon.com/apache2.0
- * 
+ *
  * or in the "license" file accompanying this file. This file is distributed
  * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
  * express or implied. See the License for the specific language governing
@@ -37,16 +37,17 @@ public abstract class BasePersistenceStore : IPersistenceStore
     /// Idempotency Options
     /// </summary>
     private IdempotencyOptions _idempotencyOptions = null!;
-    
+
     /// <summary>
     /// Function name
     /// </summary>
     private string _functionName;
+
     /// <summary>
     /// Boolean to indicate whether or not payload validation is enabled
     /// </summary>
     protected bool PayloadValidationEnabled;
-    
+
     /// <summary>
     /// LRUCache
     /// </summary>
@@ -57,34 +58,45 @@ public abstract class BasePersistenceStore : IPersistenceStore
     /// </summary>
     /// <param name="idempotencyOptions">Idempotency configuration settings</param>
     /// <param name="functionName">The name of the function being decorated</param>
-    public void Configure(IdempotencyOptions idempotencyOptions, string functionName)
+    /// <param name="keyPrefix"></param>
+    public void Configure(IdempotencyOptions idempotencyOptions, string functionName, string keyPrefix)
     {
-        var funcEnv = Environment.GetEnvironmentVariable(Constants.LambdaFunctionNameEnv);
-        _functionName = funcEnv ?? "testFunction";
-        if (!string.IsNullOrWhiteSpace(functionName))
+        if (!string.IsNullOrEmpty(keyPrefix))
         {
-            _functionName += "." + functionName;
+            _functionName = keyPrefix;
         }
+        else
+        {
+            var funcEnv = Environment.GetEnvironmentVariable(Constants.LambdaFunctionNameEnv);
+
+            _functionName = funcEnv ?? "testFunction";
+            if (!string.IsNullOrWhiteSpace(functionName))
+            {
+                _functionName += "." + functionName;
+            }
+        }
+
         _idempotencyOptions = idempotencyOptions;
-        
+
         if (!string.IsNullOrWhiteSpace(_idempotencyOptions.PayloadValidationJmesPath))
         {
             PayloadValidationEnabled = true;
         }
-        
+
         var useLocalCache = _idempotencyOptions.UseLocalCache;
         if (useLocalCache)
         {
             _cache = new LRUCache<string, DataRecord>(_idempotencyOptions.LocalCacheMaxItems);
         }
     }
-    
+
     /// <summary>
     /// For test purpose only (adding a cache to mock)
     /// </summary>
-    internal void Configure(IdempotencyOptions options, string functionName, LRUCache<string, DataRecord> cache)
+    internal void Configure(IdempotencyOptions options, string functionName, string keyPrefix,
+        LRUCache<string, DataRecord> cache)
     {
-        Configure(options, functionName);
+        Configure(options, functionName, keyPrefix);
         _cache = cache;
     }
 
@@ -118,12 +130,12 @@ public virtual async Task SaveSuccess(JsonDocument data, object result, DateTime
     public virtual async Task SaveInProgress(JsonDocument data, DateTimeOffset now, double? remainingTimeInMs)
     {
         var idempotencyKey = GetHashedIdempotencyKey(data);
-        
+
         if (RetrieveFromCache(idempotencyKey, now) != null)
         {
             throw new IdempotencyItemAlreadyExistsException();
         }
-        
+
         long? inProgressExpirationMsTimestamp = null;
         if (remainingTimeInMs.HasValue)
         {
@@ -137,11 +149,10 @@ public virtual async Task SaveInProgress(JsonDocument data, DateTimeOffset now,
             null,
             GetHashedPayload(data),
             inProgressExpirationMsTimestamp
-            
         );
         await PutRecord(record, now);
     }
-    
+
     /// <summary>
     /// Delete record from the persistence store
     /// </summary>
@@ -152,14 +163,14 @@ public virtual async Task DeleteRecord(JsonDocument data, Exception throwable)
         var idemPotencyKey = GetHashedIdempotencyKey(data);
 
         Console.WriteLine("Function raised an exception {0}. " +
-                  "Clearing in progress record in persistence store for idempotency key: {1}",
+                          "Clearing in progress record in persistence store for idempotency key: {1}",
             throwable.GetType().Name,
             idemPotencyKey);
 
         await DeleteRecord(idemPotencyKey);
         DeleteFromCache(idemPotencyKey);
     }
-    
+
     /// <summary>
     /// Retrieve idempotency key for data provided, fetch from persistence store, and convert to DataRecord.
     /// </summary>
@@ -182,7 +193,7 @@ public virtual async Task<DataRecord> GetRecord(JsonDocument data, DateTimeOffse
         ValidatePayload(data, record);
         return record;
     }
-    
+
     /// <summary>
     /// Save data_record to local cache except when status is "INPROGRESS"
     /// NOTE: We can't cache "INPROGRESS" records as we have no way to reflect updates that can happen outside of the
@@ -198,7 +209,7 @@ private void SaveToCache(DataRecord dataRecord)
 
         _cache.Set(dataRecord.IdempotencyKey, dataRecord);
     }
-    
+
     /// <summary>
     /// Validate that the hashed payload matches data provided and stored data record
     /// </summary>
@@ -215,7 +226,7 @@ private void ValidatePayload(JsonDocument data, DataRecord dataRecord)
             throw new IdempotencyValidationException("Payload does not match stored record for this event key");
         }
     }
-    
+
     /// <summary>
     /// Retrieve data record from cache
     /// </summary>
@@ -228,14 +239,15 @@ private DataRecord RetrieveFromCache(string idempotencyKey, DateTimeOffset now)
             return null;
 
         if (!_cache.TryGet(idempotencyKey, out var record) || record == null) return null;
-        if (!record.IsExpired(now)) 
+        if (!record.IsExpired(now))
         {
             return record;
         }
+
         DeleteFromCache(idempotencyKey);
         return null;
     }
-    
+
     /// <summary>
     /// Deletes item from cache
     /// </summary>
@@ -244,10 +256,10 @@ private void DeleteFromCache(string idempotencyKey)
     {
         if (!_idempotencyOptions.UseLocalCache)
             return;
-        
+
         _cache.Delete(idempotencyKey);
     }
-    
+
     /// <summary>
     /// Extract payload using validation key jmespath and return a hashed representation
     /// </summary>
@@ -259,12 +271,12 @@ private string GetHashedPayload(JsonDocument data)
         {
             return "";
         }
-        
+
         var transformer = JsonTransformer.Parse(_idempotencyOptions.PayloadValidationJmesPath);
         var result = transformer.Transform(data.RootElement);
         return GenerateHash(result.RootElement);
     }
-    
+
     /// <summary>
     /// Calculate unix timestamp of expiry date for idempotency record
     /// </summary>
@@ -285,7 +297,7 @@ private string GetHashedIdempotencyKey(JsonDocument data)
     {
         var node = data.RootElement;
         var eventKeyJmesPath = _idempotencyOptions.EventKeyJmesPath;
-        if (eventKeyJmesPath != null) 
+        if (eventKeyJmesPath != null)
         {
             var transformer = JsonTransformer.Parse(eventKeyJmesPath);
             var result = transformer.Transform(node);
@@ -298,7 +310,9 @@ private string GetHashedIdempotencyKey(JsonDocument data)
             {
                 throw new IdempotencyKeyException("No data found to create a hashed idempotency key");
             }
-            Console.WriteLine("No data found to create a hashed idempotency key. JMESPath: {0}", _idempotencyOptions.EventKeyJmesPath ?? string.Empty);
+
+            Console.WriteLine("No data found to create a hashed idempotency key. JMESPath: {0}",
+                _idempotencyOptions.EventKeyJmesPath ?? string.Empty);
         }
 
         var hash = GenerateHash(node);
@@ -313,9 +327,10 @@ private string GetHashedIdempotencyKey(JsonDocument data)
     private static bool IsMissingIdempotencyKey(JsonElement data)
     {
         return data.ValueKind == JsonValueKind.Null || data.ValueKind == JsonValueKind.Undefined
-            || (data.ValueKind == JsonValueKind.String && data.ToString() == string.Empty);
+                                                    || (data.ValueKind == JsonValueKind.String &&
+                                                        data.ToString() == string.Empty);
     }
-    
+
     /// <summary>
     /// Generate a hash value from the provided data
     /// </summary>
@@ -328,16 +343,16 @@ internal string GenerateHash(JsonElement data)
         // starting .NET 8 no option to change hash algorithm
         using var hashAlgorithm = MD5.Create();
 #else
-
         using var hashAlgorithm = HashAlgorithm.Create(_idempotencyOptions.HashFunction);
 #endif
         if (hashAlgorithm == null)
         {
             throw new ArgumentException("Invalid HashAlgorithm");
         }
+
         var stringToHash = data.ToString();
         var hash = GetHash(hashAlgorithm, stringToHash);
-        
+
         return hash;
     }
 
@@ -351,18 +366,18 @@ private static string GetHash(HashAlgorithm hashAlgorithm, string input)
     {
         // Convert the input string to a byte array and compute the hash.
         var data = hashAlgorithm.ComputeHash(Encoding.UTF8.GetBytes(input));
-        
+
         // Create a new Stringbuilder to collect the bytes
         // and create a string.
         var sBuilder = new StringBuilder();
-        
+
         // Loop through each byte of the hashed data
         // and format each one as a hexadecimal string.
         for (var i = 0; i < data.Length; i++)
         {
             sBuilder.Append(data[i].ToString("x2"));
         }
-        
+
         // Return the hexadecimal string.
         return sBuilder.ToString();
     }

libraries/tests/AWS.Lambda.Powertools.Idempotency.Tests/Handlers/IdempotencyAttributeWithCustomKeyPrefix.cs

@@ -0,0 +1,21 @@
+using System;
+using Amazon.Lambda.Core;
+
+namespace AWS.Lambda.Powertools.Idempotency.Tests.Handlers;
+
+/// <summary>
+/// Simple Lambda function with Idempotent attribute on a sub method with a custom prefix key
+/// </summary>
+public class IdempotencyAttributeWithCustomKeyPrefix
+{
+    public string HandleRequest(string input, ILambdaContext context) 
+    {
+        return ReturnGuid(input);
+    }
+
+    [Idempotent(KeyPrefix = "MyMethod")]
+    private string ReturnGuid(string p)
+    {
+        return Guid.NewGuid().ToString();
+    }
+}

libraries/tests/AWS.Lambda.Powertools.Idempotency.Tests/Handlers/IdempotencyFunctionMethodDecorated.cs

@@ -13,7 +13,6 @@
  * permissions and limitations under the License.
  */
 
-using System;
 using System.Collections.Generic;
 using System.IO;
 using System.Net.Http;
@@ -22,7 +21,6 @@
 using Amazon.DynamoDBv2;
 using Amazon.Lambda.APIGatewayEvents;
 using Amazon.Lambda.Core;
-using AWS.Lambda.Powertools.Idempotency.Tests.Model;
 
 namespace AWS.Lambda.Powertools.Idempotency.Tests.Handlers;
 
@@ -34,9 +32,6 @@ public IdempotencyFunctionMethodDecorated(AmazonDynamoDBClient client)
     {
         Idempotency.Configure(builder =>
             builder
-#if NET8_0_OR_GREATER
-                .WithJsonSerializationContext(TestJsonSerializerContext.Default)
-#endif
                 .UseDynamoDb(storeBuilder =>
                     storeBuilder
                         .WithTableName("idempotency_table")