Optimistic locking for delete scenario with DeleteItemEnhancedRequest and TransactWriteItemsEnhancedRequest

Author: anasatirbasa

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

@@ -0,0 +1,6 @@
+{
+  "type": "bugfix",
+  "category": "Amazon DynamoDB Enhanced Client",
+  "contributor": "",
+  "description": "Optimistic delete while using DynamoDbEnhancedClient - DeleteItem with DeleteItemEnhancedRequest and TransactWriteItemsEnhancedRequest"
+}

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

@@ -15,6 +15,7 @@
 
 package software.amazon.awssdk.enhanced.dynamodb;
 
+import static software.amazon.awssdk.enhanced.dynamodb.extensions.VersionedRecordExtension.AttributeTags.versionAttribute;
 import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.primaryPartitionKey;
 import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.primarySortKey;
 import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.secondaryPartitionKey;
@@ -27,6 +28,7 @@
 import java.util.stream.IntStream;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableSchema;
 import software.amazon.awssdk.enhanced.dynamodb.model.Record;
+import software.amazon.awssdk.enhanced.dynamodb.model.VersionedRecord;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.testutils.service.AwsIntegrationTestBase;
@@ -75,6 +77,37 @@ protected static DynamoDbAsyncClient createAsyncDynamoDbClient() {
                                                            .setter(Record::setStringAttribute))
                          .build();
 
+    protected static final TableSchema<VersionedRecord> VERSIONED_RECORD_TABLE_SCHEMA =
+        StaticTableSchema.builder(VersionedRecord.class)
+                         .newItemSupplier(VersionedRecord::new)
+                         .addAttribute(String.class, a -> a.name("id")
+                                                           .getter(VersionedRecord::getId)
+                                                           .setter(VersionedRecord::setId)
+                                                           .tags(primaryPartitionKey(), secondaryPartitionKey("index1")))
+                         .addAttribute(Integer.class, a -> a.name("sort")
+                                                            .getter(VersionedRecord::getSort)
+                                                            .setter(VersionedRecord::setSort)
+                                                            .tags(primarySortKey(), secondarySortKey("index1")))
+                         .addAttribute(Integer.class, a -> a.name("value")
+                                                            .getter(VersionedRecord::getValue)
+                                                            .setter(VersionedRecord::setValue))
+                         .addAttribute(String.class, a -> a.name("gsi_id")
+                                                           .getter(VersionedRecord::getGsiId)
+                                                           .setter(VersionedRecord::setGsiId)
+                                                           .tags(secondaryPartitionKey("gsi_keys_only")))
+                         .addAttribute(Integer.class, a -> a.name("gsi_sort")
+                                                            .getter(VersionedRecord::getGsiSort)
+                                                            .setter(VersionedRecord::setGsiSort)
+                                                            .tags(secondarySortKey("gsi_keys_only")))
+                         .addAttribute(String.class, a -> a.name("stringAttribute")
+                                                           .getter(VersionedRecord::getStringAttribute)
+                                                           .setter(VersionedRecord::setStringAttribute))
+                         .addAttribute(Integer.class, a -> a.name("version")
+                                                            .getter(VersionedRecord::getVersion)
+                                                            .setter(VersionedRecord::setVersion)
+                                                            .tags(versionAttribute()))
+                         .build();
+
 
     protected static final List<Record> RECORDS =
         IntStream.range(0, 9)

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

