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.operations;
17 
18 import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.readAndTransformSingleItem;
19 import static software.amazon.awssdk.enhanced.dynamodb.internal.update.UpdateExpressionUtils.operationExpression;
20 import static software.amazon.awssdk.utils.CollectionUtils.filterMap;
21 
22 import java.util.Collection;
23 import java.util.List;
24 import java.util.Map;
25 import java.util.Optional;
26 import java.util.concurrent.CompletableFuture;
27 import java.util.function.Function;
28 import software.amazon.awssdk.annotations.SdkInternalApi;
29 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
30 import software.amazon.awssdk.enhanced.dynamodb.Expression;
31 import software.amazon.awssdk.enhanced.dynamodb.OperationContext;
32 import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
33 import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
34 import software.amazon.awssdk.enhanced.dynamodb.extensions.WriteModification;
35 import software.amazon.awssdk.enhanced.dynamodb.internal.extensions.DefaultDynamoDbExtensionContext;
36 import software.amazon.awssdk.enhanced.dynamodb.internal.update.UpdateExpressionConverter;
37 import software.amazon.awssdk.enhanced.dynamodb.model.TransactUpdateItemEnhancedRequest;
38 import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedRequest;
39 import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedResponse;
40 import software.amazon.awssdk.enhanced.dynamodb.update.UpdateExpression;
41 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
42 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
43 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
44 import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
45 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
46 import software.amazon.awssdk.services.dynamodb.model.Update;
47 import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
48 import software.amazon.awssdk.services.dynamodb.model.UpdateItemResponse;
49 import software.amazon.awssdk.utils.CollectionUtils;
50 import software.amazon.awssdk.utils.Either;
51 
52 @SdkInternalApi
53 public class UpdateItemOperation<T>
54     implements TableOperation<T, UpdateItemRequest, UpdateItemResponse, UpdateItemEnhancedResponse<T>>,
55                TransactableWriteOperation<T> {
56 
57     private final Either<UpdateItemEnhancedRequest<T>, TransactUpdateItemEnhancedRequest<T>> request;
58 
UpdateItemOperation(UpdateItemEnhancedRequest<T> request)59     private UpdateItemOperation(UpdateItemEnhancedRequest<T> request) {
60         this.request = Either.left(request);
61     }
62 
UpdateItemOperation(TransactUpdateItemEnhancedRequest<T> request)63     private UpdateItemOperation(TransactUpdateItemEnhancedRequest<T> request) {
64         this.request = Either.right(request);
65     }
66 
create(UpdateItemEnhancedRequest<T> request)67     public static <T> UpdateItemOperation<T> create(UpdateItemEnhancedRequest<T> request) {
68         return new UpdateItemOperation<>(request);
69     }
70 
create(TransactUpdateItemEnhancedRequest<T> request)71     public static <T> UpdateItemOperation<T> create(TransactUpdateItemEnhancedRequest<T> request) {
72         return new UpdateItemOperation<>(request);
73     }
74 
75     @Override
operationName()76     public OperationName operationName() {
77         return OperationName.UPDATE_ITEM;
78     }
79 
80     @Override
generateRequest(TableSchema<T> tableSchema, OperationContext operationContext, DynamoDbEnhancedClientExtension extension)81     public UpdateItemRequest generateRequest(TableSchema<T> tableSchema,
82                                              OperationContext operationContext,
83                                              DynamoDbEnhancedClientExtension extension) {
84         if (!TableMetadata.primaryIndexName().equals(operationContext.indexName())) {
85             throw new IllegalArgumentException("UpdateItem cannot be executed against a secondary index.");
86         }
87 
88         T item = request.map(UpdateItemEnhancedRequest::item, TransactUpdateItemEnhancedRequest::item);
89         Boolean ignoreNulls = request.map(r -> Optional.ofNullable(r.ignoreNulls()),
90                                           r -> Optional.ofNullable(r.ignoreNulls()))
91                                      .orElse(null);
92 
93         Map<String, AttributeValue> itemMap = tableSchema.itemToMap(item, Boolean.TRUE.equals(ignoreNulls));
94         TableMetadata tableMetadata = tableSchema.tableMetadata();
95 
96         WriteModification transformation =
97             extension != null
98             ? extension.beforeWrite(DefaultDynamoDbExtensionContext.builder()
99                                                                    .items(itemMap)
100                                                                    .operationContext(operationContext)
101                                                                    .tableMetadata(tableMetadata)
102                                                                    .tableSchema(tableSchema)
103                                                                    .operationName(operationName())
104                                                                    .build())
105             : null;
106 
107         if (transformation != null && transformation.transformedItem() != null) {
108             itemMap = transformation.transformedItem();
109         }
110 
111         Collection<String> primaryKeys = tableSchema.tableMetadata().primaryKeys();
112 
113         Map<String, AttributeValue> keyAttributes = filterMap(itemMap, entry -> primaryKeys.contains(entry.getKey()));
114         Map<String, AttributeValue> nonKeyAttributes = filterMap(itemMap, entry -> !primaryKeys.contains(entry.getKey()));
115 
116         Expression updateExpression = generateUpdateExpressionIfExist(tableMetadata, transformation, nonKeyAttributes);
117         Expression conditionExpression = generateConditionExpressionIfExist(transformation, request);
118 
119         Map<String, String> expressionNames = coalesceExpressionNames(updateExpression, conditionExpression);
120         Map<String, AttributeValue> expressionValues = coalesceExpressionValues(updateExpression, conditionExpression);
121 
122         UpdateItemRequest.Builder requestBuilder = UpdateItemRequest.builder()
123             .tableName(operationContext.tableName())
124             .key(keyAttributes)
125             .returnValues(ReturnValue.ALL_NEW);
126 
127         if (request.left().isPresent()) {
128             addPlainUpdateItemParameters(requestBuilder, request.left().get());
129         }
130         if (updateExpression != null) {
131             requestBuilder.updateExpression(updateExpression.expression());
132         }
133         if (conditionExpression != null) {
134             requestBuilder.conditionExpression(conditionExpression.expression());
135         }
136         if (CollectionUtils.isNotEmpty(expressionNames)) {
137             requestBuilder = requestBuilder.expressionAttributeNames(expressionNames);
138         }
139         if (CollectionUtils.isNotEmpty(expressionValues)) {
140             requestBuilder = requestBuilder.expressionAttributeValues(expressionValues);
141         }
142 
143         return requestBuilder.build();
144     }
145 
146     @Override
transformResponse(UpdateItemResponse response, TableSchema<T> tableSchema, OperationContext operationContext, DynamoDbEnhancedClientExtension extension)147     public UpdateItemEnhancedResponse<T> transformResponse(UpdateItemResponse response,
148                                TableSchema<T> tableSchema,
149                                OperationContext operationContext,
150                                DynamoDbEnhancedClientExtension extension) {
151         try {
152             T attributes = readAndTransformSingleItem(response.attributes(), tableSchema, operationContext, extension);
153 
154             return UpdateItemEnhancedResponse.<T>builder(null)
155                 .attributes(attributes)
156                 .consumedCapacity(response.consumedCapacity())
157                 .itemCollectionMetrics(response.itemCollectionMetrics())
158                 .build();
159         } catch (RuntimeException e) {
160             // With a partial update it's possible to update the record into a state that the mapper can no longer
161             // read or validate. This is more likely to happen with signed and encrypted records that undergo partial
162             // updates (that practice is discouraged for this reason).
163             throw new IllegalStateException("Unable to read the new item returned by UpdateItem after the update "
164                                             + "occurred. Rollbacks are not supported by this operation, therefore the "
165                                             + "record may no longer be readable using this model.", e);
166         }
167     }
168 
169     @Override
serviceCall(DynamoDbClient dynamoDbClient)170     public Function<UpdateItemRequest, UpdateItemResponse> serviceCall(DynamoDbClient dynamoDbClient) {
171         return dynamoDbClient::updateItem;
172     }
173 
174     @Override
asyncServiceCall( DynamoDbAsyncClient dynamoDbAsyncClient)175     public Function<UpdateItemRequest, CompletableFuture<UpdateItemResponse>> asyncServiceCall(
176         DynamoDbAsyncClient dynamoDbAsyncClient) {
177 
178         return dynamoDbAsyncClient::updateItem;
179     }
180 
181     @Override
generateTransactWriteItem(TableSchema<T> tableSchema, OperationContext operationContext, DynamoDbEnhancedClientExtension dynamoDbEnhancedClientExtension)182     public TransactWriteItem generateTransactWriteItem(TableSchema<T> tableSchema, OperationContext operationContext,
183                                                        DynamoDbEnhancedClientExtension dynamoDbEnhancedClientExtension) {
184         UpdateItemRequest updateItemRequest = generateRequest(tableSchema, operationContext, dynamoDbEnhancedClientExtension);
185 
186         Update.Builder builder = Update.builder()
187                                        .key(updateItemRequest.key())
188                                        .tableName(updateItemRequest.tableName())
189                                        .updateExpression(updateItemRequest.updateExpression())
190                                        .conditionExpression(updateItemRequest.conditionExpression())
191                                        .expressionAttributeValues(updateItemRequest.expressionAttributeValues())
192                                        .expressionAttributeNames(updateItemRequest.expressionAttributeNames());
193 
194         request.right()
195                .map(TransactUpdateItemEnhancedRequest::returnValuesOnConditionCheckFailureAsString)
196                .ifPresent(builder::returnValuesOnConditionCheckFailure);
197 
198         return TransactWriteItem.builder()
199                                 .update(builder.build())
200                                 .build();
201     }
202 
203     /**
204      * Retrieves the UpdateExpression from extensions if existing, and then creates an UpdateExpression for the request POJO
205      * if there are attributes to be updated (most likely). If both exist, they are merged and the code generates a final
206      * Expression that represent the result.
207      */
generateUpdateExpressionIfExist(TableMetadata tableMetadata, WriteModification transformation, Map<String, AttributeValue> attributes)208     private Expression generateUpdateExpressionIfExist(TableMetadata tableMetadata,
209                                                        WriteModification transformation,
210                                                        Map<String, AttributeValue> attributes) {
211         UpdateExpression updateExpression = null;
212         if (transformation != null && transformation.updateExpression() != null) {
213             updateExpression = transformation.updateExpression();
214         }
215         if (!attributes.isEmpty()) {
216             List<String> nonRemoveAttributes = UpdateExpressionConverter.findAttributeNames(updateExpression);
217             UpdateExpression operationUpdateExpression = operationExpression(attributes, tableMetadata, nonRemoveAttributes);
218             if (updateExpression == null) {
219                 updateExpression = operationUpdateExpression;
220             } else {
221                 updateExpression = UpdateExpression.mergeExpressions(updateExpression, operationUpdateExpression);
222             }
223         }
224         return UpdateExpressionConverter.toExpression(updateExpression);
225     }
226 
227     /**
228      * Retrieves the ConditionExpression from extensions if existing, and retrieves the ConditionExpression from the request
229      * if existing. If both exist, they are merged.
230      */
generateConditionExpressionIfExist( WriteModification transformation, Either<UpdateItemEnhancedRequest<T>, TransactUpdateItemEnhancedRequest<T>> request)231     private Expression generateConditionExpressionIfExist(
232             WriteModification transformation,
233             Either<UpdateItemEnhancedRequest<T>, TransactUpdateItemEnhancedRequest<T>> request) {
234 
235         Expression conditionExpression = null;
236 
237         if (transformation != null && transformation.additionalConditionalExpression() != null) {
238             conditionExpression = transformation.additionalConditionalExpression();
239         }
240 
241         Expression operationConditionExpression = request.map(r -> Optional.ofNullable(r.conditionExpression()),
242                                                               r -> Optional.ofNullable(r.conditionExpression()))
243                                                          .orElse(null);
244         if (operationConditionExpression != null) {
245             conditionExpression = operationConditionExpression.and(conditionExpression);
246         }
247         return conditionExpression;
248     }
249 
addPlainUpdateItemParameters(UpdateItemRequest.Builder requestBuilder, UpdateItemEnhancedRequest<?> enhancedRequest)250     private UpdateItemRequest.Builder addPlainUpdateItemParameters(UpdateItemRequest.Builder requestBuilder,
251                                                                    UpdateItemEnhancedRequest<?> enhancedRequest) {
252         requestBuilder = requestBuilder.returnConsumedCapacity(enhancedRequest.returnConsumedCapacityAsString());
253         requestBuilder = requestBuilder.returnItemCollectionMetrics(enhancedRequest.returnItemCollectionMetricsAsString());
254         requestBuilder =
255             requestBuilder.returnValuesOnConditionCheckFailure(enhancedRequest.returnValuesOnConditionCheckFailureAsString());
256         return requestBuilder;
257     }
258 
coalesceExpressionNames(Expression firstExpression, Expression secondExpression)259     private static Map<String, String> coalesceExpressionNames(Expression firstExpression, Expression secondExpression) {
260         Map<String, String> expressionNames = null;
261         if (firstExpression != null && !CollectionUtils.isNullOrEmpty(firstExpression.expressionNames())) {
262             expressionNames = firstExpression.expressionNames();
263         }
264         if (secondExpression != null && !CollectionUtils.isNullOrEmpty(secondExpression.expressionNames())) {
265             expressionNames = Expression.joinNames(expressionNames, secondExpression.expressionNames());
266         }
267         return expressionNames;
268     }
269 
coalesceExpressionValues(Expression firstExpression, Expression secondExpression)270     private static Map<String, AttributeValue> coalesceExpressionValues(Expression firstExpression, Expression secondExpression) {
271         Map<String, AttributeValue> expressionValues = null;
272         if (firstExpression != null && !CollectionUtils.isNullOrEmpty(firstExpression.expressionValues())) {
273             expressionValues = firstExpression.expressionValues();
274         }
275         if (secondExpression != null && !CollectionUtils.isNullOrEmpty(secondExpression.expressionValues())) {
276             expressionValues = Expression.joinValues(expressionValues, secondExpression.expressionValues());
277         }
278         return expressionValues;
279     }
280 }
281