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;
@@ -34,6 +36,7 @@
 import software.amazon.awssdk.enhanced.dynamodb.extensions.WriteModification;
 import software.amazon.awssdk.enhanced.dynamodb.internal.extensions.DefaultDynamoDbExtensionContext;
 import software.amazon.awssdk.enhanced.dynamodb.internal.update.UpdateExpressionConverter;
+import software.amazon.awssdk.enhanced.dynamodb.model.IgnoreNullsMode;
 import software.amazon.awssdk.enhanced.dynamodb.model.TransactUpdateItemEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedResponse;
@@ -53,7 +56,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) {
@@ -90,7 +94,22 @@ public UpdateItemRequest generateRequest(TableSchema<T> tableSchema,
                                           r -> Optional.ofNullable(r.ignoreNulls()))
                                      .orElse(null);
 
-        Map<String, AttributeValue> itemMap = tableSchema.itemToMap(item, Boolean.TRUE.equals(ignoreNulls));
+        IgnoreNullsMode ignoreNullsMode = request.map(r -> Optional.ofNullable(r.ignoreNullsMode()),
+                                                     r -> Optional.ofNullable(r.ignoreNullsMode()))
+                                                        .orElse(IgnoreNullsMode.DEFAULT);
+
+        if (ignoreNullsMode == IgnoreNullsMode.SCALAR_ONLY
+            || ignoreNullsMode == IgnoreNullsMode.MAPS_ONLY) {
+            ignoreNulls = true;
+        }
+        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 = ignoreNullsMode == IgnoreNullsMode.SCALAR_ONLY ?
+                                              transformItemToMapForUpdateExpression(itemMapImmutable) : itemMapImmutable;
+
         TableMetadata tableMetadata = tableSchema.tableMetadata();
 
         WriteModification transformation =
@@ -141,6 +160,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);
+    }
 }

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/IgnoreNullsMode.java

@@ -0,0 +1,37 @@
+/*
+ * 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 software.amazon.awssdk.annotations.SdkPublicApi;
+
+/**
+ * <p>
+ * The SCALAR_ONLY mode supports updates to scalar attributes to any level (top level, first nested level, second nested level,
+ * etc.) when the user wants to update scalar attributes by providing only the delta of changes to be updated. This mode
+ * does not support updates to maps and is expected to throw a 4xx DynamoDB exception if done so.
+ * <p>
+ * In the MAPS_ONLY mode, creation of new map/bean structures through update statements are supported, i.e. setting
+ * null/non-existent maps to non-null values. If users try to update scalar attributes in this mode, it will overwrite
+ * existing values in the table.
+ * <p>
+ * The DEFAULT mode disables any special handling around null values in the update query expression
+ */
+@SdkPublicApi
+public enum IgnoreNullsMode {
+    SCALAR_ONLY,
+    MAPS_ONLY,
+    DEFAULT
+}

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/TransactUpdateItemEnhancedRequest.java

@@ -40,12 +40,14 @@ public class TransactUpdateItemEnhancedRequest<T> {
 
     private final T item;
     private final Boolean ignoreNulls;
+    private final IgnoreNullsMode ignoreNullsMode;
     private final Expression conditionExpression;
     private final String returnValuesOnConditionCheckFailure;
 
     private TransactUpdateItemEnhancedRequest(Builder<T> builder) {
         this.item = builder.item;
         this.ignoreNulls = builder.ignoreNulls;
+        this.ignoreNullsMode = builder.ignoreNullsMode;
         this.conditionExpression = builder.conditionExpression;
         this.returnValuesOnConditionCheckFailure = builder.returnValuesOnConditionCheckFailure;
     }
@@ -67,6 +69,7 @@ public static <T> Builder<T> builder(Class<? extends T> itemClass) {
     public Builder<T> toBuilder() {
         return new Builder<T>().item(item)
                                .ignoreNulls(ignoreNulls)
+                               .ignoreNullsMode(ignoreNullsMode)
                                .conditionExpression(conditionExpression)
                                .returnValuesOnConditionCheckFailure(returnValuesOnConditionCheckFailure);
     }
@@ -80,11 +83,20 @@ public T item() {
 
     /**
      * Returns if the update operation should ignore attributes with null values, or false if it has not been set.
+     * This is deprecated in favour of ignoreNullsMode()
      */
+    @Deprecated
     public Boolean ignoreNulls() {
         return ignoreNulls;
     }
 
+    /**
+     * Returns the mode of update to be performed
+     */
+    public IgnoreNullsMode ignoreNullsMode() {
+        return ignoreNullsMode;
+    }
+
     /**
      * Returns the condition {@link Expression} set on this request object, or null if it doesn't exist.
      */
@@ -161,6 +173,7 @@ public int hashCode() {
     public static final class Builder<T> {
         private T item;
         private Boolean ignoreNulls;
+        private IgnoreNullsMode ignoreNullsMode;
         private Expression conditionExpression;
         private String returnValuesOnConditionCheckFailure;
 
@@ -178,11 +191,17 @@ private Builder() {
          * @param ignoreNulls the boolean value
          * @return a builder of this type
          */
+        @Deprecated
         public Builder<T> ignoreNulls(Boolean ignoreNulls) {
             this.ignoreNulls = ignoreNulls;
             return this;
         }
 
+        public Builder<T> ignoreNullsMode(IgnoreNullsMode ignoreNullsMode) {
+            this.ignoreNullsMode = ignoreNullsMode;
+            return this;
+        }
+
         /**
          * Defines a logical expression on an item's attribute values which, if evaluating to true,
          * will allow the update operation to succeed. If evaluating to false, the operation will not succeed.