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 java .util .Objects .requireNonNull ;
19+ import static software .amazon .awssdk .enhanced .dynamodb .internal .EnhancedClientUtils .isNullAttributeValue ;
20+ import static software .amazon .awssdk .enhanced .dynamodb .internal .update .UpdateExpressionUtils .removeActionsFor ;
21+ import static software .amazon .awssdk .enhanced .dynamodb .internal .update .UpdateExpressionUtils .setActionsFor ;
22+ import static software .amazon .awssdk .utils .CollectionUtils .filterMap ;
23+
24+ import java .util .Arrays ;
25+ import java .util .Collection ;
26+ import java .util .Collections ;
27+ import java .util .HashMap ;
28+ import java .util .List ;
29+ import java .util .Map ;
30+ import java .util .Objects ;
31+ import java .util .Set ;
32+ import java .util .stream .Collectors ;
33+ import java .util .stream .Stream ;
34+ import software .amazon .awssdk .annotations .SdkInternalApi ;
35+ import software .amazon .awssdk .enhanced .dynamodb .TableMetadata ;
36+ import software .amazon .awssdk .enhanced .dynamodb .update .UpdateExpression ;
37+ import software .amazon .awssdk .services .dynamodb .model .AttributeValue ;
38+
39+ /**
40+ * Resolves and merges UpdateExpressions from multiple sources (item attributes, extensions, requests) with priority-based
41+ * conflict resolution and smart filtering to prevent attribute conflicts.
42+ */
43+ @ SdkInternalApi
44+ public final class UpdateExpressionResolver {
45+
46+ private final TableMetadata tableMetadata ;
47+ private final Map <String , AttributeValue > nonKeyAttributes ;
48+ private final UpdateExpression extensionExpression ;
49+ private final UpdateExpression requestExpression ;
50+
51+ private UpdateExpressionResolver (Builder builder ) {
52+ this .tableMetadata = builder .tableMetadata ;
53+ this .nonKeyAttributes = builder .nonKeyAttributes ;
54+ this .extensionExpression = builder .extensionExpression ;
55+ this .requestExpression = builder .requestExpression ;
56+ }
57+
58+ public static Builder builder () {
59+ return new Builder ();
60+ }
61+
62+ /**
63+ * Merges UpdateExpressions from three sources with priority: item attributes (lowest), extension expressions (medium),
64+ * request expressions (highest).
65+ *
66+ * <p><b>Steps:</b> Identify attributes used by extensions/requests to prevent REMOVE conflicts →
67+ * create item SET/REMOVE actions → merge extensions (override item) → merge request (override all).
68+ *
69+ * <p><b>Backward compatibility:</b> Without request expressions, behavior is identical to previous versions.
70+ * <p><b>Exceptions:</b> DynamoDbException may be thrown when the same attribute is updated by multiple sources.
71+ *
72+ * @return merged UpdateExpression, or empty if no updates needed
73+ */
74+ public UpdateExpression resolve () {
75+ UpdateExpression itemExpression = null ;
76+
77+ if (!nonKeyAttributes .isEmpty ()) {
78+ Set <String > attributesExcludedFromRemoval = attributesPresentInOtherExpressions (
79+ Arrays .asList (extensionExpression , requestExpression ));
80+
81+ itemExpression = UpdateExpression .mergeExpressions (
82+ generateItemSetExpression (nonKeyAttributes , tableMetadata ),
83+ generateItemRemoveExpression (nonKeyAttributes , attributesExcludedFromRemoval ));
84+ }
85+
86+ return Stream .of (itemExpression , extensionExpression , requestExpression )
87+ .filter (Objects ::nonNull )
88+ .reduce (UpdateExpression ::mergeExpressions )
89+ .orElse (null );
90+ }
91+
92+ private static Set <String > attributesPresentInOtherExpressions (Collection <UpdateExpression > updateExpressions ) {
93+ return updateExpressions .stream ()
94+ .filter (Objects ::nonNull )
95+ .map (UpdateExpressionConverter ::findAttributeNames )
96+ .flatMap (List ::stream )
97+ .collect (Collectors .toSet ());
98+ }
99+
100+ public static UpdateExpression generateItemSetExpression (Map <String , AttributeValue > itemMap ,
101+ TableMetadata tableMetadata ) {
102+
103+ Map <String , AttributeValue > setAttributes = filterMap (itemMap , e -> !isNullAttributeValue (e .getValue ()));
104+ return UpdateExpression .builder ()
105+ .actions (setActionsFor (setAttributes , tableMetadata ))
106+ .build ();
107+ }
108+
109+ public static UpdateExpression generateItemRemoveExpression (Map <String , AttributeValue > itemMap ,
110+ Collection <String > nonRemoveAttributes ) {
111+ Map <String , AttributeValue > removeAttributes =
112+ filterMap (itemMap , e -> isNullAttributeValue (e .getValue ()) && !nonRemoveAttributes .contains (e .getKey ()));
113+
114+ return UpdateExpression .builder ()
115+ .actions (removeActionsFor (removeAttributes ))
116+ .build ();
117+ }
118+
119+ public static final class Builder {
120+
121+ private TableMetadata tableMetadata ;
122+ private Map <String , AttributeValue > nonKeyAttributes ;
123+ private UpdateExpression extensionExpression ;
124+ private UpdateExpression requestExpression ;
125+
126+ public Builder tableMetadata (TableMetadata tableMetadata ) {
127+ this .tableMetadata = requireNonNull (
128+ tableMetadata , "A TableMetadata is required when generating an Update Expression" );
129+ return this ;
130+ }
131+
132+ public Builder nonKeyAttributes (Map <String , AttributeValue > nonKeyAttributes ) {
133+ if (nonKeyAttributes == null ) {
134+ this .nonKeyAttributes = Collections .emptyMap ();
135+ } else {
136+ this .nonKeyAttributes = Collections .unmodifiableMap (new HashMap <>(nonKeyAttributes ));
137+ }
138+ return this ;
139+ }
140+
141+ public Builder extensionExpression (UpdateExpression extensionExpression ) {
142+ this .extensionExpression = extensionExpression ;
143+ return this ;
144+ }
145+
146+ public Builder requestExpression (UpdateExpression requestExpression ) {
147+ this .requestExpression = requestExpression ;
148+ return this ;
149+ }
150+
151+ public UpdateExpressionResolver build () {
152+ return new UpdateExpressionResolver (this );
153+ }
154+
155+ }
156+ }
0 commit comments