Delete optimistic locking POC

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

@@ -13,9 +13,9 @@
  * permissions and limitations under the License.
  */
 
+
 package software.amazon.awssdk.enhanced.dynamodb;
 
-import static org.assertj.core.api.Assertions.as;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
@@ -31,6 +31,8 @@
 import software.amazon.awssdk.enhanced.dynamodb.model.PutItemEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.PutItemEnhancedResponse;
 import software.amazon.awssdk.enhanced.dynamodb.model.Record;
+import software.amazon.awssdk.enhanced.dynamodb.model.RecordWithVersion;
+import software.amazon.awssdk.enhanced.dynamodb.model.TransactWriteItemsEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedResponse;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
@@ -41,6 +43,7 @@
 import software.amazon.awssdk.services.dynamodb.model.ReturnItemCollectionMetrics;
 import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
 import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
+import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledException;
 
 public class AsyncCrudWithResponseIntegrationTest extends DynamoDbEnhancedIntegrationTestBase {
 
@@ -56,13 +59,15 @@ public class AsyncCrudWithResponseIntegrationTest extends DynamoDbEnhancedIntegr
     private static DynamoDbAsyncClient dynamoDbClient;
     private static DynamoDbEnhancedAsyncClient enhancedClient;
     private static DynamoDbAsyncTable<Record> mappedTable;
+    private static DynamoDbAsyncTable<RecordWithVersion> recordWithVersionMappedTable;
 
     @BeforeClass
     public static void beforeClass() {
         dynamoDbClient = createAsyncDynamoDbClient();
         enhancedClient = DynamoDbEnhancedAsyncClient.builder().dynamoDbClient(dynamoDbClient).build();
         mappedTable = enhancedClient.table(TABLE_NAME, TABLE_SCHEMA);
         mappedTable.createTable(r -> r.localSecondaryIndices(LOCAL_SECONDARY_INDEX)).join();
+        recordWithVersionMappedTable = enhancedClient.table(TABLE_NAME, RECORD_WITH_VERSION_TABLE_SCHEMA);
         dynamoDbClient.waiter().waitUntilTableExists(r -> r.tableName(TABLE_NAME)).join();
     }
 
@@ -341,4 +346,167 @@ public void getItem_withoutReturnConsumedCapacity() {
         GetItemEnhancedResponse<Record> response = mappedTable.getItemWithResponse(req -> req.key(key)).join();
         assertThat(response.consumedCapacity()).isNull();
     }
+
+    @Test
+    public void deleteItemWithoutVersion_andOptimisticLockingEnabled_shouldSucceed() {
+        Record originalItem = new Record().setId("123").setSort(10).setStringAttribute("Original Item");
+        Key recordKey = Key.builder()
+                           .partitionValue(originalItem.getId())
+                           .sortValue(originalItem.getSort())
+                           .build();
+        mappedTable.putItem(originalItem).join();
+
+        // Retrieve the item
+        Record retrievedItem = mappedTable.getItem(r -> r.key(recordKey)).join();
+
+        // Delete the item using a transaction
+        TransactWriteItemsEnhancedRequest request =
+            TransactWriteItemsEnhancedRequest.builder()
+                                             .addDeleteItem(mappedTable, retrievedItem)
+                                             .build();
+
+        enhancedClient.transactWriteItems(request).join();
+
+        Record deletedItem = mappedTable.getItem(r -> r.key(recordKey)).join();
+        assertThat(deletedItem).isNull();
+    }
+
+    @Test
+    public void deleteItemWithoutVersion_andOptimisticLockingDisabled_shouldSucceed() {
+        Record originalItem = new Record().setId("123").setSort(10).setStringAttribute("Original Item");
+        Key recordKey = Key.builder()
+                           .partitionValue(originalItem.getId())
+                           .sortValue(originalItem.getSort())
+                           .build();
+        mappedTable.putItem(originalItem).join();
+
+        // Retrieve the item
+        Record retrievedItem = mappedTable.getItem(r -> r.key(recordKey)).join();
+
+        // Delete the item using a transaction
+        TransactWriteItemsEnhancedRequest request =
+            TransactWriteItemsEnhancedRequest.builder()
+                                             .addDeleteItem(mappedTable, retrievedItem)
+                                             .build();
+
+        enhancedClient.transactWriteItems(request).join();
+
+        Record deletedItem = mappedTable.getItem(r -> r.key(recordKey)).join();
+        assertThat(deletedItem).isNull();
+    }
+
+    @Test
+    public void deleteItemWithVersion_andOptimisticLockingEnabled_ifVersionMatch_shouldSucceed() {
+        RecordWithVersion originalItem = new RecordWithVersion().setId("123").setSort(10).setStringAttribute("Original Item");
+        Key recordKey = Key.builder()
+                           .partitionValue(originalItem.getId())
+                           .sortValue(originalItem.getSort())
+                           .build();
+        recordWithVersionMappedTable.putItem(originalItem).join();
+
+        // Retrieve the item
+        RecordWithVersion retrievedItem = recordWithVersionMappedTable.getItem(r -> r.key(recordKey)).join();
+
+        // Delete the item using a transaction
+        TransactWriteItemsEnhancedRequest request =
+            TransactWriteItemsEnhancedRequest.builder()
+                                             .addDeleteItem(recordWithVersionMappedTable, retrievedItem)
+                                             .build();
+
+        enhancedClient.transactWriteItems(request).join();
+
+        RecordWithVersion deletedItem = recordWithVersionMappedTable.getItem(r -> r.key(recordKey)).join();
+        assertThat(deletedItem).isNull();
+    }
+
+    @Test
+    public void deleteItemWithVersion_andOptimisticLockingEnabled_ifVersionMismatch_shouldFail() {
+        RecordWithVersion originalItem = new RecordWithVersion().setId("123").setSort(10).setStringAttribute("Original Item");
+        Key recordKey = Key.builder()
+                           .partitionValue(originalItem.getId())
+                           .sortValue(originalItem.getSort())
+                           .build();
+
+        recordWithVersionMappedTable.putItem(originalItem).join();
+
+        // Retrieve the item and modify it separately
+        RecordWithVersion modifiedItem = recordWithVersionMappedTable.getItem(r -> r.key(recordKey)).join();
+        modifiedItem.setStringAttribute("Updated Item");
+
+        // Update the item, which will increment the version
+        recordWithVersionMappedTable.updateItem(modifiedItem);
+
+
+        //  Now attempt to delete the original item using a transaction
+        TransactWriteItemsEnhancedRequest request =
+            TransactWriteItemsEnhancedRequest.builder()
+                                             .addDeleteItem(recordWithVersionMappedTable, modifiedItem)
+                                             .build();
+
+        // enhancedClient.transactWriteItems(request).join();
+
+        assertThatThrownBy(() -> enhancedClient.transactWriteItems(request).join())
+            .isInstanceOf(CompletionException.class)
+            .satisfies(e ->
+                           assertThat(((TransactionCanceledException) e.getCause())
+                                          .cancellationReasons()
+                                          .stream()
+                                          .anyMatch(reason ->
+                                                        "ConditionalCheckFailed".equals(reason.code())
+                                                        && "The conditional request failed".equals(reason.message())))
+                               .isTrue());
+    }
+
+    @Test
+    public void deleteItemWithVersion_andOptimisticLockingDisabled_ifVersionMatch_shouldSucceed() {
+        RecordWithVersion originalItem = new RecordWithVersion().setId("123").setSort(10).setStringAttribute("Original Item");
+        Key recordKey = Key.builder()
+                           .partitionValue(originalItem.getId())
+                           .sortValue(originalItem.getSort())
+                           .build();
+        recordWithVersionMappedTable.putItem(originalItem).join();
+
+        // Retrieve the item
+        RecordWithVersion retrievedItem = recordWithVersionMappedTable.getItem(r -> r.key(recordKey)).join();
+
+        // Delete the item using a transaction
+        TransactWriteItemsEnhancedRequest request =
+            TransactWriteItemsEnhancedRequest.builder()
+                                             .addDeleteItem(recordWithVersionMappedTable, retrievedItem)
+                                             .build();
+
+        enhancedClient.transactWriteItems(request).join();
+
+        RecordWithVersion deletedItem = recordWithVersionMappedTable.getItem(r -> r.key(recordKey)).join();
+        assertThat(deletedItem).isNull();
+    }
+
+    @Test
+    public void deleteItemWithVersion_andOptimisticLockingDisabled_ifVersionMismatch_shouldSucceed() {
+        RecordWithVersion originalItem = new RecordWithVersion().setId("123").setSort(10).setStringAttribute("Original Item");
+        Key recordKey = Key.builder()
+                           .partitionValue(originalItem.getId())
+                           .sortValue(originalItem.getSort())
+                           .build();
+
+        recordWithVersionMappedTable.putItem(originalItem).join();
+
+        // Retrieve the item and modify it separately
+        RecordWithVersion modifiedItem = recordWithVersionMappedTable.getItem(r -> r.key(recordKey)).join();
+        modifiedItem.setStringAttribute("Updated Item");
+
+        // Update the item, which will increment the version
+        recordWithVersionMappedTable.updateItem(modifiedItem);
+
+        //  Now attempt to delete the original item using a transaction
+        TransactWriteItemsEnhancedRequest request =
+            TransactWriteItemsEnhancedRequest.builder()
+                                             .addDeleteItem(recordWithVersionMappedTable, modifiedItem)
+                                             .build();
+
+        enhancedClient.transactWriteItems(request).join();
+
+        RecordWithVersion deletedItem = recordWithVersionMappedTable.getItem(r -> r.key(recordKey)).join();
+        assertThat(deletedItem).isNull();
+    }
 }

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

