Skip to content

Commit 0e85b6b

Browse files
anasatirbasaanasatirbasa
committed
Added support for @DynamoDbAutoGeneratedTimestampAttribute and @DynamoDbUpdateBehavior on attributes within nested objects
1 parent 2800f91 commit 0e85b6b

File tree

13 files changed

+951
-140
lines changed

13 files changed

+951
-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: 109 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,21 @@
1515

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

18+
import static software.amazon.awssdk.enhanced.dynamodb.extensions.utility.NestedRecordUtils.getTableSchemaForListElement;
19+
import static software.amazon.awssdk.enhanced.dynamodb.extensions.utility.NestedRecordUtils.reconstructCompositeKey;
20+
import static software.amazon.awssdk.enhanced.dynamodb.extensions.utility.NestedRecordUtils.resolveSchemasPerPath;
21+
import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.getNestedSchema;
22+
1823
import java.time.Clock;
1924
import java.time.Instant;
2025
import java.util.Collection;
2126
import java.util.Collections;
2227
import java.util.HashMap;
28+
import java.util.List;
2329
import java.util.Map;
30+
import java.util.Optional;
2431
import java.util.function.Consumer;
32+
import java.util.stream.Collectors;
2533
import software.amazon.awssdk.annotations.NotThreadSafe;
2634
import software.amazon.awssdk.annotations.SdkPublicApi;
2735
import software.amazon.awssdk.annotations.ThreadSafe;
@@ -30,6 +38,7 @@
3038
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
3139
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbExtensionContext;
3240
import software.amazon.awssdk.enhanced.dynamodb.EnhancedType;
41+
import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
3342
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
3443
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
3544
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
@@ -64,6 +73,10 @@
6473
* <p>
6574
* Every time a new update of the record is successfully written to the database, the timestamp at which it was modified will
6675
* be automatically updated. This extension applies the conversions as defined in the attribute convertor.
76+
* The implementation handles both flattened nested parameters (identified by keys separated with
77+
* {@code "_NESTED_ATTR_UPDATE_"}) and entire nested maps or lists, ensuring consistent behavior across both representations.
78+
* If a nested object or list is {@code null}, no timestamp values will be generated for any of its annotated fields.
79+
* The same timestamp value is used for both top-level attributes and all applicable nested fields.
6780
*/
6881
@SdkPublicApi
6982
@ThreadSafe
@@ -126,26 +139,113 @@ public static AutoGeneratedTimestampRecordExtension create() {
126139
*/
127140
@Override
128141
public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
142+
Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
143+
144+
Map<String, AttributeValue> updatedItems = new HashMap<>();
145+
Instant currentInstant = clock.instant();
146+
147+
itemToTransform.forEach((key, value) -> {
148+
if (value.hasM() && value.m() != null) {
149+
Optional<? extends TableSchema<?>> nestedSchema = getNestedSchema(context.tableSchema(), key);
150+
if (nestedSchema.isPresent()) {
151+
Map<String, AttributeValue> processed = processNestedObject(value.m(), nestedSchema.get(), currentInstant);
152+
updatedItems.put(key, AttributeValue.builder().m(processed).build());
153+
}
154+
} else if (value.hasL() && !value.l().isEmpty() && value.l().get(0).hasM()) {
155+
TableSchema<?> elementListSchema = getTableSchemaForListElement(context.tableSchema(), key);
156+
157+
List<AttributeValue> updatedList = value.l()
158+
.stream()
159+
.map(listItem -> listItem.hasM() ?
160+
AttributeValue.builder()
161+
.m(processNestedObject(listItem.m(),
162+
elementListSchema,
163+
currentInstant))
164+
.build() : listItem)
165+
.collect(Collectors.toList());
166+
updatedItems.put(key, AttributeValue.builder().l(updatedList).build());
167+
}
168+
});
169+
170+
Map<String, TableSchema<?>> stringTableSchemaMap = resolveSchemasPerPath(itemToTransform, context.tableSchema());
129171

130-
Collection<String> customMetadataObject = context.tableMetadata()
131-
.customMetadataObject(CUSTOM_METADATA_KEY, Collection.class).orElse(null);
172+
stringTableSchemaMap.forEach((path, schema) -> {
173+
Collection<String> customMetadataObject = schema.tableMetadata()
174+
.customMetadataObject(CUSTOM_METADATA_KEY, Collection.class)
175+
.orElse(null);
132176

133-
if (customMetadataObject == null) {
177+
if (customMetadataObject != null) {
178+
customMetadataObject.forEach(
179+
key -> insertTimestampInItemToTransform(updatedItems, reconstructCompositeKey(path, key),
180+
schema.converterForAttribute(key), currentInstant));
181+
}
182+
});
183+
184+
if (updatedItems.isEmpty()) {
134185
return WriteModification.builder().build();
135186
}
136-
Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
137-
customMetadataObject.forEach(
138-
key -> insertTimestampInItemToTransform(itemToTransform, key,
139-
context.tableSchema().converterForAttribute(key)));
187+
188+
itemToTransform.putAll(updatedItems);
189+
140190
return WriteModification.builder()
141191
.transformedItem(Collections.unmodifiableMap(itemToTransform))
142192
.build();
143193
}
144194

195+
private Map<String, AttributeValue> processNestedObject(Map<String, AttributeValue> nestedMap, TableSchema<?> nestedSchema,
196+
Instant currentInstant) {
197+
Map<String, AttributeValue> updatedNestedMap = new HashMap<>(nestedMap);
198+
Collection<String> customMetadataObject = nestedSchema.tableMetadata()
199+
.customMetadataObject(CUSTOM_METADATA_KEY, Collection.class).orElse(null);
200+
201+
if (customMetadataObject != null) {
202+
customMetadataObject.forEach(
203+
key -> insertTimestampInItemToTransform(updatedNestedMap, String.valueOf(key),
204+
nestedSchema.converterForAttribute(key), currentInstant));
205+
}
206+
207+
nestedMap.forEach((nestedKey, nestedValue) -> {
208+
if (nestedValue.hasM()) {
209+
// Resolve the child-attribute schema; default to the current (parent) schema if none is registered
210+
Optional<? extends TableSchema<?>> childSchemaOptional = getNestedSchema(nestedSchema, nestedKey);
211+
TableSchema<?> schemaToUse = childSchemaOptional.isPresent()
212+
? childSchemaOptional.get()
213+
: nestedSchema;
214+
215+
updatedNestedMap.put(nestedKey, AttributeValue.builder()
216+
.m(processNestedObject(
217+
nestedValue.m(), schemaToUse, currentInstant))
218+
.build());
219+
220+
} else if (nestedValue.hasL() && !nestedValue.l().isEmpty()
221+
&& nestedValue.l().get(0).hasM()) {
222+
try {
223+
TableSchema<?> listElementSchema = TableSchema.fromClass(
224+
Class.forName(nestedSchema.converterForAttribute(nestedKey)
225+
.type().rawClassParameters().get(0).rawClass().getName()));
226+
List<AttributeValue> updatedList = nestedValue
227+
.l()
228+
.stream()
229+
.map(listItem -> listItem.hasM() ?
230+
AttributeValue.builder()
231+
.m(processNestedObject(listItem.m(),
232+
listElementSchema,
233+
currentInstant)).build() : listItem)
234+
.collect(Collectors.toList());
235+
updatedNestedMap.put(nestedKey, AttributeValue.builder().l(updatedList).build());
236+
} catch (ClassNotFoundException e) {
237+
throw new IllegalArgumentException("Class not found for field name: " + nestedKey, e);
238+
}
239+
}
240+
});
241+
return updatedNestedMap;
242+
}
243+
145244
private void insertTimestampInItemToTransform(Map<String, AttributeValue> itemToTransform,
146245
String key,
147-
AttributeConverter converter) {
148-
itemToTransform.put(key, converter.transformFrom(clock.instant()));
246+
AttributeConverter converter,
247+
Instant instant) {
248+
itemToTransform.put(key, converter.transformFrom(instant));
149249
}
150250

151251
/**

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

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
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.extensions.utility;
17+
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+
21+
import java.util.HashMap;
22+
import java.util.Map;
23+
import java.util.Optional;
24+
import java.util.regex.Pattern;
25+
import software.amazon.awssdk.annotations.SdkInternalApi;
26+
import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
27+
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
28+
29+
@SdkInternalApi
30+
public final class NestedRecordUtils {
31+
32+
private static final Pattern NESTED_OBJECT_PATTERN = Pattern.compile(NESTED_OBJECT_UPDATE);
33+
34+
private NestedRecordUtils() {
35+
}
36+
37+
/**
38+
* Resolves and returns the {@link TableSchema} for the element type of a list attribute from the provided root schema.
39+
* <p>
40+
* This method is useful when dealing with lists of nested objects in a DynamoDB-enhanced table schema,
41+
* particularly in scenarios where the list is part of a flattened nested structure.
42+
* <p>
43+
* If the provided key contains the nested object delimiter (e.g., {@code _NESTED_ATTR_UPDATE_}), the method traverses
44+
* the nested hierarchy based on that path to locate the correct schema for the target attribute.
45+
* Otherwise, it directly resolves the list element type from the root schema using reflection.
46+
*
47+
* @param rootSchema The root {@link TableSchema} representing the top-level entity.
48+
* @param key The key representing the list attribute, either flat or nested (using a delimiter).
49+
* @return The {@link TableSchema} representing the list element type of the specified attribute.
50+
* @throws IllegalArgumentException If the list element class cannot be found via reflection.
51+
*/
52+
public static TableSchema<?> getTableSchemaForListElement(TableSchema<?> rootSchema, String key) {
53+
TableSchema<?> listElementSchema;
54+
try {
55+
if (!key.contains(NESTED_OBJECT_UPDATE)) {
56+
listElementSchema = TableSchema.fromClass(
57+
Class.forName(rootSchema.converterForAttribute(key).type().rawClassParameters().get(0).rawClass().getName()));
58+
} else {
59+
String[] parts = NESTED_OBJECT_PATTERN.split(key);
60+
TableSchema<?> currentSchema = rootSchema;
61+
62+
for (int i = 0; i < parts.length - 1; i++) {
63+
Optional<? extends TableSchema<?>> nestedSchema = getNestedSchema(currentSchema, parts[i]);
64+
if (nestedSchema.isPresent()) {
65+
currentSchema = nestedSchema.get();
66+
}
67+
}
68+
String attributeName = parts[parts.length - 1];
69+
listElementSchema = TableSchema.fromClass(
70+
Class.forName(currentSchema.converterForAttribute(attributeName)
71+
.type().rawClassParameters().get(0).rawClass().getName()));
72+
}
73+
} catch (ClassNotFoundException e) {
74+
throw new IllegalArgumentException("Class not found for field name: " + key, e);
75+
}
76+
return listElementSchema;
77+
}
78+
79+
/**
80+
* Traverses the attribute keys representing flattened nested structures and resolves the corresponding
81+
* {@link TableSchema} for each nested path.
82+
* <p>
83+
* The method constructs a mapping between each unique nested path (represented as dot-delimited strings)
84+
* and the corresponding {@link TableSchema} object derived from the root schema. It supports resolving schemas
85+
* for arbitrarily deep nesting, using the {@code _NESTED_ATTR_UPDATE_} pattern as a path delimiter.
86+
* <p>
87+
* This is typically used in update or transformation flows where fields from nested objects are represented
88+
* as flattened keys in the attribute map (e.g., {@code parent_NESTED_ATTR_UPDATE_child}).
89+
*
90+
* @param attributesToSet A map of flattened attribute keys to values, where keys may represent paths to nested attributes.
91+
* @param rootSchema The root {@link TableSchema} of the top-level entity.
92+
* @return A map where the key is the nested path (e.g., {@code "parent.child"}) and the value is the {@link TableSchema}
93+
* corresponding to that level in the object hierarchy.
94+
*/
95+
public static Map<String, TableSchema<?>> resolveSchemasPerPath(Map<String, AttributeValue> attributesToSet,
96+
TableSchema<?> rootSchema) {
97+
Map<String, TableSchema<?>> schemaMap = new HashMap<>();
98+
schemaMap.put("", rootSchema);
99+
100+
for (String key : attributesToSet.keySet()) {
101+
String[] parts = NESTED_OBJECT_PATTERN.split(key);
102+
103+
StringBuilder pathBuilder = new StringBuilder();
104+
TableSchema<?> currentSchema = rootSchema;
105+
106+
for (int i = 0; i < parts.length - 1; i++) {
107+
if (pathBuilder.length() > 0) {
108+
pathBuilder.append(".");
109+
}
110+
pathBuilder.append(parts[i]);
111+
112+
String path = pathBuilder.toString();
113+
114+
if (!schemaMap.containsKey(path)) {
115+
Optional<? extends TableSchema<?>> nestedSchema = getNestedSchema(currentSchema, parts[i]);
116+
if (nestedSchema.isPresent()) {
117+
schemaMap.put(path, nestedSchema.get());
118+
currentSchema = nestedSchema.get();
119+
}
120+
} else {
121+
currentSchema = schemaMap.get(path);
122+
}
123+
}
124+
}
125+
return schemaMap;
126+
}
127+
128+
public static String reconstructCompositeKey(String path, String attributeName) {
129+
if (path == null || path.isEmpty()) {
130+
return attributeName;
131+
}
132+
return String.join(NESTED_OBJECT_UPDATE, path.split("\\."))
133+
+ NESTED_OBJECT_UPDATE + attributeName;
134+
}
135+
}

‎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 {

0 commit comments

Comments
 (0)