Skip to content

Commit 1592f4d

Browse files
anasatirbasaanasatirbasa
committed
Support update expressions in single request update
1 parent 1a8be21 commit 1592f4d

File tree

3 files changed

+42
-31
lines changed

3 files changed

+42
-31
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -296,7 +296,7 @@ private Expression generateUpdateExpressionIfExist(
296296
UpdateExpressionResolver updateExpressionResolver =
297297
UpdateExpressionResolver.builder()
298298
.tableMetadata(tableMetadata)
299-
.itemNonKeyAttributes(attributes)
299+
.nonKeyAttributes(attributes)
300300
.requestExpression(requestUpdateExpression)
301301
.extensionExpression(transformation != null ? transformation.updateExpression() : null)
302302
.build();

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

Lines changed: 31 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,13 @@
2323

2424
import java.util.Arrays;
2525
import java.util.Collections;
26+
import java.util.HashMap;
2627
import java.util.List;
2728
import java.util.Map;
2829
import java.util.Objects;
2930
import java.util.Set;
3031
import java.util.stream.Collectors;
32+
import java.util.stream.Stream;
3133
import software.amazon.awssdk.annotations.SdkInternalApi;
3234
import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
3335
import software.amazon.awssdk.enhanced.dynamodb.update.UpdateExpression;
@@ -40,16 +42,16 @@
4042
@SdkInternalApi
4143
public final class UpdateExpressionResolver {
4244

45+
private final TableMetadata tableMetadata;
46+
private final Map<String, AttributeValue> nonKeyAttributes;
4347
private final UpdateExpression extensionExpression;
4448
private final UpdateExpression requestExpression;
45-
private final Map<String, AttributeValue> itemNonKeyAttributes;
46-
private final TableMetadata tableMetadata;
4749

4850
private UpdateExpressionResolver(Builder builder) {
51+
this.tableMetadata = builder.tableMetadata;
52+
this.nonKeyAttributes = builder.nonKeyAttributes;
4953
this.extensionExpression = builder.extensionExpression;
5054
this.requestExpression = builder.requestExpression;
51-
this.itemNonKeyAttributes = builder.nonKeyAttributes;
52-
this.tableMetadata = builder.tableMetadata;
5355
}
5456

5557
public static Builder builder() {
@@ -69,17 +71,24 @@ public static Builder builder() {
6971
* @return merged UpdateExpression, or empty if no updates needed
7072
*/
7173
public UpdateExpression resolve() {
72-
Set<String> excludedFromRemoval = attributesPresentInExpressions(Arrays.asList(extensionExpression, requestExpression));
74+
UpdateExpression itemExpression = null;
7375

74-
UpdateExpression itemSetExpression = generateItemSetExpression(itemNonKeyAttributes, tableMetadata);
75-
UpdateExpression itemRemoveExpression = generateItemRemoveExpression(itemNonKeyAttributes, excludedFromRemoval);
76-
UpdateExpression itemFinalExpression = UpdateExpression.mergeExpressions(itemSetExpression, itemRemoveExpression);
76+
if (!nonKeyAttributes.isEmpty()) {
77+
Set<String> attributesExcludedFromRemoval = attributesPresentInOtherExpressions(
78+
Arrays.asList(extensionExpression, requestExpression));
7779

78-
UpdateExpression itemAndExtensionExpression = UpdateExpression.mergeExpressions(extensionExpression, itemFinalExpression);
79-
return UpdateExpression.mergeExpressions(requestExpression, itemAndExtensionExpression);
80+
itemExpression = UpdateExpression.mergeExpressions(
81+
generateItemSetExpression(nonKeyAttributes, tableMetadata),
82+
generateItemRemoveExpression(nonKeyAttributes, attributesExcludedFromRemoval));
83+
}
84+
85+
return Stream.of(itemExpression, extensionExpression, requestExpression)
86+
.filter(Objects::nonNull)
87+
.reduce(UpdateExpression::mergeExpressions)
88+
.orElse(null);
8089
}
8190

82-
private static Set<String> attributesPresentInExpressions(List<UpdateExpression> updateExpressions) {
91+
private static Set<String> attributesPresentInOtherExpressions(List<UpdateExpression> updateExpressions) {
8392
return updateExpressions.stream()
8493
.filter(Objects::nonNull)
8594
.map(UpdateExpressionConverter::findAttributeNames)
@@ -109,26 +118,27 @@ public static UpdateExpression generateItemRemoveExpression(Map<String, Attribut
109118
public static final class Builder {
110119

111120
private TableMetadata tableMetadata;
121+
private Map<String, AttributeValue> nonKeyAttributes;
112122
private UpdateExpression extensionExpression;
113123
private UpdateExpression requestExpression;
114-
private Map<String, AttributeValue> nonKeyAttributes;
115124

116125
public Builder tableMetadata(TableMetadata tableMetadata) {
117-
requireNonNull(tableMetadata, "A TableMetadata is required when generating an Update Expression");
118-
this.tableMetadata = tableMetadata;
126+
this.tableMetadata = requireNonNull(
127+
tableMetadata, "A TableMetadata is required when generating an Update Expression");
119128
return this;
120129
}
121130

122-
public Builder extensionExpression(UpdateExpression extensionExpression) {
123-
this.extensionExpression = extensionExpression;
131+
public Builder nonKeyAttributes(Map<String, AttributeValue> nonKeyAttributes) {
132+
if (nonKeyAttributes == null) {
133+
this.nonKeyAttributes = Collections.emptyMap();
134+
} else {
135+
this.nonKeyAttributes = Collections.unmodifiableMap(new HashMap<>(nonKeyAttributes));
136+
}
124137
return this;
125138
}
126139

127-
public Builder itemNonKeyAttributes(Map<String, AttributeValue> nonKeyAttributes) {
128-
if (nonKeyAttributes == null) {
129-
nonKeyAttributes = Collections.emptyMap();
130-
}
131-
this.nonKeyAttributes = nonKeyAttributes;
140+
public Builder extensionExpression(UpdateExpression extensionExpression) {
141+
this.extensionExpression = extensionExpression;
132142
return this;
133143
}
134144

‎services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/update/UpdateExpressionResolverTest.java‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
package software.amazon.awssdk.enhanced.dynamodb.internal.update;
1717

1818
import static org.assertj.core.api.Assertions.assertThat;
19+
import static org.junit.Assert.assertNull;
1920
import static org.mockito.Mockito.mock;
2021

2122
import java.util.Collections;
@@ -37,12 +38,12 @@ public class UpdateExpressionResolverTest {
3738
public void resolve_emptyInputs_returnsEmptyUpdateExpression() {
3839
UpdateExpressionResolver resolver = UpdateExpressionResolver.builder()
3940
.tableMetadata(mockTableMetadata)
40-
.itemNonKeyAttributes(Collections.emptyMap())
41+
.nonKeyAttributes(Collections.emptyMap())
4142
.build();
4243

4344
UpdateExpression result = resolver.resolve();
4445

45-
assertThat(result).isEqualTo(UpdateExpression.builder().build());
46+
assertNull(result);
4647
}
4748

4849
@Test
@@ -53,7 +54,7 @@ public void resolve_nonNullAttributes_generatesSetActions() {
5354

5455
UpdateExpressionResolver resolver = UpdateExpressionResolver.builder()
5556
.tableMetadata(mockTableMetadata)
56-
.itemNonKeyAttributes(itemMap)
57+
.nonKeyAttributes(itemMap)
5758
.build();
5859

5960
UpdateExpression result = resolver.resolve();
@@ -87,7 +88,7 @@ public void resolve_nullAttributes_generatesRemoveActions() {
8788

8889
UpdateExpressionResolver resolver = UpdateExpressionResolver.builder()
8990
.tableMetadata(mockTableMetadata)
90-
.itemNonKeyAttributes(itemMap)
91+
.nonKeyAttributes(itemMap)
9192
.build();
9293

9394
UpdateExpression result = resolver.resolve();
@@ -117,7 +118,7 @@ public void resolve_mixedAttributes_generatesBothActions() {
117118

118119
UpdateExpressionResolver resolver = UpdateExpressionResolver.builder()
119120
.tableMetadata(mockTableMetadata)
120-
.itemNonKeyAttributes(itemMap)
121+
.nonKeyAttributes(itemMap)
121122
.build();
122123

123124
UpdateExpression result = resolver.resolve();
@@ -158,7 +159,7 @@ public void resolve_withItemAndExtensionExpression_mergesActions() {
158159

159160
UpdateExpressionResolver resolver = UpdateExpressionResolver.builder()
160161
.tableMetadata(mockTableMetadata)
161-
.itemNonKeyAttributes(itemMap)
162+
.nonKeyAttributes(itemMap)
162163
.extensionExpression(extensionExpression)
163164
.build();
164165

@@ -211,7 +212,7 @@ public void resolve_withAllExpressionTypes_mergesInCorrectOrder() {
211212

212213
UpdateExpressionResolver resolver = UpdateExpressionResolver.builder()
213214
.tableMetadata(mockTableMetadata)
214-
.itemNonKeyAttributes(itemMap)
215+
.nonKeyAttributes(itemMap)
215216
.extensionExpression(extensionExpression)
216217
.requestExpression(requestExpression)
217218
.build();
@@ -262,7 +263,7 @@ public void resolve_attributeUsedInOtherExpression_filteredOutFromRemoveActions(
262263

263264
UpdateExpressionResolver resolver = UpdateExpressionResolver.builder()
264265
.tableMetadata(mockTableMetadata)
265-
.itemNonKeyAttributes(itemMap)
266+
.nonKeyAttributes(itemMap)
266267
.requestExpression(requestExpression)
267268
.build();
268269

@@ -360,7 +361,7 @@ public void builder_allFields_buildsSuccessfully() {
360361

361362
UpdateExpressionResolver resolver = UpdateExpressionResolver.builder()
362363
.tableMetadata(mockTableMetadata)
363-
.itemNonKeyAttributes(itemMap)
364+
.nonKeyAttributes(itemMap)
364365
.extensionExpression(extensionExpr)
365366
.requestExpression(requestExpr)
366367
.build();

0 commit comments

Comments
 (0)