Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
{
"type": "feature",
"category": "Amazon DynamoDB Enhanced Client",
"contributor": "",
"description": "DynamoDb enhanced client: support UpdateExpressions in single-request update"
}
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,10 @@

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;
import java.util.concurrent.CompletableFuture;
Expand All @@ -36,6 +34,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.IgnoreNullsMode;
import software.amazon.awssdk.enhanced.dynamodb.model.TransactUpdateItemEnhancedRequest;
import software.amazon.awssdk.enhanced.dynamodb.model.UpdateItemEnhancedRequest;
Expand Down Expand Up @@ -132,7 +131,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, nonKeyAttributes, request);
Expression conditionExpression = generateConditionExpressionIfExist(transformation, request);

Map<String, String> expressionNames = coalesceExpressionNames(updateExpression, conditionExpression);
Expand Down Expand Up @@ -271,27 +270,38 @@ public TransactWriteItem generateTransactWriteItem(TableSchema<T> tableSchema, O
}

/**
* Retrieves the UpdateExpression from extensions if existing, and then creates an UpdateExpression for the request POJO
* 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.
* Merges UpdateExpressions from three sources in priority order: POJO attributes (lowest),
* extensions (medium), request (highest). Higher priority sources override conflicting actions.
*
* <p>Null POJO attributes normally generate REMOVE actions, but are skipped if the same
* attribute is referenced in extension/request expressions to avoid DynamoDB conflicts.
*
* @param tableMetadata metadata about the table structure
* @param transformation write modification from extensions containing UpdateExpression
* @param attributes non-key attributes from the POJO item
* @param request the update request containing optional explicit UpdateExpression
* @return merged Expression containing the final update expression, or null if no updates needed
*/
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,
Map<String, AttributeValue> attributes,
Either<UpdateItemEnhancedRequest<T>, TransactUpdateItemEnhancedRequest<T>> request) {

UpdateExpression requestUpdateExpression =
request.left().map(UpdateItemEnhancedRequest::updateExpression)
.orElseGet(() -> request.right().map(TransactUpdateItemEnhancedRequest::updateExpression).orElse(null));

UpdateExpressionResolver updateExpressionResolver =
UpdateExpressionResolver.builder()
.tableMetadata(tableMetadata)
.nonKeyAttributes(attributes)
.requestExpression(requestUpdateExpression)
.extensionExpression(transformation != null ? transformation.updateExpression() : null)
.build();

UpdateExpression mergedUpdateExpression = updateExpressionResolver.resolve();
return UpdateExpressionConverter.toExpression(mergedUpdateExpression);
}

/**
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,156 @@
/*
* 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 java.util.Objects.requireNonNull;
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.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
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;

/**
* Resolves and merges UpdateExpressions from multiple sources (item attributes, extensions, requests) with priority-based
* conflict resolution and smart filtering to prevent attribute conflicts.
*/
@SdkInternalApi
public final class UpdateExpressionResolver {

private final TableMetadata tableMetadata;
private final Map<String, AttributeValue> nonKeyAttributes;
private final UpdateExpression extensionExpression;
private final UpdateExpression requestExpression;

private UpdateExpressionResolver(Builder builder) {
this.tableMetadata = builder.tableMetadata;
this.nonKeyAttributes = builder.nonKeyAttributes;
this.extensionExpression = builder.extensionExpression;
this.requestExpression = builder.requestExpression;
}

public static Builder builder() {
return new Builder();
}

/**
* Merges UpdateExpressions from three sources with priority: item attributes (lowest), extension expressions (medium),
* request expressions (highest).
*
* <p><b>Steps:</b> Identify attributes used by extensions/requests to prevent REMOVE conflicts →
* create item SET/REMOVE actions → merge extensions (override item) → merge request (override all).
*
* <p><b>Backward compatibility:</b> Without request expressions, behavior is identical to previous versions.
* <p><b>Exceptions:</b> DynamoDbException may be thrown when the same attribute is updated by multiple sources.
*
* @return merged UpdateExpression, or empty if no updates needed
*/
public UpdateExpression resolve() {
UpdateExpression itemExpression = null;

if (!nonKeyAttributes.isEmpty()) {
Set<String> attributesExcludedFromRemoval = attributesPresentInOtherExpressions(
Arrays.asList(extensionExpression, requestExpression));

itemExpression = UpdateExpression.mergeExpressions(
generateItemSetExpression(nonKeyAttributes, tableMetadata),
generateItemRemoveExpression(nonKeyAttributes, attributesExcludedFromRemoval));
}

return Stream.of(itemExpression, extensionExpression, requestExpression)
.filter(Objects::nonNull)
.reduce(UpdateExpression::mergeExpressions)
.orElse(null);
}

private static Set<String> attributesPresentInOtherExpressions(Collection<UpdateExpression> updateExpressions) {
return updateExpressions.stream()
.filter(Objects::nonNull)
.map(UpdateExpressionConverter::findAttributeNames)
.flatMap(List::stream)
.collect(Collectors.toSet());
}

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,
Collection<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 Map<String, AttributeValue> nonKeyAttributes;
private UpdateExpression extensionExpression;
private UpdateExpression requestExpression;

public Builder tableMetadata(TableMetadata tableMetadata) {
this.tableMetadata = requireNonNull(
tableMetadata, "A TableMetadata is required when generating an Update Expression");
return this;
}

public Builder nonKeyAttributes(Map<String, AttributeValue> nonKeyAttributes) {
if (nonKeyAttributes == null) {
this.nonKeyAttributes = Collections.emptyMap();
} else {
this.nonKeyAttributes = Collections.unmodifiableMap(new HashMap<>(nonKeyAttributes));
}
return this;
}

public Builder extensionExpression(UpdateExpression extensionExpression) {
this.extensionExpression = extensionExpression;
return this;
}

public Builder requestExpression(UpdateExpression requestExpression) {
this.requestExpression = requestExpression;
return this;
}

public UpdateExpressionResolver build() {
return new UpdateExpressionResolver(this);
}

}
}
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,9 @@

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.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;
Expand All @@ -35,7 +33,6 @@
import software.amazon.awssdk.enhanced.dynamodb.mapper.UpdateBehavior;
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.AttributeValue;

@SdkInternalApi
Expand All @@ -53,32 +50,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) {
static List<SetAction> setActionsFor(Map<String, AttributeValue> attributesToSet, TableMetadata tableMetadata) {
return attributesToSet.entrySet()
.stream()
.map(entry -> setValue(entry.getKey(),
Expand All @@ -90,7 +65,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) {
static List<RemoveAction> removeActionsFor(Map<String, AttributeValue> attributesToSet) {
return attributesToSet.entrySet()
.stream()
.map(entry -> remove(entry.getKey()))
Expand Down
Loading