@@ -0,0 +1,120 @@
+/*
+ * 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 VersionedRecord {
+
+    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 VersionedRecord setId(String id) {
+        this.id = id;
+        return this;
+    }
+
+    public Integer getSort() {
+        return sort;
+    }
+
+    public VersionedRecord setSort(Integer sort) {
+        this.sort = sort;
+        return this;
+    }
+
+    public Integer getValue() {
+        return value;
+    }
+
+    public VersionedRecord setValue(Integer value) {
+        this.value = value;
+        return this;
+    }
+
+    public String getGsiId() {
+        return gsiId;
+    }
+
+    public VersionedRecord setGsiId(String gsiId) {
+        this.gsiId = gsiId;
+        return this;
+    }
+
+    public Integer getGsiSort() {
+        return gsiSort;
+    }
+
+    public VersionedRecord setGsiSort(Integer gsiSort) {
+        this.gsiSort = gsiSort;
+        return this;
+    }
+
+    public String getStringAttribute() {
+        return stringAttribute;
+    }
+
+    public VersionedRecord setStringAttribute(String stringAttribute) {
+        this.stringAttribute = stringAttribute;
+        return this;
+    }
+
+    @DynamoDbVersionAttribute
+    public Integer getVersion() {
+        return version;
+    }
+
+    public VersionedRecord 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;
+        }
+        VersionedRecord versionedRecord = (VersionedRecord) o;
+        return Objects.equals(id, versionedRecord.id) &&
+               Objects.equals(sort, versionedRecord.sort) &&
+               Objects.equals(value, versionedRecord.value) &&
+               Objects.equals(gsiId, versionedRecord.gsiId) &&
+               Objects.equals(stringAttribute, versionedRecord.stringAttribute) &&
+               Objects.equals(gsiSort, versionedRecord.gsiSort) &&
+               Objects.equals(version, versionedRecord.version);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(id, sort, value, gsiId, gsiSort, stringAttribute, version);
+    }
+}

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

@@ -31,6 +31,7 @@
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbExtensionContext;
 import software.amazon.awssdk.enhanced.dynamodb.Expression;
+import software.amazon.awssdk.enhanced.dynamodb.internal.operations.OperationName;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
@@ -64,8 +65,9 @@ public final class VersionedRecordExtension implements DynamoDbEnhancedClientExt
 
     private final long startAt;
     private final long incrementBy;
+    private final boolean optimisticLockingOnDelete;
 
-    private VersionedRecordExtension(Long startAt, Long incrementBy) {
+    private VersionedRecordExtension(Long startAt, Long incrementBy, Boolean optimisticLockingOnDelete) {
         Validate.isNotNegativeOrNull(startAt, "startAt");
 
         if (incrementBy != null && incrementBy < 1) {
@@ -74,6 +76,7 @@ private VersionedRecordExtension(Long startAt, Long incrementBy) {
 
         this.startAt = startAt != null ? startAt : 0L;
         this.incrementBy = incrementBy != null ? incrementBy : 1L;
+        this.optimisticLockingOnDelete = optimisticLockingOnDelete != null && optimisticLockingOnDelete;
     }
 
     public static Builder builder() {
@@ -91,23 +94,37 @@ public static StaticAttributeTag versionAttribute() {
         public static StaticAttributeTag versionAttribute(Long startAt, Long incrementBy) {
             return new VersionAttribute(startAt, incrementBy);
         }
+
+        public static StaticAttributeTag versionAttribute(Long startAt, Long incrementBy, Boolean optimisticLockingOnDelete) {
+            return new VersionAttribute(startAt, incrementBy, optimisticLockingOnDelete);
+        }
     }
 
     private static final class VersionAttribute implements StaticAttributeTag {
         private static final String START_AT_METADATA_KEY = "VersionedRecordExtension:StartAt";
         private static final String INCREMENT_BY_METADATA_KEY = "VersionedRecordExtension:IncrementBy";
+        private static final String OPTIMISTIC_LOCKING_ON_DELETE_METADATA_KEY = "VersionedRecordExtension:OptimisticLockingOnDelete";
 
         private final Long startAt;
         private final Long incrementBy;
+        private final Boolean optimisticLockingOnDelete;
 
         private VersionAttribute() {
             this.startAt = null;
             this.incrementBy = null;
+            this.optimisticLockingOnDelete = null;
         }
 
         private VersionAttribute(Long startAt, Long incrementBy) {
             this.startAt = startAt;
             this.incrementBy = incrementBy;
+            this.optimisticLockingOnDelete = null;
+        }
+
+        private VersionAttribute(Long startAt, Long incrementBy, Boolean optimisticLockingOnDelete) {
+            this.startAt = startAt;
+            this.incrementBy = incrementBy;
+            this.optimisticLockingOnDelete = optimisticLockingOnDelete;
         }
 
         @Override
@@ -128,6 +145,7 @@ public Consumer<StaticTableMetadata.Builder> modifyMetadata(String attributeName
             return metadata -> metadata.addCustomMetadataObject(CUSTOM_METADATA_KEY, attributeName)
                                        .addCustomMetadataObject(START_AT_METADATA_KEY, startAt)
                                        .addCustomMetadataObject(INCREMENT_BY_METADATA_KEY, incrementBy)
+                                       .addCustomMetadataObject(OPTIMISTIC_LOCKING_ON_DELETE_METADATA_KEY, optimisticLockingOnDelete)
                                        .markAttributeAsKey(attributeName, attributeValueType);
         }
     }
@@ -141,6 +159,27 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
             return WriteModification.builder().build();
         }
 
+        // Check if optimistic locking is enabled for delete operations
+        // First check attribute-level setting, then fall back to extension-level setting
+        Boolean attributeLevelOptimisticLocking = context.tableMetadata()
+                                                         .customMetadataObject(VersionAttribute.OPTIMISTIC_LOCKING_ON_DELETE_METADATA_KEY, Boolean.class)
+                                                         .orElse(null);
+
+        boolean shouldApplyOptimisticLocking = attributeLevelOptimisticLocking != null
+                                               ? attributeLevelOptimisticLocking
+                                               : this.optimisticLockingOnDelete;
+
+        // Handle DELETE operations with optimistic locking if enabled
+        if (context.operationName() == OperationName.DELETE_ITEM && shouldApplyOptimisticLocking) {
+            return handleOptimisticDelete(context, versionAttributeKey.get());
+        }
+
+        // For non-delete operations, skip version handling if it's a delete
+        if (context.operationName() == OperationName.DELETE_ITEM) {
+            return WriteModification.builder().build();
+        }
+
+        // Existing logic for other operations
         Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
 
         String attributeKeyRef = keyRef(versionAttributeKey.get());
@@ -204,6 +243,30 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
                                 .build();
     }
 
+    private WriteModification handleOptimisticDelete(DynamoDbExtensionContext.BeforeWrite context, String versionAttributeKey) {
+        // Look for version in the items map
+        AttributeValue versionValue = context.items().get(versionAttributeKey);
+
+        if (versionValue != null && versionValue.n() != null) {
+            // Build condition for the specific version
+            String attributeKeyRef = keyRef(versionAttributeKey);
+            String valueKey = VERSIONED_RECORD_EXPRESSION_VALUE_KEY_MAPPER.apply(versionAttributeKey);
+
+            Expression condition = Expression.builder()
+                                             .expression(String.format("%s = %s", attributeKeyRef, valueKey))
+                                             .expressionNames(Collections.singletonMap(attributeKeyRef, versionAttributeKey))
+                                             .expressionValues(Collections.singletonMap(valueKey, versionValue))
+                                             .build();
+
+            return WriteModification.builder()
+                                    .additionalConditionalExpression(condition)
+                                    .build();
+        }
+
+        // If no version value is provided, don't add any condition (backward compatible)
+        return WriteModification.builder().build();
+    }
+
     private boolean isInitialVersion(AttributeValue existingVersionValue, Long versionStartAtFromAnnotation) {
         if (existingVersionValue == null || isNullAttributeValue(existingVersionValue)) {
             return true;
@@ -223,6 +286,7 @@ private boolean isInitialVersion(AttributeValue existingVersionValue, Long versi
     public static final class Builder {
         private Long startAt;
         private Long incrementBy;
+        private Boolean optimisticLockingOnDelete;
 
         private Builder() {
         }
@@ -251,8 +315,20 @@ public Builder incrementBy(Long incrementBy) {
             return this;
         }
 
+        /**
+         * Enables or disables optimistic locking for delete operations. When enabled, delete operations will include a condition
+         * to check the version attribute. Default value - {@code false} (for backward compatibility).
+         *
+         * @param optimisticLockingOnDelete true to enable optimistic locking on deletes, false to disable
+         * @return the builder instance
+         */
+        public Builder optimisticLockingOnDelete(Boolean optimisticLockingOnDelete) {
+            this.optimisticLockingOnDelete = optimisticLockingOnDelete;
+            return this;
+        }
+
         public VersionedRecordExtension build() {
-            return new VersionedRecordExtension(this.startAt, this.incrementBy);
+            return new VersionedRecordExtension(this.startAt, this.incrementBy, this.optimisticLockingOnDelete);
         }
     }
 }

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

