Skip to content

Commit 12779d9

Browse files
roamarieiroamariei
committed
updated transact delete item flow to check the version for optimistic locking
1 parent 3ec9f08 commit 12779d9

File tree

16 files changed

+383
-16
lines changed

16 files changed

+383
-16
lines changed

‎.changes/next-release/bugfix-AmazonDynamoDBEnhancedClient-64b922e.json

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"type": "bugfix",
3+
"category": "Amazon DynamoDB Enhanced Client",
4+
"contributor": "",
5+
"description": "Optimistic locking while using DynamoDbEnhancedClient - DeleteItem with TransactWriteItemsEnhancedRequest"
6+
}

‎services-custom/dynamodb-enhanced/src/it/java/software/amazon/awssdk/enhanced/dynamodb/AsyncCrudWithResponseIntegrationTest.java

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515

1616
package software.amazon.awssdk.enhanced.dynamodb;
1717

18-
import static org.assertj.core.api.Assertions.as;
1918
import static org.assertj.core.api.Assertions.assertThat;
2019
import static org.assertj.core.api.Assertions.assertThatThrownBy;
2120

@@ -192,7 +191,7 @@ public void updateItem_returnValues_all_old() {
192191
Record record = new Record().setId("1").setSort(10);
193192
mappedTable.putItem(record).join();
194193

195-
Record updatedRecord = new Record().setId("1").setSort(10).setValue(11);
194+
Record updatedRecord = new Record().setId("1").setSort(10).setValue(11).setVersion(1);
196195

197196

198197
UpdateItemEnhancedResponse<Record> response = mappedTable.updateItemWithResponse(r -> r.item(updatedRecord)
@@ -202,14 +201,15 @@ public void updateItem_returnValues_all_old() {
202201
assertThat(response.attributes().getId()).isEqualTo(record.getId());
203202
assertThat(response.attributes().getSort()).isEqualTo(record.getSort());
204203
assertThat(response.attributes().getValue()).isEqualTo(null);
204+
assertThat(response.attributes().getVersion()).isEqualTo(1);
205205
}
206206

207207
@Test
208208
public void updateItem_returnValues_all_new() {
209209
Record record = new Record().setId("1").setSort(10);
210210
mappedTable.putItem(record).join();
211211

212-
Record updatedRecord = new Record().setId("1").setSort(10).setValue(11);
212+
Record updatedRecord = new Record().setId("1").setSort(10).setValue(11).setVersion(1);
213213

214214

215215
UpdateItemEnhancedResponse<Record> response = mappedTable.updateItemWithResponse(r -> r.item(updatedRecord)
@@ -219,14 +219,15 @@ public void updateItem_returnValues_all_new() {
219219
assertThat(response.attributes().getId()).isEqualTo(updatedRecord.getId());
220220
assertThat(response.attributes().getSort()).isEqualTo(updatedRecord.getSort());
221221
assertThat(response.attributes().getValue()).isEqualTo(updatedRecord.getValue());
222+
assertThat(response.attributes().getVersion()).isEqualTo(updatedRecord.getVersion() + 1);
222223
}
223224

224225
@Test
225226
public void updateItem_returnValues_not_set() {
226227
Record record = new Record().setId("1").setSort(10);
227228
mappedTable.putItem(record).join();
228229

229-
Record updatedRecord = new Record().setId("1").setSort(10).setValue(11);
230+
Record updatedRecord = new Record().setId("1").setSort(10).setValue(11).setVersion(1);
230231

231232

232233
UpdateItemEnhancedResponse<Record> response = mappedTable.updateItemWithResponse(r -> r.item(updatedRecord))
@@ -235,6 +236,7 @@ public void updateItem_returnValues_not_set() {
235236
assertThat(response.attributes().getId()).isEqualTo(updatedRecord.getId());
236237
assertThat(response.attributes().getSort()).isEqualTo(updatedRecord.getSort());
237238
assertThat(response.attributes().getValue()).isEqualTo(updatedRecord.getValue());
239+
assertThat(response.attributes().getVersion()).isEqualTo(updatedRecord.getVersion() + 1);
238240
}
239241

240242
@Test

‎services-custom/dynamodb-enhanced/src/it/java/software/amazon/awssdk/enhanced/dynamodb/CrudWithResponseIntegrationTest.java

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import static org.assertj.core.api.Assertions.assertThat;
1919
import static org.assertj.core.api.Assertions.assertThatThrownBy;
20+
import static org.junit.Assert.assertThrows;
2021

2122
import org.assertj.core.data.Offset;
2223
import org.junit.After;
@@ -30,6 +31,7 @@
3031
import software.amazon.awssdk.enhanced.dynamodb.model.PutItemEnhancedRequest;
3132
import software.amazon.awssdk.enhanced.dynamodb.model.PutItemEnhancedResponse;
3233
import software.amazon.awssdk.enhanced.dynamodb.model.Record;
34+
import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedRequest;
3335
import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedRequest;
3436
import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedResponse;
3537
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
@@ -40,6 +42,7 @@
4042
import software.amazon.awssdk.services.dynamodb.model.ReturnConsumedCapacity;
4143
import software.amazon.awssdk.services.dynamodb.model.ReturnItemCollectionMetrics;
4244
import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
45+
import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledException;
4346

4447
public class CrudWithResponseIntegrationTest extends DynamoDbEnhancedIntegrationTestBase {
4548

@@ -218,6 +221,54 @@ public void deleteItem_returnValuesOnConditionCheckFailure_set_returnValuesOnCon
218221
.satisfies(e -> assertThat(((ConditionalCheckFailedException) e).hasItem()).isFalse());
219222
}
220223

224+
@Test
225+
public void deleteItemWithTransactWrite_shouldFailIfVersionMismatch() {
226+
Record originalItem = new Record().setId("123").setSort(10).setStringAttribute("Original Item");
227+
Key recordKey = Key.builder()
228+
.partitionValue(originalItem.getId())
229+
.sortValue(originalItem.getSort())
230+
.build();
231+
232+
mappedTable.putItem(originalItem);
233+
234+
// Retrieve the item and modify it separately
235+
Record modifiedItem = mappedTable.getItem(r -> r.key(recordKey));
236+
modifiedItem.setStringAttribute("Updated Item");
237+
238+
// Update the item, which will increment the version
239+
mappedTable.updateItem(modifiedItem);
240+
241+
// Now attempt to delete the original item using a transaction
242+
TransactWriteItemsEnhancedRequest request = TransactWriteItemsEnhancedRequest.builder()
243+
.addDeleteItem(mappedTable, modifiedItem)
244+
.build();
245+
246+
assertThrows(TransactionCanceledException.class, () -> enhancedClient.transactWriteItems(request));
247+
}
248+
249+
@Test
250+
public void deleteItemWithTransactWrite_shouldSucceedIfVersionMatch() {
251+
Record originalItem = new Record().setId("123").setSort(10).setStringAttribute("Original Item");
252+
Key recordKey = Key.builder()
253+
.partitionValue(originalItem.getId())
254+
.sortValue(originalItem.getSort())
255+
.build();
256+
mappedTable.putItem(originalItem);
257+
258+
// Retrieve the item
259+
Record retrievedItem = mappedTable.getItem(r -> r.key(recordKey));
260+
261+
// Delete the item using a transaction
262+
TransactWriteItemsEnhancedRequest request = TransactWriteItemsEnhancedRequest.builder()
263+
.addDeleteItem(mappedTable, retrievedItem)
264+
.build();
265+
266+
enhancedClient.transactWriteItems(request);
267+
268+
Record deletedItem = mappedTable.getItem(r -> r.key(recordKey));
269+
assertThat(deletedItem).isNull();
270+
}
271+
221272
@Test
222273
public void deleteItem_returnValuesOnConditionCheckFailure_set_returnValuesOnConditionCheckFailureNotNull() {
223274
Record record = new Record().setId("1").setSort(10);

‎services-custom/dynamodb-enhanced/src/it/java/software/amazon/awssdk/enhanced/dynamodb/DynamoDbEnhancedIntegrationTestBase.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515

1616
package software.amazon.awssdk.enhanced.dynamodb;
1717

18+
import static software.amazon.awssdk.enhanced.dynamodb.extensions.VersionedRecordExtension.AttributeTags.versionAttribute;
1819
import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.primaryPartitionKey;
1920
import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.primarySortKey;
2021
import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.secondaryPartitionKey;
@@ -73,6 +74,10 @@ protected static DynamoDbAsyncClient createAsyncDynamoDbClient() {
7374
.addAttribute(String.class, a -> a.name("stringAttribute")
7475
.getter(Record::getStringAttribute)
7576
.setter(Record::setStringAttribute))
77+
.addAttribute(Integer.class, a -> a.name("version")
78+
.getter(Record::getVersion)
79+
.setter(Record::setVersion)
80+
.tags(versionAttribute()))
7681
.build();
7782

7883

‎services-custom/dynamodb-enhanced/src/it/java/software/amazon/awssdk/enhanced/dynamodb/ScanQueryIntegrationTest.java

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,8 @@ public static void setup() {
5757
mappedTable = enhancedClient.table(TABLE_NAME, TABLE_SCHEMA);
5858
mappedTable.createTable();
5959
dynamoDbClient.waiter().waitUntilTableExists(r -> r.tableName(TABLE_NAME));
60+
RECORDS.forEach(record -> mappedTable.putItem(r -> r.item(record)));
61+
RECORDS.forEach(record -> record.setVersion(1));
6062
}
6163

6264
@AfterClass
@@ -74,8 +76,6 @@ private void insertRecords() {
7476

7577
@Test
7678
public void scan_withoutReturnConsumedCapacity_checksPageCount() {
77-
insertRecords();
78-
7979
Iterator<Page<Record>> results = mappedTable.scan(ScanEnhancedRequest.builder().limit(5).build())
8080
.iterator();
8181
Page<Record> page1 = results.next();
@@ -97,8 +97,6 @@ public void scan_withoutReturnConsumedCapacity_checksPageCount() {
9797

9898
@Test
9999
public void scan_withReturnConsumedCapacityAndDifferentReadConsistency_checksConsumedCapacity() {
100-
insertRecords();
101-
102100
Iterator<Page<Record>> eventualConsistencyResult =
103101
mappedTable.scan(ScanEnhancedRequest.builder().returnConsumedCapacity(ReturnConsumedCapacity.TOTAL).build())
104102
.iterator();
@@ -122,8 +120,6 @@ public void scan_withReturnConsumedCapacityAndDifferentReadConsistency_checksCon
122120

123121
@Test
124122
public void query_withoutReturnConsumedCapacity_checksPageCount() {
125-
insertRecords();
126-
127123
Iterator<Page<Record>> results =
128124
mappedTable.query(QueryEnhancedRequest.builder()
129125
.queryConditional(sortBetween(k-> k.partitionValue("id-value").sortValue(2),
@@ -151,8 +147,6 @@ public void query_withoutReturnConsumedCapacity_checksPageCount() {
151147

152148
@Test
153149
public void query_withReturnConsumedCapacityAndDifferentReadConsistency_checksConsumedCapacity() {
154-
insertRecords();
155-
156150
Iterator<Page<Record>> eventualConsistencyResult =
157151
mappedTable.query(QueryEnhancedRequest.builder()
158152
.queryConditional(sortGreaterThan(k -> k.partitionValue("id-value").sortValue(3)))

‎services-custom/dynamodb-enhanced/src/it/java/software/amazon/awssdk/enhanced/dynamodb/model/Record.java

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,10 @@
1616
package software.amazon.awssdk.enhanced.dynamodb.model;
1717

1818
import java.util.Objects;
19+
import software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbVersionAttribute;
20+
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbBean;
1921

22+
@DynamoDbBean
2023
public class Record {
2124

2225
private String id;
@@ -26,6 +29,7 @@ public class Record {
2629
private Integer gsiSort;
2730

2831
private String stringAttribute;
32+
private Integer version;
2933

3034
public String getId() {
3135
return id;
@@ -81,6 +85,16 @@ public Record setStringAttribute(String stringAttribute) {
8185
return this;
8286
}
8387

88+
@DynamoDbVersionAttribute
89+
public Integer getVersion() {
90+
return version;
91+
}
92+
93+
public Record setVersion(Integer version) {
94+
this.version = version;
95+
return this;
96+
}
97+
8498
@Override
8599
public boolean equals(Object o) {
86100
if (this == o) return true;
@@ -91,11 +105,12 @@ public boolean equals(Object o) {
91105
Objects.equals(value, record.value) &&
92106
Objects.equals(gsiId, record.gsiId) &&
93107
Objects.equals(stringAttribute, record.stringAttribute) &&
94-
Objects.equals(gsiSort, record.gsiSort);
108+
Objects.equals(gsiSort, record.gsiSort) &&
109+
Objects.equals(version, record.version);
95110
}
96111

97112
@Override
98113
public int hashCode() {
99-
return Objects.hash(id, sort, value, gsiId, gsiSort, stringAttribute);
114+
return Objects.hash(id, sort, value, gsiId, gsiSort, stringAttribute, version);
100115
}
101116
}

‎services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/DynamoDbEnhancedClientExtension.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,4 +55,15 @@ default WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite conte
5555
default ReadModification afterRead(DynamoDbExtensionContext.AfterRead context) {
5656
return ReadModification.builder().build();
5757
}
58+
59+
/**
60+
* This hook is called just before an operation is going to delete data from the database. The extension that
61+
* implements this method can add a condition to the delete operation.
62+
*
63+
* @param context The {@link DynamoDbExtensionContext.BeforeDelete} context containing the state of the execution.
64+
* @return A {@link Expression} object that can alter the behavior of the delete operation.
65+
*/
66+
default Expression beforeDelete(DynamoDbExtensionContext.BeforeDelete context) {
67+
return Expression.builder().build();
68+
}
5869
}

‎services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/DynamoDbExtensionContext.java

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,4 +76,12 @@ public interface BeforeWrite extends Context {
7676
@ThreadSafe
7777
public interface AfterRead extends Context {
7878
}
79+
80+
/**
81+
* The state of the execution when the {@link DynamoDbEnhancedClientExtension#beforeDelete} method is invoked.
82+
*/
83+
@SdkPublicApi
84+
@ThreadSafe
85+
public interface BeforeDelete extends Context {
86+
}
7987
}

‎services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/extensions/VersionedRecordExtension.java

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,40 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
142142
.build();
143143
}
144144

145+
@Override
146+
public Expression beforeDelete(DynamoDbExtensionContext.BeforeDelete context) {
147+
Optional<String> versionAttributeKey = context.tableMetadata()
148+
.customMetadataObject(CUSTOM_METADATA_KEY, String.class);
149+
150+
if (!versionAttributeKey.isPresent()) {
151+
return Expression.builder().build();
152+
}
153+
154+
String attributeKeyRef = keyRef(versionAttributeKey.get());
155+
Expression condition;
156+
Optional<AttributeValue> existingVersionValue =
157+
Optional.ofNullable(context.items().get(versionAttributeKey.get()));
158+
159+
if (!existingVersionValue.isPresent() || isNullAttributeValue(existingVersionValue.get())) {
160+
throw new IllegalArgumentException("Version attribute is null.");
161+
} else {
162+
if (existingVersionValue.get().n() == null) {
163+
// In this case a non-null version attribute is present, but it's not an N
164+
throw new IllegalArgumentException("Version attribute appears to be the wrong type. N is required.");
165+
}
166+
167+
String existingVersionValueKey = VERSIONED_RECORD_EXPRESSION_VALUE_KEY_MAPPER.apply(versionAttributeKey.get());
168+
condition = Expression.builder()
169+
.expression(String.format("%s = %s", attributeKeyRef, existingVersionValueKey))
170+
.expressionNames(Collections.singletonMap(attributeKeyRef, versionAttributeKey.get()))
171+
.expressionValues(Collections.singletonMap(existingVersionValueKey,
172+
existingVersionValue.get()))
173+
.build();
174+
}
175+
176+
return condition;
177+
}
178+
145179
@NotThreadSafe
146180
public static final class Builder {
147181
private Builder() {

‎services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/ChainExtension.java

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,4 +184,35 @@ public ReadModification afterRead(DynamoDbExtensionContext.AfterRead context) {
184184
.transformedItem(transformedItem)
185185
.build();
186186
}
187+
188+
/**
189+
* Implementation of the {@link DynamoDbEnhancedClientExtension} interface that will call all the chained extensions
190+
* in forward order, passing the results of each one to the next and coalescing the results into a single expression.
191+
* Multiple conditional statements will be separated by the string " AND ".
192+
*
193+
* @param context A {@link DynamoDbExtensionContext.BeforeDelete} context
194+
* @return A single {@link Expression} representing the coalesced results of all the chained extensions.
195+
*/
196+
@Override
197+
public Expression beforeDelete(DynamoDbExtensionContext.BeforeDelete context) {
198+
Expression conditionalExpression = null;
199+
200+
for (DynamoDbEnhancedClientExtension extension : this.extensionChain) {
201+
202+
DynamoDbExtensionContext.BeforeDelete beforeDelete =
203+
DefaultDynamoDbExtensionContext.builder()
204+
.items(context.items())
205+
.operationContext(context.operationContext())
206+
.tableMetadata(context.tableMetadata())
207+
.tableSchema(context.tableSchema())
208+
.build();
209+
210+
Expression additionalConditionExpression = extension.beforeDelete(beforeDelete);
211+
212+
conditionalExpression = mergeConditionalExpressions(conditionalExpression,
213+
additionalConditionExpression);
214+
}
215+
216+
return conditionalExpression;
217+
}
187218
}

‎services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/DefaultDynamoDbExtensionContext.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,8 @@
3232
*/
3333
@SdkInternalApi
3434
public final class DefaultDynamoDbExtensionContext implements DynamoDbExtensionContext.BeforeWrite,
35-
DynamoDbExtensionContext.AfterRead {
35+
DynamoDbExtensionContext.AfterRead,
36+
DynamoDbExtensionContext.BeforeDelete {
3637
private final Map<String, AttributeValue> items;
3738
private final OperationContext operationContext;
3839
private final TableMetadata tableMetadata;

0 commit comments

Comments
 (0)