Skip to content

Commit e719423

Browse files
roamarieiroamariei
committed
Added support for @DynamoDbAutoGeneratedTimestampAttribute and @DynamoDbUpdateBehavior on attributes within nested objects
1 parent 683cff8 commit e719423

File tree

11 files changed

+704
-140
lines changed

11 files changed

+704
-140
lines changed

‎.changes/next-release/feature-AmazonDynamoDBEnhancedClient-96aff9e.json

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"type": "feature",
3+
"category": "Amazon DynamoDB Enhanced Client",
4+
"contributor": "",
5+
"description": "Added support for @DynamoDbAutoGeneratedTimestampAttribute and @DynamoDbUpdateBehavior on attributes within nested objects. The @DynamoDbUpdateBehavior annotation will only take effect for nested attributes when using IgnoreNullsMode.SCALAR_ONLY."
6+
}

‎services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/extensions/AutoGeneratedTimestampRecordExtension.java

Lines changed: 150 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,20 @@
1515

1616
package software.amazon.awssdk.enhanced.dynamodb.extensions;
1717

18+
import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.getNestedSchema;
19+
import static software.amazon.awssdk.enhanced.dynamodb.internal.operations.UpdateItemOperation.NESTED_OBJECT_UPDATE;
20+
1821
import java.time.Clock;
1922
import java.time.Instant;
2023
import java.util.Collection;
2124
import java.util.Collections;
2225
import java.util.HashMap;
26+
import java.util.List;
2327
import java.util.Map;
28+
import java.util.Optional;
2429
import java.util.function.Consumer;
30+
import java.util.regex.Pattern;
31+
import java.util.stream.Collectors;
2532
import software.amazon.awssdk.annotations.NotThreadSafe;
2633
import software.amazon.awssdk.annotations.SdkPublicApi;
2734
import software.amazon.awssdk.annotations.ThreadSafe;
@@ -30,6 +37,7 @@
3037
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
3138
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbExtensionContext;
3239
import software.amazon.awssdk.enhanced.dynamodb.EnhancedType;
40+
import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
3341
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
3442
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
3543
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
@@ -64,13 +72,17 @@
6472
* <p>
6573
* Every time a new update of the record is successfully written to the database, the timestamp at which it was modified will
6674
* be automatically updated. This extension applies the conversions as defined in the attribute convertor.
75+
* The implementation handles both flattened nested parameters (identified by keys separated with
76+
* {@code "_NESTED_ATTR_UPDATE_"}) and entire nested maps or lists, ensuring consistent behavior across both representations.
77+
* The same timestamp value is used for both top-level attributes and all applicable nested fields.
6778
*/
6879
@SdkPublicApi
6980
@ThreadSafe
7081
public final class AutoGeneratedTimestampRecordExtension implements DynamoDbEnhancedClientExtension {
7182
private static final String CUSTOM_METADATA_KEY = "AutoGeneratedTimestampExtension:AutoGeneratedTimestampAttribute";
7283
private static final AutoGeneratedTimestampAttribute
7384
AUTO_GENERATED_TIMESTAMP_ATTRIBUTE = new AutoGeneratedTimestampAttribute();
85+
private static final Pattern NESTED_OBJECT_PATTERN = Pattern.compile(NESTED_OBJECT_UPDATE);
7486
private final Clock clock;
7587

7688
private AutoGeneratedTimestampRecordExtension() {
@@ -126,26 +138,155 @@ public static AutoGeneratedTimestampRecordExtension create() {
126138
*/
127139
@Override
128140
public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
141+
Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
142+
143+
Map<String, AttributeValue> updatedItems = new HashMap<>();
144+
Instant currentInstant = clock.instant();
129145

130-
Collection<String> customMetadataObject = context.tableMetadata()
131-
.customMetadataObject(CUSTOM_METADATA_KEY, Collection.class).orElse(null);
146+
itemToTransform.forEach((key, value) -> {
147+
if (value.hasM() && value.m() != null) {
148+
Optional<? extends TableSchema<?>> nestedSchema = getNestedSchema(context.tableSchema(), key);
149+
if (nestedSchema.isPresent()) {
150+
Map<String, AttributeValue> processed = processNestedObject(value.m(), nestedSchema.get(), currentInstant);
151+
updatedItems.put(key, AttributeValue.builder().m(processed).build());
152+
}
153+
} else if (value.hasL() && !value.l().isEmpty() && value.l().get(0).hasM()) {
154+
TableSchema<?> elementListSchema = getTableSchemaForElementList(context.tableSchema(), key);
155+
156+
List<AttributeValue> updatedList = value.l().stream()
157+
.map(listItem -> listItem.hasM() ?
158+
AttributeValue.builder().m(processNestedObject(listItem.m(), elementListSchema, currentInstant)).build() : listItem)
159+
.collect(Collectors.toList());
160+
updatedItems.put(key, AttributeValue.builder().l(updatedList).build());
161+
}
162+
});
163+
164+
Map<String, TableSchema<?>> stringTableSchemaMap = resolveSchemasPerPath(itemToTransform, context.tableSchema());
165+
166+
stringTableSchemaMap.forEach((path, schema) -> {
167+
Collection<String> customMetadataObject = schema.tableMetadata()
168+
.customMetadataObject(CUSTOM_METADATA_KEY, Collection.class).orElse(null);
169+
170+
if (customMetadataObject != null) {
171+
customMetadataObject.forEach(
172+
key -> insertTimestampInItemToTransform(updatedItems, reconstructCompositeKey(path, key),
173+
schema.converterForAttribute(key), currentInstant));
174+
}
175+
});
132176

133-
if (customMetadataObject == null) {
177+
if (updatedItems.isEmpty()) {
134178
return WriteModification.builder().build();
135179
}
136-
Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
137-
customMetadataObject.forEach(
138-
key -> insertTimestampInItemToTransform(itemToTransform, key,
139-
context.tableSchema().converterForAttribute(key)));
180+
181+
itemToTransform.putAll(updatedItems);
182+
140183
return WriteModification.builder()
141184
.transformedItem(Collections.unmodifiableMap(itemToTransform))
142185
.build();
143186
}
144187

188+
private TableSchema<?> getTableSchemaForElementList(TableSchema<?> rootSchema, String key) {
189+
TableSchema<?> elementListSchema;
190+
try {
191+
if (!key.contains(NESTED_OBJECT_UPDATE)) {
192+
elementListSchema = TableSchema.fromClass(
193+
Class.forName(rootSchema.converterForAttribute(key).type().rawClassParameters().get(0).rawClass().getName()));
194+
} else {
195+
String[] parts = NESTED_OBJECT_PATTERN.split(key);
196+
TableSchema<?> currentSchema = rootSchema;
197+
198+
for (int i = 0; i < parts.length - 1; i++) {
199+
Optional<? extends TableSchema<?>> nestedSchema = getNestedSchema(currentSchema, parts[i]);
200+
if (nestedSchema.isPresent()) {
201+
currentSchema = nestedSchema.get();
202+
}
203+
}
204+
String attributeName = parts[parts.length - 1];
205+
elementListSchema = TableSchema.fromClass(
206+
Class.forName(currentSchema.converterForAttribute(attributeName).type().rawClassParameters().get(0).rawClass().getName()));
207+
}
208+
} catch (ClassNotFoundException e) {
209+
throw new IllegalArgumentException("Class not found for field name: " + key, e);
210+
}
211+
return elementListSchema;
212+
}
213+
214+
private Map<String, TableSchema<?>> resolveSchemasPerPath(Map<String, AttributeValue> attributesToSet,
215+
TableSchema<?> rootSchema) {
216+
Map<String, TableSchema<?>> schemaMap = new HashMap<>();
217+
schemaMap.put("", rootSchema);
218+
219+
for (String key : attributesToSet.keySet()) {
220+
String[] parts = NESTED_OBJECT_PATTERN.split(key);
221+
222+
String path = "";
223+
TableSchema<?> currentSchema = rootSchema;
224+
225+
for (int i = 0; i < parts.length - 1; i++) {
226+
path = path.isEmpty() ? parts[i] : path + "." + parts[i];
227+
228+
if (!schemaMap.containsKey(path)) {
229+
Optional<? extends TableSchema<?>> nestedSchema = getNestedSchema(currentSchema, parts[i]);
230+
if (nestedSchema.isPresent()) {
231+
schemaMap.put(path, nestedSchema.get());
232+
currentSchema = nestedSchema.get();
233+
}
234+
} else {
235+
currentSchema = schemaMap.get(path);
236+
}
237+
}
238+
}
239+
return schemaMap;
240+
}
241+
242+
private static String reconstructCompositeKey(String path, String attributeName) {
243+
if (path == null || path.isEmpty()) {
244+
return attributeName;
245+
}
246+
return String.join(NESTED_OBJECT_UPDATE, path.split("\\."))
247+
+ NESTED_OBJECT_UPDATE + attributeName;
248+
}
249+
250+
private Map<String, AttributeValue> processNestedObject(Map<String, AttributeValue> nestedMap, TableSchema<?> nestedSchema,
251+
Instant currentInstant) {
252+
Map<String, AttributeValue> updatedNestedMap = new HashMap<>(nestedMap);
253+
Collection<String> customMetadataObject = nestedSchema.tableMetadata()
254+
.customMetadataObject(CUSTOM_METADATA_KEY, Collection.class).orElse(null);
255+
256+
if (customMetadataObject!= null) {
257+
customMetadataObject.forEach(
258+
key -> insertTimestampInItemToTransform(updatedNestedMap, String.valueOf(key),
259+
nestedSchema.converterForAttribute(key), currentInstant));
260+
}
261+
262+
nestedMap.forEach((nestedKey, nestedValue) -> {
263+
if (nestedValue.hasM()) {
264+
updatedNestedMap.put(nestedKey,
265+
AttributeValue.builder().m(processNestedObject(nestedValue.m(), nestedSchema,
266+
currentInstant)).build());
267+
} else if (nestedValue.hasL() && !nestedValue.l().isEmpty()
268+
&& nestedValue.l().get(0).hasM()) {
269+
try {
270+
TableSchema<?> listElementSchema = TableSchema.fromClass(
271+
Class.forName(nestedSchema.converterForAttribute(nestedKey).type().rawClassParameters().get(0).rawClass().getName()));
272+
List<AttributeValue> updatedList = nestedValue.l().stream()
273+
.map(listItem -> listItem.hasM() ?
274+
AttributeValue.builder().m(processNestedObject(listItem.m(), listElementSchema, currentInstant)).build() : listItem)
275+
.collect(Collectors.toList());
276+
updatedNestedMap.put(nestedKey, AttributeValue.builder().l(updatedList).build());
277+
} catch (ClassNotFoundException e) {
278+
throw new IllegalArgumentException("Class not found for field name: " + nestedKey, e);
279+
}
280+
}
281+
});
282+
return updatedNestedMap;
283+
}
284+
145285
private void insertTimestampInItemToTransform(Map<String, AttributeValue> itemToTransform,
146286
String key,
147-
AttributeConverter converter) {
148-
itemToTransform.put(key, converter.transformFrom(clock.instant()));
287+
AttributeConverter converter,
288+
Instant instant) {
289+
itemToTransform.put(key, converter.transformFrom(instant));
149290
}
150291

151292
/**

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,4 +204,15 @@ public static <T> List<T> getItemsFromSupplier(List<Supplier<T>> itemSupplierLis
204204
public static boolean isNullAttributeValue(AttributeValue attributeValue) {
205205
return attributeValue.nul() != null && attributeValue.nul();
206206
}
207+
208+
/**
209+
* Retrieves the {@link TableSchema} for a nested attribute within the given parent schema.
210+
*
211+
* @param parentSchema the schema of the parent bean class
212+
* @param attributeName the name of the nested attribute
213+
* @return an {@link Optional} containing the nested attribute's {@link TableSchema}, or empty if unavailable
214+
*/
215+
public static Optional<? extends TableSchema<?>> getNestedSchema(TableSchema<?> parentSchema, String attributeName) {
216+
return parentSchema.converterForAttribute(attributeName).type().tableSchema();
217+
}
207218
}

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -131,8 +131,7 @@ public UpdateItemRequest generateRequest(TableSchema<T> tableSchema,
131131

132132
Map<String, AttributeValue> keyAttributes = filterMap(itemMap, entry -> primaryKeys.contains(entry.getKey()));
133133
Map<String, AttributeValue> nonKeyAttributes = filterMap(itemMap, entry -> !primaryKeys.contains(entry.getKey()));
134-
135-
Expression updateExpression = generateUpdateExpressionIfExist(tableMetadata, transformation, nonKeyAttributes);
134+
Expression updateExpression = generateUpdateExpressionIfExist(tableSchema, transformation, nonKeyAttributes);
136135
Expression conditionExpression = generateConditionExpressionIfExist(transformation, request);
137136

138137
Map<String, String> expressionNames = coalesceExpressionNames(updateExpression, conditionExpression);
@@ -275,7 +274,7 @@ public TransactWriteItem generateTransactWriteItem(TableSchema<T> tableSchema, O
275274
* if there are attributes to be updated (most likely). If both exist, they are merged and the code generates a final
276275
* Expression that represent the result.
277276
*/
278-
private Expression generateUpdateExpressionIfExist(TableMetadata tableMetadata,
277+
private Expression generateUpdateExpressionIfExist(TableSchema<T> tableSchema,
279278
WriteModification transformation,
280279
Map<String, AttributeValue> attributes) {
281280
UpdateExpression updateExpression = null;
@@ -284,7 +283,7 @@ private Expression generateUpdateExpressionIfExist(TableMetadata tableMetadata,
284283
}
285284
if (!attributes.isEmpty()) {
286285
List<String> nonRemoveAttributes = UpdateExpressionConverter.findAttributeNames(updateExpression);
287-
UpdateExpression operationUpdateExpression = operationExpression(attributes, tableMetadata, nonRemoveAttributes);
286+
UpdateExpression operationUpdateExpression = operationExpression(attributes, tableSchema, nonRemoveAttributes);
288287
if (updateExpression == null) {
289288
updateExpression = operationUpdateExpression;
290289
} else {

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

Lines changed: 30 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -15,21 +15,24 @@
1515

1616
package software.amazon.awssdk.enhanced.dynamodb.internal.update;
1717

18+
import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.getNestedSchema;
1819
import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.isNullAttributeValue;
1920
import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.keyRef;
2021
import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.valueRef;
2122
import static software.amazon.awssdk.enhanced.dynamodb.internal.operations.UpdateItemOperation.NESTED_OBJECT_UPDATE;
2223
import static software.amazon.awssdk.utils.CollectionUtils.filterMap;
2324

25+
import java.util.ArrayList;
2426
import java.util.Arrays;
2527
import java.util.Collections;
2628
import java.util.List;
2729
import java.util.Map;
30+
import java.util.Optional;
2831
import java.util.function.Function;
2932
import java.util.regex.Pattern;
3033
import java.util.stream.Collectors;
3134
import software.amazon.awssdk.annotations.SdkInternalApi;
32-
import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
35+
import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
3336
import software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils;
3437
import software.amazon.awssdk.enhanced.dynamodb.internal.mapper.UpdateBehaviorTag;
3538
import software.amazon.awssdk.enhanced.dynamodb.mapper.UpdateBehavior;
@@ -57,12 +60,12 @@ public static String ifNotExists(String key, String initValue) {
5760
* Generates an UpdateExpression representing a POJO, with only SET and REMOVE actions.
5861
*/
5962
public static UpdateExpression operationExpression(Map<String, AttributeValue> itemMap,
60-
TableMetadata tableMetadata,
63+
TableSchema tableSchema,
6164
List<String> nonRemoveAttributes) {
6265

6366
Map<String, AttributeValue> setAttributes = filterMap(itemMap, e -> !isNullAttributeValue(e.getValue()));
6467
UpdateExpression setAttributeExpression = UpdateExpression.builder()
65-
.actions(setActionsFor(setAttributes, tableMetadata))
68+
.actions(setActionsFor(setAttributes, tableSchema))
6669
.build();
6770

6871
Map<String, AttributeValue> removeAttributes =
@@ -78,13 +81,30 @@ public static UpdateExpression operationExpression(Map<String, AttributeValue> i
7881
/**
7982
* Creates a list of SET actions for all attributes supplied in the map.
8083
*/
81-
private static List<SetAction> setActionsFor(Map<String, AttributeValue> attributesToSet, TableMetadata tableMetadata) {
82-
return attributesToSet.entrySet()
83-
.stream()
84-
.map(entry -> setValue(entry.getKey(),
85-
entry.getValue(),
86-
UpdateBehaviorTag.resolveForAttribute(entry.getKey(), tableMetadata)))
87-
.collect(Collectors.toList());
84+
private static List<SetAction> setActionsFor(Map<String, AttributeValue> attributesToSet, TableSchema tableSchema) {
85+
List<SetAction> actions = new ArrayList<>();
86+
for (Map.Entry<String, AttributeValue> entry : attributesToSet.entrySet()) {
87+
String key = entry.getKey();
88+
AttributeValue value = entry.getValue();
89+
90+
if (key.contains(NESTED_OBJECT_UPDATE)) {
91+
TableSchema currentSchema = tableSchema;
92+
List<String> pathFieldNames = Arrays.asList(PATTERN.split(key));
93+
String attributeName = pathFieldNames.get(pathFieldNames.size()-1);
94+
95+
for(int i = 0; i < pathFieldNames.size() - 1; i++ ) {
96+
Optional<? extends TableSchema<?>> nestedSchema = getNestedSchema(currentSchema, pathFieldNames.get(i));
97+
if (nestedSchema.isPresent()) {
98+
currentSchema = nestedSchema.get();
99+
}
100+
}
101+
102+
actions.add(setValue(key, value, UpdateBehaviorTag.resolveForAttribute(attributeName, currentSchema.tableMetadata())));
103+
} else {
104+
actions.add(setValue(key, value, UpdateBehaviorTag.resolveForAttribute(key, tableSchema.tableMetadata())));
105+
}
106+
}
107+
return actions;
88108
}
89109

90110
/**

0 commit comments

Comments
 (0)