Fixed EnhancedClient UpdateItem operation to make it work on nested attributes as well

.changes/next-release/bugfix-AWSSDKforJavav2-2a80679.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2",
+    "contributor": "anirudh9391",
+    "description": "Introduce a new method to transform input to be able to perform update operations on nested DynamoDB object attributes."
+}

services-custom/dynamodb-enhanced/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.internal.operations.UpdateItemOperation.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 NESTED_OBJECT_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) ?
+                             NESTED_OBJECT_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) ?
+                             NESTED_OBJECT_PATTERN.matcher(cleanAttributeName).replaceAll("_")
+                                                                   : cleanAttributeName;
+        return ":AMZN_MAPPED_" + cleanAttributeName;
     }
 
     public static <T> T readAndTransformSingleItem(Map<String, AttributeValue> itemMap,

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/UpdateItemOperation.java

@@ -15,11 +15,13 @@
 
 package software.amazon.awssdk.enhanced.dynamodb.internal.operations;
 
+import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.isNullAttributeValue;
 import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.readAndTransformSingleItem;
 import static software.amazon.awssdk.enhanced.dynamodb.internal.update.UpdateExpressionUtils.operationExpression;
 import static software.amazon.awssdk.utils.CollectionUtils.filterMap;
 
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -53,7 +55,8 @@
 public class UpdateItemOperation<T>
     implements TableOperation<T, UpdateItemRequest, UpdateItemResponse, UpdateItemEnhancedResponse<T>>,
                TransactableWriteOperation<T> {
-
+
+    public static final String NESTED_OBJECT_UPDATE = "_NESTED_ATTR_UPDATE_";
     private final Either<UpdateItemEnhancedRequest<T>, TransactUpdateItemEnhancedRequest<T>> request;
 
     private UpdateItemOperation(UpdateItemEnhancedRequest<T> request) {
@@ -89,8 +92,14 @@ public UpdateItemRequest generateRequest(TableSchema<T> tableSchema,
         Boolean ignoreNulls = request.map(r -> Optional.ofNullable(r.ignoreNulls()),
                                           r -> Optional.ofNullable(r.ignoreNulls()))
                                      .orElse(null);
-
-        Map<String, AttributeValue> itemMap = tableSchema.itemToMap(item, Boolean.TRUE.equals(ignoreNulls));
+
+        Map<String, AttributeValue> itemMapImmutable = tableSchema.itemToMap(item, Boolean.TRUE.equals(ignoreNulls));
+
+        // If ignoreNulls is set to true, check for nested params to be updated
+        // If needed, Transform itemMap for it to be able to handle them.
+        Map<String, AttributeValue> itemMap = Boolean.TRUE.equals(ignoreNulls) ?
+                                             transformItemToMapForUpdateExpression(itemMapImmutable) : itemMapImmutable;
+
         TableMetadata tableMetadata = tableSchema.tableMetadata();
 
         WriteModification transformation =
@@ -141,6 +150,58 @@ public UpdateItemRequest generateRequest(TableSchema<T> tableSchema,
 
         return requestBuilder.build();
     }
+
+    /**
+     * Method checks if a nested object parameter requires an update
+     * If so flattens out nested params separated by "_NESTED_ATTR_UPDATE_"
+     * this is consumed by @link EnhancedClientUtils to form the appropriate UpdateExpression
+     */
+    public Map<String, AttributeValue> transformItemToMapForUpdateExpression(Map<String, AttributeValue> itemToMap) {
+
+        Map<String, AttributeValue> nestedAttributes = new HashMap<>();
+
+        itemToMap.forEach((key, value) -> {
+            if (value.hasM() && isNotEmptyMap(value.m())) {
+                nestedAttributes.put(key, value);
+            }
+        });
+
+        if (!nestedAttributes.isEmpty()) {
+            Map<String, AttributeValue> itemToMapMutable = new HashMap<>(itemToMap);
+            nestedAttributes.forEach((key, value) -> {
+                itemToMapMutable.remove(key);
+                nestedItemToMap(itemToMapMutable, key, value);
+            });
+            return itemToMapMutable;
+        }
+
+        return itemToMap;
+    }
+
+    private Map<String, AttributeValue> nestedItemToMap(Map<String, AttributeValue> itemToMap,
+                                 String key,
+                                 AttributeValue attributeValue) {
+        attributeValue.m().forEach((mapKey, mapValue) -> {
+            String nestedAttributeKey = key + NESTED_OBJECT_UPDATE + mapKey;
+            if (attributeValueNonNullOrShouldWriteNull(mapValue)) {
+                if (mapValue.hasM()) {
+                    nestedItemToMap(itemToMap, nestedAttributeKey, mapValue);
+                } else {
+                    itemToMap.put(nestedAttributeKey, mapValue);
+                }
+            }
+        });
+        return itemToMap;
+    }
+
+    private boolean isNotEmptyMap(Map<String, AttributeValue> map) {
+        return !map.isEmpty() && map.values().stream()
+                                    .anyMatch(this::attributeValueNonNullOrShouldWriteNull);
+    }
+
+    private boolean attributeValueNonNullOrShouldWriteNull(AttributeValue attributeValue) {
+        return !isNullAttributeValue(attributeValue);
+    }
 
     @Override
     public UpdateItemEnhancedResponse<T> transformResponse(UpdateItemResponse response,

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/update/UpdateExpressionUtils.java

@@ -18,13 +18,15 @@
 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.internal.operations.UpdateItemOperation.NESTED_OBJECT_UPDATE;
 import static software.amazon.awssdk.utils.CollectionUtils.filterMap;
 
 import java.util.Arrays;
 import java.util.Collections;
 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 +41,8 @@
 @SdkInternalApi
 public final class UpdateExpressionUtils {
 
+    private static final Pattern PATTERN = Pattern.compile(NESTED_OBJECT_UPDATE);
+
     private UpdateExpressionUtils() {
     }
 
@@ -136,9 +140,12 @@ 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()));
-    }
+    private static Map<String, String> expressionNamesFor(String attributeNames) {
+        if (attributeNames.contains(NESTED_OBJECT_UPDATE)) {
+            return Arrays.stream(PATTERN.split(attributeNames)).distinct()
+                         .collect(Collectors.toMap(EnhancedClientUtils::keyRef, Function.identity()));
+        }
 
+        return Collections.singletonMap(keyRef(attributeNames), attributeNames);
+    }
 }