Skip to content

Commit 7b18029

Browse files
phipagphipag
authored
feat(idempotency): Add support for ReturnValuesOnConditionCheckFailure in Idempotency. (#1821)
1 parent 98b36b5 commit 7b18029

File tree

9 files changed

+136
-55
lines changed

9 files changed

+136
-55
lines changed

‎powertools-idempotency/powertools-idempotency-core/src/main/java/software/amazon/lambda/powertools/idempotency/exceptions/IdempotencyItemAlreadyExistsException.java

+15
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,17 @@
1414

1515
package software.amazon.lambda.powertools.idempotency.exceptions;
1616

17+
import java.util.Optional;
18+
19+
import software.amazon.lambda.powertools.idempotency.persistence.DataRecord;
20+
1721
/**
1822
* Exception thrown when trying to store an item which already exists.
1923
*/
2024
public class IdempotencyItemAlreadyExistsException extends RuntimeException {
2125
private static final long serialVersionUID = 9027152772149436500L;
26+
// transient because we don't want to accidentally dump any payloads in logs / stack traces
27+
private transient Optional<DataRecord> dr = Optional.empty();
2228

2329
public IdempotencyItemAlreadyExistsException() {
2430
super();
@@ -27,4 +33,13 @@ public IdempotencyItemAlreadyExistsException() {
2733
public IdempotencyItemAlreadyExistsException(String msg, Throwable e) {
2834
super(msg, e);
2935
}
36+
37+
public IdempotencyItemAlreadyExistsException(String msg, Throwable e, DataRecord dr) {
38+
super(msg, e);
39+
this.dr = Optional.ofNullable(dr);
40+
}
41+
42+
public Optional<DataRecord> getDataRecord() {
43+
return dr;
44+
}
3045
}

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

+5-3
Original file line numberDiff line numberDiff line change
@@ -94,9 +94,11 @@ private Object processIdempotency() throws Throwable {
9494
// already exists. If it succeeds, there's no need to call getRecord.
9595
persistenceStore.saveInProgress(data, Instant.now(), getRemainingTimeInMillis());
9696
} catch (IdempotencyItemAlreadyExistsException iaee) {
97-
DataRecord record = getIdempotencyRecord();
98-
if (record != null) {
99-
return handleForStatus(record);
97+
// If a DataRecord is already present on the Exception we can immediately take that one instead of trying
98+
// to fetch it first.
99+
DataRecord dr = iaee.getDataRecord().orElseGet(this::getIdempotencyRecord);
100+
if (dr != null) {
101+
return handleForStatus(dr);
100102
}
101103
} catch (IdempotencyKeyException ike) {
102104
throw ike;

‎powertools-idempotency/powertools-idempotency-core/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DataRecord.java

+14-5
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ public class DataRecord {
5353
private final OptionalLong inProgressExpiryTimestamp;
5454

5555
public DataRecord(String idempotencyKey, Status status, long expiryTimestamp, String responseData,
56-
String payloadHash) {
56+
String payloadHash) {
5757
this.idempotencyKey = idempotencyKey;
5858
this.status = status.toString();
5959
this.expiryTimestamp = expiryTimestamp;
@@ -63,7 +63,7 @@ public DataRecord(String idempotencyKey, Status status, long expiryTimestamp, St
6363
}
6464

6565
public DataRecord(String idempotencyKey, Status status, long expiryTimestamp, String responseData,
66-
String payloadHash, OptionalLong inProgressExpiryTimestamp) {
66+
String payloadHash, OptionalLong inProgressExpiryTimestamp) {
6767
this.idempotencyKey = idempotencyKey;
6868
this.status = status.toString();
6969
this.expiryTimestamp = expiryTimestamp;
@@ -131,13 +131,22 @@ public int hashCode() {
131131
return Objects.hash(idempotencyKey, status, expiryTimestamp, responseData, payloadHash);
132132
}
133133

134+
@Override
135+
public String toString() {
136+
return "DataRecord{" +
137+
"idempotencyKey='" + idempotencyKey + '\'' +
138+
", status='" + status + '\'' +
139+
", expiryTimestamp=" + expiryTimestamp +
140+
", payloadHash='" + payloadHash + '\'' +
141+
'}';
142+
}
134143

135144
/**
136145
* Status of the record:
137146
* <ul>
138-
* <li>INPROGRESS: record initialized when function starts</li>
139-
* <li>COMPLETED: record updated with the result of the function when it ends</li>
140-
* <li>EXPIRED: record expired, idempotency will not happen</li>
147+
* <li>INPROGRESS: record initialized when function starts</li>
148+
* <li>COMPLETED: record updated with the result of the function when it ends</li>
149+
* <li>EXPIRED: record expired, idempotency will not happen</li>
141150
* </ul>
142151
*/
143152
public enum Status {

‎powertools-idempotency/powertools-idempotency-core/src/test/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyAspectTest.java

+45-6
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import static org.mockito.ArgumentMatchers.any;
2121
import static org.mockito.Mockito.doReturn;
2222
import static org.mockito.Mockito.doThrow;
23+
import static org.mockito.Mockito.never;
2324
import static org.mockito.Mockito.spy;
2425
import static org.mockito.Mockito.verify;
2526
import static org.mockito.Mockito.verifyNoInteractions;
@@ -154,7 +155,7 @@ public void secondCall_notExpired_shouldGetFromStore() throws JsonProcessingExce
154155
.build())
155156
.configure();
156157

157-
doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any(), any());
158+
doThrow(new IdempotencyItemAlreadyExistsException()).when(store).saveInProgress(any(), any(), any());
158159

159160
Product p = new Product(42, "fake product", 12);
160161
Basket b = new Basket(p);
@@ -175,6 +176,44 @@ public void secondCall_notExpired_shouldGetFromStore() throws JsonProcessingExce
175176
assertThat(function.handlerCalled()).isFalse();
176177
}
177178

179+
@Test
180+
public void secondCall_notExpired_shouldNotGetFromStoreIfPresentOnIdempotencyException()
181+
throws JsonProcessingException {
182+
// GIVEN
183+
Idempotency.config()
184+
.withPersistenceStore(store)
185+
.withConfig(IdempotencyConfig.builder()
186+
.withEventKeyJMESPath("id")
187+
.build())
188+
.configure();
189+
190+
Product p = new Product(42, "fake product", 12);
191+
Basket b = new Basket(p);
192+
DataRecord dr = new DataRecord(
193+
"42",
194+
DataRecord.Status.COMPLETED,
195+
Instant.now().plus(356, SECONDS).getEpochSecond(),
196+
JsonConfig.get().getObjectMapper().writer().writeValueAsString(b),
197+
null);
198+
199+
// A data record on this exception should take precedence over fetching a record from the store / cache
200+
doThrow(new IdempotencyItemAlreadyExistsException(
201+
"Test message",
202+
new RuntimeException("Test Cause"),
203+
dr))
204+
.when(store).saveInProgress(any(), any(), any());
205+
206+
// WHEN
207+
IdempotencyEnabledFunction function = new IdempotencyEnabledFunction();
208+
Basket basket = function.handleRequest(p, context);
209+
210+
// THEN
211+
assertThat(basket).isEqualTo(b);
212+
assertThat(function.handlerCalled()).isFalse();
213+
// Should never call the store because item is already present on IdempotencyItemAlreadyExistsException
214+
verify(store, never()).getRecord(any(), any());
215+
}
216+
178217
@Test
179218
public void secondCall_notExpired_shouldGetStringFromStore() {
180219
// GIVEN
@@ -185,7 +224,7 @@ public void secondCall_notExpired_shouldGetStringFromStore() {
185224
.build())
186225
.configure();
187226

188-
doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any(), any());
227+
doThrow(new IdempotencyItemAlreadyExistsException()).when(store).saveInProgress(any(), any(), any());
189228

190229
Product p = new Product(42, "fake product", 12);
191230
DataRecord dr = new DataRecord(
@@ -220,7 +259,7 @@ public void secondCall_notExpired_shouldGetStringFromStoreWithResponseHook() {
220259
.build())
221260
.configure();
222261

223-
doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any(), any());
262+
doThrow(new IdempotencyItemAlreadyExistsException()).when(store).saveInProgress(any(), any(), any());
224263

225264
Product p = new Product(42, "fake product", 12);
226265
DataRecord dr = new DataRecord(
@@ -251,7 +290,7 @@ public void secondCall_inProgress_shouldThrowIdempotencyAlreadyInProgressExcepti
251290
.build())
252291
.configure();
253292

254-
doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any(), any());
293+
doThrow(new IdempotencyItemAlreadyExistsException()).when(store).saveInProgress(any(), any(), any());
255294

256295
Product p = new Product(42, "fake product", 12);
257296
Basket b = new Basket(p);
@@ -283,7 +322,7 @@ public void secondCall_inProgress_lambdaTimeout_timeoutExpired_shouldThrowIncons
283322
.build())
284323
.configure();
285324

286-
doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any(), any());
325+
doThrow(new IdempotencyItemAlreadyExistsException()).when(store).saveInProgress(any(), any(), any());
287326

288327
Product p = new Product(42, "fake product", 12);
289328
Basket b = new Basket(p);
@@ -412,7 +451,7 @@ public void idempotencyOnSubMethodAnnotated_secondCall_notExpired_shouldGetFromS
412451
.withPersistenceStore(store)
413452
.configure();
414453

415-
doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any(), any());
454+
doThrow(new IdempotencyItemAlreadyExistsException()).when(store).saveInProgress(any(), any(), any());
416455

417456
Product p = new Product(42, "fake product", 12);
418457
Basket b = new Basket(p);

‎powertools-idempotency/powertools-idempotency-dynamodb/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/dynamodb/DynamoDBPersistenceStore.java

+9-2
Original file line numberDiff line numberDiff line change
@@ -189,10 +189,17 @@ public void putRecord(DataRecord record, Instant now) throws IdempotencyItemAlre
189189
"attribute_not_exists(#id) OR #expiry < :now OR (attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now_milliseconds AND #status = :inprogress)")
190190
.expressionAttributeNames(expressionAttributeNames)
191191
.expressionAttributeValues(expressionAttributeValues)
192-
.build()
193-
);
192+
.returnValuesOnConditionCheckFailure("ALL_OLD")
193+
.build());
194194
} catch (ConditionalCheckFailedException e) {
195195
LOG.debug("Failed to put record for already existing idempotency key: {}", record.getIdempotencyKey());
196+
if (e.hasItem()) {
197+
DataRecord existingRecord = itemToRecord(e.item());
198+
throw new IdempotencyItemAlreadyExistsException(
199+
"Failed to put record for already existing idempotency key: " + record.getIdempotencyKey()
200+
+ ". Existing record: " + existingRecord,
201+
e, existingRecord);
202+
}
196203
throw new IdempotencyItemAlreadyExistsException(
197204
"Failed to put record for already existing idempotency key: " + record.getIdempotencyKey(), e);
198205
}

