DynamoDb enhanced client: support UpdateExpressions in single-request update

Author: anasatirbasa

services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/UpdateExpressionTest.java

@@ -3,7 +3,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.keyRef;
-import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.valueRef;
+import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.primaryPartitionKey;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -305,9 +305,20 @@ public void updateExpressionInRequest_whenAttributeAlsoInPojo_shouldThrowConflic
     @Test
     public void updateExpressionInRequest_whenAttributeAlsoInExtension_shouldThrowDynamoDbError() {
         initClientWithExtensions(new ItemPreservingUpdateExtension());
-        RecordForUpdateExpressions recordForUpdateExpressions = createFullRecord();
+        RecordForUpdateExpressions recordForUpdateExpressions = createKeyOnlyRecord();
+
+        // Create an UpdateExpression that conflicts with the extension's UpdateExpression
+        // Extension modifies extensionNumberAttribute, so we create a request expression that also modifies it
+        UpdateExpression conflictingExpression = UpdateExpression.builder()
+                                                                 .addAction(SetAction.builder()
+                                                                                     .path("extensionNumberAttribute")
+                                                                                     .value(":conflictValue")
+                                                                                     .putExpressionValue(":conflictValue", AttributeValue.builder().n("99").build())
+                                                                                     .build())
+                                                                 .build();
+
         assertThatThrownBy(() -> mappedTable.updateItem(r -> r.item(recordForUpdateExpressions)
-                                                              .updateExpression(expressionWithSetExtensionAttribute())))
+                                                              .updateExpression(conflictingExpression)))
             .isInstanceOf(DynamoDbException.class)
             .hasMessageContaining("Two document paths")
             .hasMessageContaining(NUMBER_ATTRIBUTE_REF);
@@ -560,7 +571,7 @@ public void transactWriteItems_withUpdateExpression() {
                                              .addPutItem(mappedTable, record2)
                                              .build());
 
-        // Verify both operations succeeded  
+        // Verify both operations succeeded
         RecordForUpdateExpressions deletedRecord = mappedTable.getItem(record1);
         assertThat(deletedRecord).isNull();
 
@@ -569,6 +580,51 @@ public void transactWriteItems_withUpdateExpression() {
         assertThat(persistedRecord2.getRequestAttributeList()).containsExactly("attr1", "attr2");
     }
 
+    /**
+     * Tests StaticTableSchema with UpdateExpression extensions
+     */
+    @Test
+    public void staticTableSchema_withUpdateExpressions() {
+        TableSchema<RecordForUpdateExpressions> staticSchema = TableSchema.builder(RecordForUpdateExpressions.class)
+                                                                          .newItemSupplier(RecordForUpdateExpressions::new)
+                                                                          .addAttribute(String.class, a -> a.name("id")
+                                                                                                            .getter(RecordForUpdateExpressions::getId)
+                                                                                                            .setter(RecordForUpdateExpressions::setId)
+                                                                                                            .tags(primaryPartitionKey()))
+                                                                          .addAttribute(String.class, a -> a.name("stringAttribute")
+                                                                                                            .getter(RecordForUpdateExpressions::getStringAttribute)
+                                                                                                            .setter(RecordForUpdateExpressions::setStringAttribute))
+                                                                          .addAttribute(Long.class, a -> a.name("extensionNumberAttribute")
+                                                                                                          .getter(RecordForUpdateExpressions::getExtensionNumberAttribute)
+                                                                                                          .setter(RecordForUpdateExpressions::setExtensionNumberAttribute))
+                                                                          .build();
+
+        DynamoDbEnhancedClient enhancedClient = DynamoDbEnhancedClient.builder()
+                                                                      .dynamoDbClient(getDynamoDbClient())
+                                                                      .extensions(new ItemPreservingUpdateExtension())
+                                                                      .build();
+
+        String staticTableName = getConcreteTableName("static-table");
+        DynamoDbTable<RecordForUpdateExpressions> staticTable = enhancedClient.table(staticTableName, staticSchema);
+
+        try {
+            staticTable.createTable(r -> r.provisionedThroughput(getDefaultProvisionedThroughput()));
+
+            RecordForUpdateExpressions record = new RecordForUpdateExpressions();
+            record.setId("static-test");
+            record.setStringAttribute("init");
+            // Don't set requestAttributeList to avoid path conflicts with extension
+
+            staticTable.updateItem(r -> r.item(record));
+
+            RecordForUpdateExpressions persistedRecord = staticTable.getItem(record);
+            assertThat(persistedRecord.getStringAttribute()).isEqualTo("init");
+            assertThat(persistedRecord.getExtensionNumberAttribute()).isEqualTo(5L);
+        } finally {
+            getDynamoDbClient().deleteTable(r -> r.tableName(staticTableName));
+        }
+    }
+
     private void verifyDDBError(RecordForUpdateExpressions record, boolean ignoreNulls) {
         assertThatThrownBy(() -> mappedTable.updateItem(r -> r.item(record).ignoreNulls(ignoreNulls)))
             .isInstanceOf(DynamoDbException.class)
@@ -628,18 +684,6 @@ private UpdateExpression expressionWithSetListElement(int index, String value) {
         return UpdateExpression.builder().addAction(setListElement).build();
     }
 
-    private UpdateExpression expressionWithSetExtensionAttribute() {
-        String attributeName = "extensionNumberAttribute";
-        AttributeValue elementValue = AttributeValue.builder().n("11").build();
-        SetAction setAttribute = SetAction.builder()
-                                          .path(keyRef(attributeName))
-                                          .value(valueRef(attributeName))
-                                          .putExpressionValue(valueRef(attributeName), elementValue)
-                                          .putExpressionName(keyRef(attributeName), attributeName)
-                                          .build();
-        return UpdateExpression.builder().addAction(setAttribute).build();
-    }
-
     private static final class ItemPreservingUpdateExtension implements DynamoDbEnhancedClientExtension {
         private long incrementValue;
         private String valueRef;