Fixed EnhancedClient UpdateItem operation to make it work on nested a…

…ttributes as well
aws · May 30, 2024 · df37a55 · df37a55
1 parent 7424318
commit df37a55
Show file tree

Hide file tree

Showing 9 changed files with 131 additions and 18 deletions.
diff --git a/...dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/TableSchema.java b/...dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/TableSchema.java
@@ -244,6 +244,8 @@ default T mapToItem(Map<String, AttributeValue> attributeMap, boolean preserveEm
      */
     Map<String, AttributeValue> itemToMap(T item, boolean ignoreNulls);
 
+    Map<String, AttributeValue> updateItemToMap(T item, boolean ignoreNulls);
+
     /**
      * Takes a modelled object and extracts a specific set of attributes which are then returned as a map of
      * {@link AttributeValue} that the DynamoDb low-level SDK can work with. This method is typically used to extract

diff --git a/.../src/main/java/software/amazon/awssdk/enhanced/dynamodb/document/DocumentTableSchema.java b/.../src/main/java/software/amazon/awssdk/enhanced/dynamodb/document/DocumentTableSchema.java
@@ -113,6 +113,11 @@ public Map<String, AttributeValue> itemToMap(EnhancedDocument item, boolean igno
         return item.toBuilder().attributeConverterProviders(providers).build().toMap();
     }
 
+    @Override
+    public Map<String, AttributeValue> updateItemToMap(EnhancedDocument item, boolean ignoreNulls) {
+        return itemToMap(item, ignoreNulls);
+    }
+
     private List<AttributeConverterProvider> mergeAttributeConverterProviders(EnhancedDocument item) {
         if (item.attributeConverterProviders() != null && !item.attributeConverterProviders().isEmpty()) {
             Set<AttributeConverterProvider> providers = new LinkedHashSet<>();

diff --git a/.../src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/EnhancedClientUtils.java b/.../src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/EnhancedClientUtils.java
@@ -15,13 +15,16 @@
 
 package software.amazon.awssdk.enhanced.dynamodb.internal;
 
+import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticImmutableTableSchema.NESTED_OBJECT_UPDATE;
+
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.function.Supplier;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import software.amazon.awssdk.annotations.SdkInternalApi;
@@ -40,6 +43,7 @@ public final class EnhancedClientUtils {
     private static final Set<Character> SPECIAL_CHARACTERS = Stream.of(
         '*', '.', '-', '#', '+', ':', '/', '(', ')', ' ',
         '&', '<', '>', '?', '=', '!', '@', '%', '$', '|').collect(Collectors.toSet());
+    private static final Pattern PATTERN = Pattern.compile(NESTED_OBJECT_UPDATE);
 
     private EnhancedClientUtils() {
 
@@ -67,18 +71,30 @@ public static String cleanAttributeName(String key) {
         return somethingChanged ? new String(chars) : key;
     }
 
+    private static boolean isNestedAttribute(String key) {
+        return key.contains(NESTED_OBJECT_UPDATE);
+    }
+
     /**
      * Creates a key token to be used with an ExpressionNames map.
      */
     public static String keyRef(String key) {
-        return "#AMZN_MAPPED_" + cleanAttributeName(key);
+        String cleanAttributeName = cleanAttributeName(key);
+        cleanAttributeName = isNestedAttribute(cleanAttributeName) ?
+                             PATTERN.matcher(cleanAttributeName).replaceAll(".#AMZN_MAPPED_")
+                                                    : cleanAttributeName;
+        return "#AMZN_MAPPED_" + cleanAttributeName;
     }
 
     /**
      * Creates a value token to be used with an ExpressionValues map.
      */
     public static String valueRef(String value) {
-        return ":AMZN_MAPPED_" + cleanAttributeName(value);
+        String cleanAttributeName = cleanAttributeName(value);
+        cleanAttributeName = isNestedAttribute(cleanAttributeName) ?
+                             PATTERN.matcher(cleanAttributeName).replaceAll("_")
+                                                      : cleanAttributeName;
+        return ":AMZN_MAPPED_" + cleanAttributeName;
     }
 
     public static <T> T readAndTransformSingleItem(Map<String, AttributeValue> itemMap,

diff --git a/...c/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/mapper/MetaTableSchema.java b/...c/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/mapper/MetaTableSchema.java
@@ -54,6 +54,11 @@ public Map<String, AttributeValue> itemToMap(T item, boolean ignoreNulls) {
         return concreteTableSchema().itemToMap(item, ignoreNulls);
     }
 
+    @Override
+    public Map<String, AttributeValue> updateItemToMap(T item, boolean ignoreNulls) {
+        return concreteTableSchema().updateItemToMap(item, ignoreNulls);
+    }
+
     @Override
     public Map<String, AttributeValue> itemToMap(T item, Collection<String> attributes) {
         return concreteTableSchema().itemToMap(item, attributes);

diff --git a/...ava/software/amazon/awssdk/enhanced/dynamodb/internal/operations/UpdateItemOperation.java b/...ava/software/amazon/awssdk/enhanced/dynamodb/internal/operations/UpdateItemOperation.java
@@ -90,7 +90,7 @@ public UpdateItemRequest generateRequest(TableSchema<T> tableSchema,
                                           r -> Optional.ofNullable(r.ignoreNulls()))
                                      .orElse(null);
 
-        Map<String, AttributeValue> itemMap = tableSchema.itemToMap(item, Boolean.TRUE.equals(ignoreNulls));
+        Map<String, AttributeValue> itemMap = tableSchema.updateItemToMap(item, Boolean.TRUE.equals(ignoreNulls));
         TableMetadata tableMetadata = tableSchema.tableMetadata();
 
         WriteModification transformation =

diff --git a/.../java/software/amazon/awssdk/enhanced/dynamodb/internal/update/UpdateExpressionUtils.java b/.../java/software/amazon/awssdk/enhanced/dynamodb/internal/update/UpdateExpressionUtils.java
@@ -18,13 +18,16 @@
 import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.isNullAttributeValue;
 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.StaticImmutableTableSchema.NESTED_OBJECT_UPDATE;
 import static software.amazon.awssdk.utils.CollectionUtils.filterMap;
 
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
@@ -39,6 +42,8 @@
 @SdkInternalApi
 public final class UpdateExpressionUtils {
 
+    private static final Pattern PATTERN = Pattern.compile(NESTED_OBJECT_UPDATE);
+
     private UpdateExpressionUtils() {
     }
 
@@ -99,7 +104,7 @@ private static List<RemoveAction> removeActionsFor(Map<String, AttributeValue> a
     private static RemoveAction remove(String attributeName) {
         return RemoveAction.builder()
                            .path(keyRef(attributeName))
-                           .expressionNames(Collections.singletonMap(keyRef(attributeName), attributeName))
+                           .expressionNames(expressionNamesFor(attributeName))
                            .build();
     }
 
@@ -137,8 +142,18 @@ private static Function<String, String> behaviorBasedValue(UpdateBehavior update
      * Simple utility method that can create an ExpressionNames map based on a list of attribute names.
      */
     private static Map<String, String> expressionNamesFor(String... attributeNames) {
-        return Arrays.stream(attributeNames)
-                     .collect(Collectors.toMap(EnhancedClientUtils::keyRef, Function.identity()));
+        Map<String, String> map = new HashMap<>();
+        for (String attributeName : attributeNames) {
+            if (attributeName.contains(NESTED_OBJECT_UPDATE)) {
+                for (String attribute : PATTERN.split(attributeName)) {
+                    map.put(keyRef(attribute), attribute);
+                }
+            }
+        }
+        if (map.isEmpty()) {
+            return Arrays.stream(attributeNames)
+                         .collect(Collectors.toMap(EnhancedClientUtils::keyRef, Function.identity()));
+        }
+        return map;
     }
-
 }
diff --git a/...main/java/software/amazon/awssdk/enhanced/dynamodb/mapper/StaticImmutableTableSchema.java b/...main/java/software/amazon/awssdk/enhanced/dynamodb/mapper/StaticImmutableTableSchema.java
@@ -78,6 +78,7 @@
 @SdkPublicApi
 @ThreadSafe
 public final class StaticImmutableTableSchema<T, B> implements TableSchema<T> {
+    public static final String NESTED_OBJECT_UPDATE = "_NESTED_OBJECT_UPDATE_";
     private final List<ResolvedImmutableAttribute<T, B>> attributeMappers;
     private final Supplier<B> newBuilderSupplier;
     private final Function<B, T> buildItemFunction;
@@ -529,6 +530,40 @@ public Map<String, AttributeValue> itemToMap(T item, boolean ignoreNulls) {
         return unmodifiableMap(attributeValueMap);
     }
 
+    @Override
+    public Map<String, AttributeValue> updateItemToMap(T item, boolean ignoreNulls) {
+        Map<String, AttributeValue> attributeValueMap = new HashMap<>();
+
+        attributeMappers.forEach(attributeMapper -> {
+            String attributeKey = attributeMapper.attributeName();
+            AttributeValue attributeValue = attributeMapper.attributeGetterMethod().apply(item);
+
+            if (!ignoreNulls || !isNullAttributeValue(attributeValue)) {
+                if (attributeValue.hasM()) {
+                    nestedUpdateAttributeMapper(attributeValueMap, attributeValue.m(), attributeKey, ignoreNulls);
+                } else {
+                    attributeValueMap.put(attributeKey, attributeValue);
+                }
+            }
+        });
+
+        indexedFlattenedMappers.forEach((name, flattenedMapper) -> {
+            attributeValueMap.putAll(flattenedMapper.itemToMap(item, ignoreNulls));
+        });
+
+        return unmodifiableMap(attributeValueMap);
+    }
+
+    public void nestedUpdateAttributeMapper(Map<String, AttributeValue> attributeValueMap,
+                                            Map<String, AttributeValue> updateItemAttributeMap, String attributeKey,
+                                            boolean ignoreNulls) {
+        updateItemAttributeMap.forEach((mapKey, mapValue) -> {
+            if (!ignoreNulls || !isNullAttributeValue(mapValue)) {
+                attributeValueMap.put(attributeKey + NESTED_OBJECT_UPDATE + mapKey, mapValue);
+            }
+        });
+    }
+
     @Override
     public Map<String, AttributeValue> itemToMap(T item, Collection<String> attributes) {
         Map<String, AttributeValue> attributeValueMap = new HashMap<>();

diff --git a/...ced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/mapper/WrappedTableSchema.java b/...ced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/mapper/WrappedTableSchema.java
@@ -67,6 +67,11 @@ public Map<String, AttributeValue> itemToMap(T item, boolean ignoreNulls) {
         return this.delegateTableSchema.itemToMap(item, ignoreNulls);
     }
 
+    @Override
+    public Map<String, AttributeValue> updateItemToMap(T item, boolean ignoreNulls) {
+        return this.delegateTableSchema.updateItemToMap(item, ignoreNulls);
+    }
+
     @Override
     public Map<String, AttributeValue> itemToMap(T item, Collection<String> attributes) {
         return this.delegateTableSchema.itemToMap(item, attributes);

diff --git a/...est/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/UpdateBehaviorTest.java b/...est/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/UpdateBehaviorTest.java
@@ -1,10 +1,12 @@
 package software.amazon.awssdk.enhanced.dynamodb.functionaltests;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import java.time.Instant;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import javax.xml.bind.ValidationException;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -15,6 +17,7 @@
 import software.amazon.awssdk.enhanced.dynamodb.functionaltests.models.NestedRecordWithUpdateBehavior;
 import software.amazon.awssdk.enhanced.dynamodb.functionaltests.models.RecordWithUpdateBehaviors;
 import software.amazon.awssdk.enhanced.dynamodb.internal.client.ExtensionResolver;
+import software.amazon.awssdk.services.dynamodb.model.DynamoDbException;
 
 public class UpdateBehaviorTest extends LocalDynamoDbSyncTestBase {
     private static final Instant INSTANT_1 = Instant.parse("2020-05-03T10:00:00Z");
@@ -145,11 +148,45 @@ public void updateBehaviors_transactWriteItems_secondUpdate() {
         assertThat(persistedRecord.getCreatedAutoUpdateOn()).isEqualTo(firstUpdatedRecord.getCreatedAutoUpdateOn());
     }
 
+    @Test
+    public void updateBehaviors_nested() {
+
+        NestedRecordWithUpdateBehavior nestedRecord = new NestedRecordWithUpdateBehavior();
+        nestedRecord.setId("id456");
+        nestedRecord.setNestedCounter(5L);
+        nestedRecord.setNestedUpdateBehaviorAttribute("TEST_BEHAVIOUR_ATTRIBUTE");
+        nestedRecord.setNestedTimeAttribute(INSTANT_1);
+
+        RecordWithUpdateBehaviors record = new RecordWithUpdateBehaviors();
+        record.setId("id123");
+        record.setNestedRecord(nestedRecord);
+
+        mappedTable.putItem(record);
+
+        NestedRecordWithUpdateBehavior updatedNestedRecord = new NestedRecordWithUpdateBehavior();
+        updatedNestedRecord.setNestedCounter(10L);
+
+        RecordWithUpdateBehaviors update_record = new RecordWithUpdateBehaviors();
+        update_record.setId("id123");
+        update_record.setVersion(1L);
+        update_record.setNestedRecord(updatedNestedRecord);
+
+        mappedTable.updateItem(r -> r.item(update_record).ignoreNulls(true));
+
+        RecordWithUpdateBehaviors persistedRecord = mappedTable.getItem(r -> r.key(k -> k.partitionValue("id123")));
+
+        assertThat(persistedRecord.getVersion()).isEqualTo(2L);
+        assertThat(persistedRecord.getNestedRecord()).isNotNull();
+        assertThat(persistedRecord.getNestedRecord().getNestedCounter()).isEqualTo(10L);
+        assertThat(persistedRecord.getNestedRecord().getNestedUpdateBehaviorAttribute()).isEqualTo("TEST_BEHAVIOUR_ATTRIBUTE");
+        assertThat(persistedRecord.getNestedRecord().getNestedTimeAttribute()).isEqualTo(INSTANT_1);
+    }
+
     /**
      * Currently, nested records are not updated through extensions.
      */
     @Test
-    public void updateBehaviors_nested() {
+    public void updateNonexistentField_nested() {
         NestedRecordWithUpdateBehavior nestedRecord = new NestedRecordWithUpdateBehavior();
         nestedRecord.setId("id456");
 
@@ -158,15 +195,8 @@ public void updateBehaviors_nested() {
         record.setCreatedOn(INSTANT_1);
         record.setLastUpdatedOn(INSTANT_2);
         record.setNestedRecord(nestedRecord);
-        mappedTable.updateItem(record);
-
-        RecordWithUpdateBehaviors persistedRecord = mappedTable.getItem(record);
-
-        assertThat(persistedRecord.getVersion()).isEqualTo(1L);
-        assertThat(persistedRecord.getNestedRecord()).isNotNull();
-        assertThat(persistedRecord.getNestedRecord().getNestedVersionedAttribute()).isNull();
-        assertThat(persistedRecord.getNestedRecord().getNestedCounter()).isNull();
-        assertThat(persistedRecord.getNestedRecord().getNestedUpdateBehaviorAttribute()).isNull();
-        assertThat(persistedRecord.getNestedRecord().getNestedTimeAttribute()).isNull();
+        assertThatThrownBy(() ->mappedTable.updateItem(record))
+            .isInstanceOf(DynamoDbException.class)
+            .hasMessageContaining("The document path provided in the update expression is invalid for update");
     }
 }