‎powertools-idempotency/powertools-idempotency-dynamodb/src/test/java/software.amazon.lambda.powertools.idempotency.dynamodb/DynamoDBConfig.java renamed to ‎powertools-idempotency/powertools-idempotency-dynamodb/src/test/java/software/amazon/lambda/powertools/idempotency/persistence/dynamodb/DynamoDBConfig.java

+9-7
Original file line numberDiff line numberDiff line change
@@ -12,15 +12,18 @@
1212
*
1313
*/
1414

15-
package software.amazon.lambda.powertools.idempotency.dynamodb;
15+
package software.amazon.lambda.powertools.idempotency.persistence.dynamodb;
1616

17-
import com.amazonaws.services.dynamodbv2.local.main.ServerRunner;
18-
import com.amazonaws.services.dynamodbv2.local.server.DynamoDBProxyServer;
1917
import java.io.IOException;
2018
import java.net.ServerSocket;
2119
import java.net.URI;
20+
2221
import org.junit.jupiter.api.AfterAll;
2322
import org.junit.jupiter.api.BeforeAll;
23+
24+
import com.amazonaws.services.dynamodbv2.local.main.ServerRunner;
25+
import com.amazonaws.services.dynamodbv2.local.server.DynamoDBProxyServer;
26+
2427
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
2528
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
2629
import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient;
@@ -66,13 +69,12 @@ public static void setupDynamo() {
6669
.tableName(TABLE_NAME)
6770
.keySchema(KeySchemaElement.builder().keyType(KeyType.HASH).attributeName("id").build())
6871
.attributeDefinitions(
69-
AttributeDefinition.builder().attributeName("id").attributeType(ScalarAttributeType.S).build()
70-
)
72+
AttributeDefinition.builder().attributeName("id").attributeType(ScalarAttributeType.S).build())
7173
.billingMode(BillingMode.PAY_PER_REQUEST)
7274
.build());
7375