@@ -19,13 +19,15 @@
 import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.primarySortKey;
 import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.secondaryPartitionKey;
 import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.secondarySortKey;
+import static software.amazon.awssdk.enhanced.dynamodb.extensions.VersionedRecordExtension.AttributeTags.versionAttribute;
 
 import java.util.Arrays;
 import java.util.List;
 import java.util.UUID;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableSchema;
+import software.amazon.awssdk.enhanced.dynamodb.model.RecordWithVersion;
 import software.amazon.awssdk.enhanced.dynamodb.model.Record;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
@@ -102,4 +104,34 @@ protected static String getStringAttrValue(int numChars) {
         return new String(chars);
     }
 
+    protected static final TableSchema<RecordWithVersion> RECORD_WITH_VERSION_TABLE_SCHEMA =
+        StaticTableSchema.builder(RecordWithVersion.class)
+                         .newItemSupplier(RecordWithVersion::new)
+                         .addAttribute(String.class, a -> a.name("id")
+                                                           .getter(RecordWithVersion::getId)
+                                                           .setter(RecordWithVersion::setId)
+                                                           .tags(primaryPartitionKey(), secondaryPartitionKey("index1")))
+                         .addAttribute(Integer.class, a -> a.name("sort")
+                                                            .getter(RecordWithVersion::getSort)
+                                                            .setter(RecordWithVersion::setSort)
+                                                            .tags(primarySortKey(), secondarySortKey("index1")))
+                         .addAttribute(Integer.class, a -> a.name("value")
+                                                            .getter(RecordWithVersion::getValue)
+                                                            .setter(RecordWithVersion::setValue))
+                         .addAttribute(String.class, a -> a.name("gsi_id")
+                                                           .getter(RecordWithVersion::getGsiId)
+                                                           .setter(RecordWithVersion::setGsiId)
+                                                           .tags(secondaryPartitionKey("gsi_keys_only")))
+                         .addAttribute(Integer.class, a -> a.name("gsi_sort")
+                                                            .getter(RecordWithVersion::getGsiSort)
+                                                            .setter(RecordWithVersion::setGsiSort)
+                                                            .tags(secondarySortKey("gsi_keys_only")))
+                         .addAttribute(String.class, a -> a.name("stringAttribute")
+                                                           .getter(RecordWithVersion::getStringAttribute)
+                                                           .setter(RecordWithVersion::setStringAttribute))
+                         .addAttribute(Integer.class, a -> a.name("version")
+                                                            .getter(RecordWithVersion::getVersion)
+                                                            .setter(RecordWithVersion::setVersion)
+                                                            .tags(versionAttribute(0L, 1L, true)))  // startAt=0, incrementBy=1,
+                         .build();
 }

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

