1 /* 2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"). 5 * You may not use this file except in compliance with the License. 6 * A copy of the License is located at 7 * 8 * http://aws.amazon.com/apache2.0 9 * 10 * or in the "license" file accompanying this file. This file is distributed 11 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either 12 * express or implied. See the License for the specific language governing 13 * permissions and limitations under the License. 14 */ 15 16 package software.amazon.awssdk.enhanced.dynamodb.internal.update; 17 18 import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.isNullAttributeValue; 19 import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.keyRef; 20 import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.valueRef; 21 import static software.amazon.awssdk.utils.CollectionUtils.filterMap; 22 23 import java.util.Arrays; 24 import java.util.Collections; 25 import java.util.List; 26 import java.util.Map; 27 import java.util.function.Function; 28 import java.util.stream.Collectors; 29 import software.amazon.awssdk.annotations.SdkInternalApi; 30 import software.amazon.awssdk.enhanced.dynamodb.TableMetadata; 31 import software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils; 32 import software.amazon.awssdk.enhanced.dynamodb.internal.mapper.UpdateBehaviorTag; 33 import software.amazon.awssdk.enhanced.dynamodb.mapper.UpdateBehavior; 34 import software.amazon.awssdk.enhanced.dynamodb.update.RemoveAction; 35 import software.amazon.awssdk.enhanced.dynamodb.update.SetAction; 36 import software.amazon.awssdk.enhanced.dynamodb.update.UpdateExpression; 37 import software.amazon.awssdk.services.dynamodb.model.AttributeValue; 38 39 @SdkInternalApi 40 public final class UpdateExpressionUtils { 41 UpdateExpressionUtils()42 private UpdateExpressionUtils() { 43 } 44 45 /** 46 * A function to specify an initial value if the attribute represented by 'key' does not exist. 47 */ ifNotExists(String key, String initValue)48 public static String ifNotExists(String key, String initValue) { 49 return "if_not_exists(" + keyRef(key) + ", " + valueRef(initValue) + ")"; 50 } 51 52 /** 53 * Generates an UpdateExpression representing a POJO, with only SET and REMOVE actions. 54 */ operationExpression(Map<String, AttributeValue> itemMap, TableMetadata tableMetadata, List<String> nonRemoveAttributes)55 public static UpdateExpression operationExpression(Map<String, AttributeValue> itemMap, 56 TableMetadata tableMetadata, 57 List<String> nonRemoveAttributes) { 58 59 Map<String, AttributeValue> setAttributes = filterMap(itemMap, e -> !isNullAttributeValue(e.getValue())); 60 UpdateExpression setAttributeExpression = UpdateExpression.builder() 61 .actions(setActionsFor(setAttributes, tableMetadata)) 62 .build(); 63 64 Map<String, AttributeValue> removeAttributes = 65 filterMap(itemMap, e -> isNullAttributeValue(e.getValue()) && !nonRemoveAttributes.contains(e.getKey())); 66 67 UpdateExpression removeAttributeExpression = UpdateExpression.builder() 68 .actions(removeActionsFor(removeAttributes)) 69 .build(); 70 71 return UpdateExpression.mergeExpressions(setAttributeExpression, removeAttributeExpression); 72 } 73 74 /** 75 * Creates a list of SET actions for all attributes supplied in the map. 76 */ setActionsFor(Map<String, AttributeValue> attributesToSet, TableMetadata tableMetadata)77 private static List<SetAction> setActionsFor(Map<String, AttributeValue> attributesToSet, TableMetadata tableMetadata) { 78 return attributesToSet.entrySet() 79 .stream() 80 .map(entry -> setValue(entry.getKey(), 81 entry.getValue(), 82 UpdateBehaviorTag.resolveForAttribute(entry.getKey(), tableMetadata))) 83 .collect(Collectors.toList()); 84 } 85 86 /** 87 * Creates a list of REMOVE actions for all attributes supplied in the map. 88 */ removeActionsFor(Map<String, AttributeValue> attributesToSet)89 private static List<RemoveAction> removeActionsFor(Map<String, AttributeValue> attributesToSet) { 90 return attributesToSet.entrySet() 91 .stream() 92 .map(entry -> remove(entry.getKey())) 93 .collect(Collectors.toList()); 94 } 95 96 /** 97 * Creates a REMOVE action for an attribute, using a token as a placeholder for the attribute name. 98 */ remove(String attributeName)99 private static RemoveAction remove(String attributeName) { 100 return RemoveAction.builder() 101 .path(keyRef(attributeName)) 102 .expressionNames(Collections.singletonMap(keyRef(attributeName), attributeName)) 103 .build(); 104 } 105 106 /** 107 * Creates a SET action for an attribute, using a token as a placeholder for the attribute name. 108 * 109 * @see UpdateBehavior for information about the values available. 110 */ setValue(String attributeName, AttributeValue value, UpdateBehavior updateBehavior)111 private static SetAction setValue(String attributeName, AttributeValue value, UpdateBehavior updateBehavior) { 112 return SetAction.builder() 113 .path(keyRef(attributeName)) 114 .value(behaviorBasedValue(updateBehavior).apply(attributeName)) 115 .expressionNames(expressionNamesFor(attributeName)) 116 .expressionValues(Collections.singletonMap(valueRef(attributeName), value)) 117 .build(); 118 } 119 120 /** 121 * When we know we want to update the attribute no matter if it exists or not, we simply need to replace the value with 122 * a value token in the expression. If we only want to set the value if the attribute doesn't exist, we use 123 * the DDB function ifNotExists. 124 */ behaviorBasedValue(UpdateBehavior updateBehavior)125 private static Function<String, String> behaviorBasedValue(UpdateBehavior updateBehavior) { 126 switch (updateBehavior) { 127 case WRITE_ALWAYS: 128 return v -> valueRef(v); 129 case WRITE_IF_NOT_EXISTS: 130 return k -> ifNotExists(k, k); 131 default: 132 throw new IllegalArgumentException("Unsupported update behavior '" + updateBehavior + "'"); 133 } 134 } 135 136 /** 137 * Simple utility method that can create an ExpressionNames map based on a list of attribute names. 138 */ expressionNamesFor(String... attributeNames)139 private static Map<String, String> expressionNamesFor(String... attributeNames) { 140 return Arrays.stream(attributeNames) 141 .collect(Collectors.toMap(EnhancedClientUtils::keyRef, Function.identity())); 142 } 143 144 }