74-
DescribeTableResponse response =
75-
client.describeTable(DescribeTableRequest.builder().tableName(TABLE_NAME).build());
76+
DescribeTableResponse response = client
77+
.describeTable(DescribeTableRequest.builder().tableName(TABLE_NAME).build());
7678
if (response == null) {
7779
throw new RuntimeException("Table was not created within expected time");
7880
}

‎powertools-idempotency/powertools-idempotency-dynamodb/src/test/java/software/amazon/lambda/powertools/idempotency/persistence/dynamodb/DynamoDBPersistenceStoreTest.java

+15-11
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@
3232
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
3333
import software.amazon.lambda.powertools.idempotency.Constants;
3434
import software.amazon.lambda.powertools.idempotency.IdempotencyConfig;
35-
import software.amazon.lambda.powertools.idempotency.dynamodb.DynamoDBConfig;
35+
3636
import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyItemAlreadyExistsException;
3737
import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyItemNotFoundException;
3838
import software.amazon.lambda.powertools.idempotency.persistence.DataRecord;
@@ -155,13 +155,14 @@ public void putRecord_shouldThrowIdempotencyItemAlreadyExistsException_IfRecordA
155155
DataRecord.Status.INPROGRESS,
156156
expiry2,
157157
null,
158-
null
159-
), now)
160-
).isInstanceOf(IdempotencyItemAlreadyExistsException.class);
158+
null),
159+
now)).isInstanceOf(IdempotencyItemAlreadyExistsException.class)
160+
// DataRecord should be present due to returnValuesOnConditionCheckFailure("ALL_OLD")
161+
.matches(e -> ((IdempotencyItemAlreadyExistsException) e).getDataRecord().isPresent());
161162