@@ -49,4 +49,11 @@
      */
     long incrementBy() default 1;
 
+    /**
+     * Whether to enable optimistic locking for delete operations. When enabled, delete operations will include a condition to
+     * check the version attribute. Default value - {@code false} (for backward compatibility).
+     *
+     * @return true to enable optimistic locking on deletes, false to disable
+     */
+    boolean optimisticLockingOnDelete() default false;
 }

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

@@ -145,7 +145,9 @@ public CompletableFuture<T> deleteItem(Key key) {
 
     @Override
     public CompletableFuture<T> deleteItem(T keyItem) {
-        return deleteItem(keyFrom(keyItem));
+        return deleteItem(DeleteItemEnhancedRequest.builder()
+                                                   .key(keyFrom(keyItem))
+                                                   .items(tableSchema.itemToMap(keyItem, true)).build());
     }
 
     @Override

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/extensions/annotations/DynamoDbVersionAttribute.java

@@ -146,7 +146,9 @@ public T deleteItem(Key key) {
 
     @Override
     public T deleteItem(T keyItem) {
-        return deleteItem(keyFrom(keyItem));
+        return deleteItem(DeleteItemEnhancedRequest.builder()
+                                                   .key(keyFrom(keyItem))
+                                                   .items(tableSchema.itemToMap(keyItem, true)).build());
     }
 
     @Override

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/client/DefaultDynamoDbAsyncTable.java

@@ -26,6 +26,10 @@ private VersionRecordAttributeTags() {
     }
 
     public static StaticAttributeTag attributeTagFor(DynamoDbVersionAttribute annotation) {
-        return VersionedRecordExtension.AttributeTags.versionAttribute(annotation.startAt(), annotation.incrementBy());
+        return VersionedRecordExtension.AttributeTags.versionAttribute(
+            annotation.startAt(),
+            annotation.incrementBy(),
+            annotation.optimisticLockingOnDelete()
+        );
     }
 }