Merge pull request #930 from aws-powertools/feature/idempotency-ReturnValuesOnConditionCheckFailure

feat: add support for ReturnValuesOnConditionCheckFailure in Idempotency

Author: hjgraca

libraries/src/AWS.Lambda.Powertools.Idempotency/Exceptions/IdempotencyItemAlreadyExistsException.cs

@@ -14,6 +14,7 @@
  */
 
 using System;
+using AWS.Lambda.Powertools.Idempotency.Persistence;
 
 namespace AWS.Lambda.Powertools.Idempotency.Exceptions;
 
@@ -22,6 +23,11 @@ namespace AWS.Lambda.Powertools.Idempotency.Exceptions;
 /// </summary>
 public class IdempotencyItemAlreadyExistsException : Exception
 {
+    /// <summary>
+    /// The record that already exists in the persistence layer.
+    /// </summary>
+    public DataRecord Record { get; set; }
+    
     /// <summary>
     /// Creates a new IdempotencyItemAlreadyExistsException
     /// </summary>

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

@@ -106,9 +106,25 @@ private async Task<T> ProcessIdempotency()
             // already exists. If it succeeds, there's no need to call getRecord.
             await _persistenceStore.SaveInProgress(_data, DateTimeOffset.UtcNow, GetRemainingTimeInMillis());
         }
-        catch (IdempotencyItemAlreadyExistsException)
+        catch (IdempotencyItemAlreadyExistsException ex)
         {
-            var record = await GetIdempotencyRecord();
+            DataRecord record;
+    
+            if(ex.Record != null)
+            {
+                // If the error includes the existing record, we can use it to validate
+                // the record being processed and cache it in memory.
+                var existingRecord = _persistenceStore.ProcessExistingRecord(ex.Record, _data);
+                record = existingRecord;
+            }
+            else
+            {
+                // If the error doesn't include the existing record, we need to fetch
+                // it from the persistence layer. In doing so, we also call the processExistingRecord
+                // method to validate the record and cache it in memory.
+                record = await GetIdempotencyRecord();
+            }
+    
             return await HandleForStatus(record);
         }
         catch (IdempotencyKeyException)

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

@@ -378,4 +378,20 @@ private static string GetHash(HashAlgorithm hashAlgorithm, string input)
 
     /// <inheritdoc />
     public abstract Task DeleteRecord(string idempotencyKey);
+
+    /// <summary>
+    /// Validates an existing record against the data payload being processed.
+    /// If the payload does not match the stored record, an `IdempotencyValidationError` error is thrown.
+    /// Whenever a record is retrieved from the persistence layer, it should be validated against the data payload
+    /// being processed. This is to ensure that the data payload being processed is the same as the one that was
+    /// used to create the record in the first place.
+    ///
+    /// The record is also saved to the local cache if local caching is enabled.
+    /// </summary>
+    public virtual DataRecord ProcessExistingRecord(DataRecord exRecord, JsonDocument data)
+    {
+        ValidatePayload(data, exRecord);
+        SaveToCache(exRecord);
+        return exRecord;
+    }
 }

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

@@ -191,6 +191,7 @@ public override async Task PutRecord(DataRecord record, DateTimeOffset now)
                 Item = item,
                 ConditionExpression = "attribute_not_exists(#id) OR #expiry < :now OR (attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now_milliseconds AND #status = :inprogress)",
                 ExpressionAttributeNames = expressionAttributeNames,
+                ReturnValuesOnConditionCheckFailure = ReturnValuesOnConditionCheckFailure.ALL_OLD,
                 ExpressionAttributeValues = new Dictionary<string, AttributeValue>
                 {
                     {":now", new AttributeValue {N = now.ToUnixTimeSeconds().ToString()}},
@@ -202,8 +203,20 @@ public override async Task PutRecord(DataRecord record, DateTimeOffset now)
         }
         catch (ConditionalCheckFailedException e)
         {
-            throw new IdempotencyItemAlreadyExistsException(
+            var ex = new IdempotencyItemAlreadyExistsException(
                 "Failed to put record for already existing idempotency key: " + record.IdempotencyKey, e);
+
+            if (e.Item != null)
+            {
+                ex.Record = new DataRecord(e.Item[_keyAttr].S,
+                    Enum.Parse<DataRecord.DataRecordStatus>(e.Item[_statusAttr].S),
+                    long.Parse(e.Item[_expiryAttr].N),
+                    e.Item.TryGetValue(_dataAttr, out var data) ? data?.S : null,
+                    e.Item.TryGetValue(_validationAttr, out var validation) ? validation?.S : null,
+                    e.Item.TryGetValue(_inProgressExpiryAttr, out var inProgExp) ? long.Parse(inProgExp.N) : null);
+            }
+
+            throw ex;
         }
     }
 

libraries/tests/AWS.Lambda.Powertools.Idempotency.Tests/Persistence/BasePersistenceStoreTests.cs

@@ -524,7 +524,7 @@ public void GenerateHash_WhenInputIsDouble_ShouldGenerateMd5ofDouble()
         // Assert
         generatedHash.Should().Be(expectedHash);
     }
-    
+
     [Fact]
     public async Task When_Key_Prefix_Set_Should_Create_With_Prefix()
     {
@@ -563,4 +563,35 @@ private static APIGatewayProxyRequest LoadApiGatewayProxyRequest()
             throw;
         }
     }
+
+    [Fact]
+    public async Task ProcessExistingRecord_WhenValidRecord_ShouldReturnRecordAndSaveToCache()
+    {
+        // Arrange
+        var persistenceStore = new InMemoryPersistenceStore();
+        var request = LoadApiGatewayProxyRequest();
+        LRUCache<string, DataRecord> cache = new(2);
+
+        persistenceStore.Configure(new IdempotencyOptionsBuilder()
+            .WithUseLocalCache(true)
+            .Build(), null, null, cache);
+
+        var now = DateTimeOffset.UtcNow;
+        var existingRecord = new DataRecord(
+            "testFunction#5eff007a9ed2789a9f9f6bc182fc6ae6",
+            DataRecord.DataRecordStatus.COMPLETED,
+            now.AddSeconds(3600).ToUnixTimeSeconds(),
+            "existing response",
+            null);
+
+        // Act
+        var result =
+            persistenceStore.ProcessExistingRecord(existingRecord, JsonSerializer.SerializeToDocument(request)!);
+
+        // Assert
+        result.Should().Be(existingRecord);
+        cache.Count.Should().Be(1);
+        cache.TryGet("testFunction#5eff007a9ed2789a9f9f6bc182fc6ae6", out var cachedRecord).Should().BeTrue();
+        cachedRecord.Should().Be(existingRecord);
+    }
 }