@@ -0,0 +1,119 @@
+/*
+ * 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
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.enhanced.dynamodb.model;
+
+import java.util.Objects;
+import software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbVersionAttribute;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbBean;
+
+@DynamoDbBean
+public class RecordWithVersion {
+
+    private String id;
+    private Integer sort;
+    private Integer value;
+    private String gsiId;
+    private Integer gsiSort;
+    private String stringAttribute;
+    private Integer version;
+
+    public String getId() {
+        return id;
+    }
+
+    public RecordWithVersion setId(String id) {
+        this.id = id;
+        return this;
+    }
+
+    public Integer getSort() {
+        return sort;
+    }
+
+    public RecordWithVersion setSort(Integer sort) {
+        this.sort = sort;
+        return this;
+    }
+
+    public Integer getValue() {
+        return value;
+    }
+
+    public RecordWithVersion setValue(Integer value) {
+        this.value = value;
+        return this;
+    }
+
+    public String getGsiId() {
+        return gsiId;
+    }
+
+    public RecordWithVersion setGsiId(String gsiId) {
+        this.gsiId = gsiId;
+        return this;
+    }
+
+    public Integer getGsiSort() {
+        return gsiSort;
+    }
+
+    public RecordWithVersion setGsiSort(Integer gsiSort) {
+        this.gsiSort = gsiSort;
+        return this;
+    }
+
+    public String getStringAttribute() {
+        return stringAttribute;
+    }
+
+    public RecordWithVersion setStringAttribute(String stringAttribute) {
+        this.stringAttribute = stringAttribute;
+        return this;
+    }
+
+    @DynamoDbVersionAttribute
+    public Integer getVersion() {
+        return version;
+    }
+
+    public RecordWithVersion setVersion(Integer version) {
+        this.version = version;
+        return this;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        RecordWithVersion recordWithVersion = (RecordWithVersion) o;
+        return Objects.equals(id, recordWithVersion.id) &&
+               Objects.equals(sort, recordWithVersion.sort) &&
+               Objects.equals(value, recordWithVersion.value) &&
+               Objects.equals(gsiId, recordWithVersion.gsiId) &&
+               Objects.equals(stringAttribute, recordWithVersion.stringAttribute) &&
+               Objects.equals(gsiSort, recordWithVersion.gsiSort) &&
+               Objects.equals(version, recordWithVersion.version);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(id, sort, value, gsiId, gsiSort, stringAttribute, version);
+    }
+}