162163
// THEN: item was not updated, retrieve the initial one
163-
Map<String, AttributeValue> itemInDb =
164-
client.getItem(GetItemRequest.builder().tableName(TABLE_NAME).key(key).build()).item();
164+
Map<String, AttributeValue> itemInDb = client
165+
.getItem(GetItemRequest.builder().tableName(TABLE_NAME).key(key).build()).item();
165166
assertThat(itemInDb).isNotNull();
166167
assertThat(itemInDb.get("status").s()).isEqualTo("COMPLETED");
167168
assertThat(itemInDb.get("expiration").n()).isEqualTo(String.valueOf(expiry));
@@ -190,13 +191,16 @@ public void putRecord_shouldBlockUpdate_IfRecordAlreadyExistAndProgressNotExpire
190191
DataRecord.Status.INPROGRESS,
191192
expiry2,
192193
"Fake Data 2",
193-
null
194-
), now))
195-
.isInstanceOf(IdempotencyItemAlreadyExistsException.class);
194+
null),
195+
now))
196+
.isInstanceOf(IdempotencyItemAlreadyExistsException.class)
197+
// DataRecord should be present due to returnValuesOnConditionCheckFailure("ALL_OLD")
198+
.matches(e -> ((IdempotencyItemAlreadyExistsException) e).getDataRecord().isPresent());
199+
;
196200

