UpdateExpression at request level

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

@@ -16,11 +16,9 @@
 package software.amazon.awssdk.enhanced.dynamodb.internal.operations;
 
 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.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
@@ -34,6 +32,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.internal.update.UpdateExpressionResolver;
 import software.amazon.awssdk.enhanced.dynamodb.model.TransactUpdateItemEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedRequest;
 import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedResponse;
@@ -113,7 +112,7 @@ public UpdateItemRequest generateRequest(TableSchema<T> tableSchema,
         Map<String, AttributeValue> keyAttributes = filterMap(itemMap, entry -> primaryKeys.contains(entry.getKey()));
         Map<String, AttributeValue> nonKeyAttributes = filterMap(itemMap, entry -> !primaryKeys.contains(entry.getKey()));
 
-        Expression updateExpression = generateUpdateExpressionIfExist(tableMetadata, transformation, nonKeyAttributes);
+        Expression updateExpression = generateUpdateExpressionIfExist(tableMetadata, transformation, request, nonKeyAttributes);
         Expression conditionExpression = generateConditionExpressionIfExist(transformation, request);
 
         Map<String, String> expressionNames = coalesceExpressionNames(updateExpression, conditionExpression);
@@ -205,23 +204,26 @@ public TransactWriteItem generateTransactWriteItem(TableSchema<T> tableSchema, O
      * if there are attributes to be updated (most likely). If both exist, they are merged and the code generates a final
      * Expression that represent the result.
      */
-    private Expression generateUpdateExpressionIfExist(TableMetadata tableMetadata,
-                                                       WriteModification transformation,
-                                                       Map<String, AttributeValue> attributes) {
-        UpdateExpression updateExpression = null;
-        if (transformation != null && transformation.updateExpression() != null) {
-            updateExpression = transformation.updateExpression();
-        }
-        if (!attributes.isEmpty()) {
-            List<String> nonRemoveAttributes = UpdateExpressionConverter.findAttributeNames(updateExpression);
-            UpdateExpression operationUpdateExpression = operationExpression(attributes, tableMetadata, nonRemoveAttributes);
-            if (updateExpression == null) {
-                updateExpression = operationUpdateExpression;
-            } else {
-                updateExpression = UpdateExpression.mergeExpressions(updateExpression, operationUpdateExpression);
-            }
-        }
-        return UpdateExpressionConverter.toExpression(updateExpression);
+    private Expression generateUpdateExpressionIfExist(
+        TableMetadata tableMetadata,
+        WriteModification transformation,
+        Either<UpdateItemEnhancedRequest<T>, TransactUpdateItemEnhancedRequest<T>> request,
+        Map<String, AttributeValue> nonKeyAttributes) {
+
+        UpdateExpression requestUpdateExpression = request.map(r -> Optional.ofNullable(r.updateExpression()),
+                                                               r -> Optional.ofNullable(r.updateExpression()))
+                                                          .orElse(null);
+
+        UpdateExpressionResolver updateExpressionResolver =
+            UpdateExpressionResolver.builder()
+                                    .tableMetadata(tableMetadata)
+                                    .itemNonKeyAttributes(nonKeyAttributes)
+                                    .requestExpression(requestUpdateExpression)
+                                    .transformationExpression(transformation != null ? transformation.updateExpression() : null)
+                                    .build();
+
+        UpdateExpression mergedUpdateExpression = updateExpressionResolver.resolve();
+        return UpdateExpressionConverter.toExpression(mergedUpdateExpression);
     }
 
     /**

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

@@ -76,7 +76,7 @@ private UpdateExpressionConverter() {
      * @return an Expression representing the concatenation of all actions in this UpdateExpression
      */
     public static Expression toExpression(UpdateExpression expression) {
-        if (expression == null) {
+        if (expression == null || expression.isEmpty()) {
             return null;
         }
         Map<String, AttributeValue> expressionValues = mergeExpressionValues(expression);
@@ -91,8 +91,9 @@ public static Expression toExpression(UpdateExpression expression) {
     }
 
     /**
-     * Attempts to find the list of attribute names that will be updated for the supplied {@link UpdateExpression} by looking at
-     * the combined collection of paths and ExpressionName values. Because attribute names can be composed from nested
+     * Attempts to find the list of attributes associated with update actions for the supplied {@link UpdateExpression} by
+     * looking at
+     * the combined collection of paths and ExpressionName values. Because attribute names can be composed of nested
      * attribute references and list references, the leftmost part will be returned if composition is detected.
      * <p>
      * Examples: The expression contains a {@link DeleteAction} with a path value of 'MyAttribute[1]'; the list returned

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

@@ -0,0 +1,142 @@
+/*
+ * 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.internal.update;
+
+import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.isNullAttributeValue;
+import static software.amazon.awssdk.enhanced.dynamodb.internal.update.UpdateExpressionUtils.removeActionsFor;
+import static software.amazon.awssdk.enhanced.dynamodb.internal.update.UpdateExpressionUtils.setActionsFor;
+import static software.amazon.awssdk.utils.CollectionUtils.filterMap;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
+import software.amazon.awssdk.enhanced.dynamodb.update.UpdateExpression;
+import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+
+/**
+ *
+ */
+@SdkInternalApi
+public final class UpdateExpressionResolver {
+
+    private final UpdateExpression extensionExpression;
+    private final UpdateExpression requestExpression;
+    private final Map<String, AttributeValue> itemNonKeyAttributes;
+    private final TableMetadata tableMetadata;
+
+    private UpdateExpressionResolver(Builder builder) {
+        this.extensionExpression = builder.transformationExpression;
+        this.requestExpression = builder.requestExpression;
+        this.itemNonKeyAttributes = builder.nonKeyAttributes;
+        this.tableMetadata = builder.tableMetadata;
+    }
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /**
+     * Resolves all available and potential update expressions by priority and returns a merged update expression. It may
+     * return null, if the item attribute map is empty / does not contain non-null attributes and no other update expressions
+     * are present.
+     * <p>
+     * Conditions that will result in error:
+     * <ul>
+     *     <li>Two expressions contain actions referencing the same attribute</li>
+     * </ul>
+     * <p>
+     * <b>Note: </b> The presence of attributes in update expressions submitted through the request or generated from extensions
+     * take precedence over removing attributes based on item configuration.
+     * For example, when IGNORE_NULLS is set to true (default), the client generates REMOVE actions for all
+     * attributes in the schema that are not explicitly set in the request item submitted to the operation. If such
+     * attributes are referenced in update expressions on the request or from extensions, the remove actions are filtered
+     * out.
+     */
+    public UpdateExpression resolve() {
+        UpdateExpression itemSetExpression = generateItemSetExpression(itemNonKeyAttributes, tableMetadata);
+
+        List<String> nonRemoveAttributes = attributesPresentInExpressions(Arrays.asList(extensionExpression, requestExpression));
+        UpdateExpression itemRemoveExpression = generateItemRemoveExpression(itemNonKeyAttributes, nonRemoveAttributes);
+
+        UpdateExpression itemExpression = UpdateExpression.mergeExpressions(itemSetExpression, itemRemoveExpression);
+        UpdateExpression extensionItemExpression = UpdateExpression.mergeExpressions(extensionExpression, itemExpression);
+        return UpdateExpression.mergeExpressions(requestExpression, extensionItemExpression);
+    }
+
+    private static List<String> attributesPresentInExpressions(List<UpdateExpression> updateExpressions) {
+        return updateExpressions.stream()
+                                .map(UpdateExpressionConverter::findAttributeNames)
+                                .flatMap(List::stream)
+                                .collect(Collectors.toList());
+    }
+
+    public static UpdateExpression generateItemSetExpression(Map<String, AttributeValue> itemMap,
+                                                             TableMetadata tableMetadata) {
+
+        Map<String, AttributeValue> setAttributes = filterMap(itemMap, e -> !isNullAttributeValue(e.getValue()));
+        return UpdateExpression.builder()
+                               .actions(setActionsFor(setAttributes, tableMetadata))
+                               .build();
+    }
+
+    public static UpdateExpression generateItemRemoveExpression(Map<String, AttributeValue> itemMap,
+                                                                List<String> nonRemoveAttributes) {
+        Map<String, AttributeValue> removeAttributes =
+            filterMap(itemMap, e -> isNullAttributeValue(e.getValue()) && !nonRemoveAttributes.contains(e.getKey()));
+
+        return UpdateExpression.builder()
+                               .actions(removeActionsFor(removeAttributes))
+                               .build();
+    }
+
+    public static final class Builder {
+
+        private TableMetadata tableMetadata;
+        private UpdateExpression transformationExpression;
+        private UpdateExpression requestExpression;
+        private Map<String, AttributeValue> nonKeyAttributes;
+
+        public Builder tableMetadata(TableMetadata tableMetadata) {
+            this.tableMetadata = tableMetadata;
+            return this;
+        }
+
+        public Builder transformationExpression(UpdateExpression transformationExpression) {
+            this.transformationExpression = transformationExpression;
+            return this;
+        }
+
+        public Builder itemNonKeyAttributes(Map<String, AttributeValue> nonKeyAttributes) {
+            this.nonKeyAttributes = nonKeyAttributes;
+            return this;
+        }
+
+        public Builder requestExpression(UpdateExpression requestExpression) {
+            this.requestExpression = requestExpression;
+            return this;
+        }
+
+        public UpdateExpressionResolver build() {
+            return new UpdateExpressionResolver(this);
+        }
+
+    }
+
+
+}

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

@@ -49,32 +49,10 @@ public static String ifNotExists(String key, String initValue) {
         return "if_not_exists(" + keyRef(key) + ", " + valueRef(initValue) + ")";
     }
 
-    /**
-     * Generates an UpdateExpression representing a POJO, with only SET and REMOVE actions.
-     */
-    public static UpdateExpression operationExpression(Map<String, AttributeValue> itemMap,
-                                                       TableMetadata tableMetadata,
-                                                       List<String> nonRemoveAttributes) {
-
-        Map<String, AttributeValue> setAttributes = filterMap(itemMap, e -> !isNullAttributeValue(e.getValue()));
-        UpdateExpression setAttributeExpression = UpdateExpression.builder()
-                                                                  .actions(setActionsFor(setAttributes, tableMetadata))
-                                                                  .build();
-
-        Map<String, AttributeValue> removeAttributes =
-            filterMap(itemMap, e -> isNullAttributeValue(e.getValue()) && !nonRemoveAttributes.contains(e.getKey()));
-
-        UpdateExpression removeAttributeExpression = UpdateExpression.builder()
-                                                                     .actions(removeActionsFor(removeAttributes))
-                                                                     .build();
-
-        return UpdateExpression.mergeExpressions(setAttributeExpression, removeAttributeExpression);
-    }
-
     /**
      * Creates a list of SET actions for all attributes supplied in the map.
      */
-    private static List<SetAction> setActionsFor(Map<String, AttributeValue> attributesToSet, TableMetadata tableMetadata) {
+    public static List<SetAction> setActionsFor(Map<String, AttributeValue> attributesToSet, TableMetadata tableMetadata) {
         return attributesToSet.entrySet()
                               .stream()
                               .map(entry -> setValue(entry.getKey(),
@@ -86,7 +64,7 @@ private static List<SetAction> setActionsFor(Map<String, AttributeValue> attribu
     /**
      * Creates a list of REMOVE actions for all attributes supplied in the map.
      */
-    private static List<RemoveAction> removeActionsFor(Map<String, AttributeValue> attributesToSet) {
+    public static List<RemoveAction> removeActionsFor(Map<String, AttributeValue> attributesToSet) {
         return attributesToSet.entrySet()
                               .stream()
                               .map(entry -> remove(entry.getKey()))

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

@@ -22,6 +22,11 @@
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedAsyncClient;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient;
 import software.amazon.awssdk.enhanced.dynamodb.Expression;
+import software.amazon.awssdk.enhanced.dynamodb.update.AddAction;
+import software.amazon.awssdk.enhanced.dynamodb.update.DeleteAction;
+import software.amazon.awssdk.enhanced.dynamodb.update.RemoveAction;
+import software.amazon.awssdk.enhanced.dynamodb.update.SetAction;
+import software.amazon.awssdk.enhanced.dynamodb.update.UpdateExpression;
 import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
 
 /**
@@ -41,12 +46,14 @@ public class TransactUpdateItemEnhancedRequest<T> {
     private final T item;
     private final Boolean ignoreNulls;
     private final Expression conditionExpression;
+    private final UpdateExpression updateExpression;
     private final String returnValuesOnConditionCheckFailure;
 
     private TransactUpdateItemEnhancedRequest(Builder<T> builder) {
         this.item = builder.item;
         this.ignoreNulls = builder.ignoreNulls;
         this.conditionExpression = builder.conditionExpression;
+        this.updateExpression = builder.updateExpression;
         this.returnValuesOnConditionCheckFailure = builder.returnValuesOnConditionCheckFailure;
     }
 
@@ -92,6 +99,13 @@ public Expression conditionExpression() {
         return conditionExpression;
     }
 
+    /**
+     * Returns the update expression {@link UpdateExpression} set on this request object, or null if it doesn't exist.
+     */
+    public UpdateExpression updateExpression() {
+        return updateExpression;
+    }
+
     /**
      * Returns what values to return if the condition check fails.
      * <p>
@@ -140,6 +154,9 @@ public boolean equals(Object o) {
         if (!Objects.equals(conditionExpression, that.conditionExpression)) {
             return false;
         }
+        if (!Objects.equals(updateExpression, that.updateExpression)) {
+            return false;
+        }
         return Objects.equals(returnValuesOnConditionCheckFailure, that.returnValuesOnConditionCheckFailure);
     }
 
@@ -148,6 +165,7 @@ public int hashCode() {
         int result = Objects.hashCode(item);
         result = 31 * result + Objects.hashCode(ignoreNulls);
         result = 31 * result + Objects.hashCode(conditionExpression);
+        result = 31 * result + Objects.hashCode(updateExpression);
         result = 31 * result + Objects.hashCode(returnValuesOnConditionCheckFailure);
         return result;
     }
@@ -162,6 +180,8 @@ public static final class Builder<T> {
         private T item;
         private Boolean ignoreNulls;
         private Expression conditionExpression;
+
+        private UpdateExpression updateExpression;
         private String returnValuesOnConditionCheckFailure;
 
         private Builder() {
@@ -208,6 +228,35 @@ public Builder<T> item(T item) {
             return this;
         }
 
+        /**
+         * Define an {@link UpdateExpression} to control updating specific parts of the item in DynamoDb. The update expression
+         * corresponds to the DynamoDb update expression format. It can be used to set, modify and delete attributes for
+         * use cases that simply supplying the item does not cover; in particular, manipulating composed attributes such as
+         * sets or lists:
+         * <ul>
+         * <li>Add/remove elements to/from list attributes</li>
+         * <li>Add/remove elements to/from set attributes</li>
+         * <li>Unset or nullify attributes without modifying the whole attribute</li>
+         * </ul>
+         * <p>
+         * This method will throw an exception if the expression references an attribute that is already present on the
+         * item, or is modified through an extension.
+         * <p>
+         * <b>Note: </b>This is a powerful mechanism that bypasses many of the abstractions and
+         * safety checks in the enhanced client, and should be used with caution. Only use it when submitting only
+         * a configured item bean/object is insufficient.
+         * <p>
+         * See {@link UpdateExpression}, {@link AddAction}, {@link DeleteAction}, {@link SetAction} and
+         * {@link RemoveAction} for syntax and examples.
+         *
+         * @param updateExpression a composed expression of type {@link UpdateExpression}
+         * @return a builder of this type
+         */
+        public Builder<T> updateExpression(UpdateExpression updateExpression) {
+            this.updateExpression = updateExpression;
+            return this;
+        }
+
         /**
          * Use <code>ReturnValuesOnConditionCheckFailure</code> to get the item attributes if the <code>ConditionCheck</code>
          * condition fails. For <code>ReturnValuesOnConditionCheckFailure</code>, the valid values are: NONE and