197201
// THEN: item was not updated, retrieve the initial one
198-
Map<String, AttributeValue> itemInDb =
199-
client.getItem(GetItemRequest.builder().tableName(TABLE_NAME).key(key).build()).item();
202+
Map<String, AttributeValue> itemInDb = client
203+
.getItem(GetItemRequest.builder().tableName(TABLE_NAME).key(key).build()).item();
200204
assertThat(itemInDb).isNotNull();
201205
assertThat(itemInDb.get("status").s()).isEqualTo("INPROGRESS");
202206
assertThat(itemInDb.get("expiration").n()).isEqualTo(String.valueOf(expiry));

‎powertools-idempotency/powertools-idempotency-dynamodb/src/test/java/software.amazon.lambda.powertools.idempotency.dynamodb/IdempotencyTest.java renamed to ‎powertools-idempotency/powertools-idempotency-dynamodb/src/test/java/software/amazon/lambda/powertools/idempotency/persistence/dynamodb/IdempotencyTest.java

+11-10
Original file line numberDiff line numberDiff line change
@@ -12,20 +12,21 @@
1212
*
1313
*/
1414

15-
package software.amazon.lambda.powertools.idempotency.dynamodb;
16-
15+
package software.amazon.lambda.powertools.idempotency.persistence.dynamodb;
1716

1817
import static org.assertj.core.api.Assertions.assertThat;
1918

20-
import com.amazonaws.services.lambda.runtime.Context;
21-
import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent;
22-
import com.amazonaws.services.lambda.runtime.tests.EventLoader;
2319
import org.junit.jupiter.api.BeforeEach;
2420
import org.junit.jupiter.api.Test;
2521
import org.mockito.Mock;
2622
import org.mockito.MockitoAnnotations;
23+
24+
import com.amazonaws.services.lambda.runtime.Context;
25+
import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent;
26+
import com.amazonaws.services.lambda.runtime.tests.EventLoader;
27+
2728
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
28-
import software.amazon.lambda.powertools.idempotency.dynamodb.handlers.IdempotencyFunction;
29+
import software.amazon.lambda.powertools.idempotency.persistence.dynamodb.handlers.IdempotencyFunction;
2930

3031
public class IdempotencyTest extends DynamoDBConfig {
3132

@@ -41,14 +42,14 @@ void setUp() {
4142
public void endToEndTest() {
4243
IdempotencyFunction function = new IdempotencyFunction(client);
4344

44-
APIGatewayProxyResponseEvent response =
45-
function.handleRequest(EventLoader.loadApiGatewayRestEvent("apigw_event2.json"), context);
45+
APIGatewayProxyResponseEvent response = function
46+
.handleRequest(EventLoader.loadApiGatewayRestEvent("apigw_event2.json"), context);
4647
assertThat(function.handlerExecuted).isTrue();
4748

4849
function.handlerExecuted = false;
4950

50-
APIGatewayProxyResponseEvent response2 =
51-
function.handleRequest(EventLoader.loadApiGatewayRestEvent("apigw_event2.json"), context);
51+
APIGatewayProxyResponseEvent response2 = function
52+
.handleRequest(EventLoader.loadApiGatewayRestEvent("apigw_event2.json"), context);
5253
assertThat(function.handlerExecuted).isFalse();
5354

5455
assertThat(response).isEqualTo(response2);

0 commit comments

Comments
 (0)