Skip to content

Support AutoGeneratedTimestamp and UpdateBehavior annotation in nested objects #6109

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
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": "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."
}
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,20 @@

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

import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.getNestedSchema;
import static software.amazon.awssdk.enhanced.dynamodb.internal.operations.UpdateItemOperation.NESTED_OBJECT_UPDATE;

import java.time.Clock;
import java.time.Instant;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Consumer;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import software.amazon.awssdk.annotations.NotThreadSafe;
import software.amazon.awssdk.annotations.SdkPublicApi;
import software.amazon.awssdk.annotations.ThreadSafe;
Expand All @@ -30,6 +37,7 @@
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbExtensionContext;
import software.amazon.awssdk.enhanced.dynamodb.EnhancedType;
import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
Expand Down Expand Up @@ -64,13 +72,17 @@
* <p>
* Every time a new update of the record is successfully written to the database, the timestamp at which it was modified will
* be automatically updated. This extension applies the conversions as defined in the attribute convertor.
* The implementation handles both flattened nested parameters (identified by keys separated with
* {@code "_NESTED_ATTR_UPDATE_"}) and entire nested maps or lists, ensuring consistent behavior across both representations.
* The same timestamp value is used for both top-level attributes and all applicable nested fields.
*/
@SdkPublicApi
@ThreadSafe
public final class AutoGeneratedTimestampRecordExtension implements DynamoDbEnhancedClientExtension {
private static final String CUSTOM_METADATA_KEY = "AutoGeneratedTimestampExtension:AutoGeneratedTimestampAttribute";
private static final AutoGeneratedTimestampAttribute
AUTO_GENERATED_TIMESTAMP_ATTRIBUTE = new AutoGeneratedTimestampAttribute();
private static final Pattern NESTED_OBJECT_PATTERN = Pattern.compile(NESTED_OBJECT_UPDATE);
private final Clock clock;

private AutoGeneratedTimestampRecordExtension() {
Expand Down Expand Up @@ -126,26 +138,155 @@ public static AutoGeneratedTimestampRecordExtension create() {
*/
@Override
public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());

Map<String, AttributeValue> updatedItems = new HashMap<>();
Instant currentInstant = clock.instant();

Collection<String> customMetadataObject = context.tableMetadata()
.customMetadataObject(CUSTOM_METADATA_KEY, Collection.class).orElse(null);
itemToTransform.forEach((key, value) -> {
if (value.hasM() && value.m() != null) {
Optional<? extends TableSchema<?>> nestedSchema = getNestedSchema(context.tableSchema(), key);
if (nestedSchema.isPresent()) {
Map<String, AttributeValue> processed = processNestedObject(value.m(), nestedSchema.get(), currentInstant);
updatedItems.put(key, AttributeValue.builder().m(processed).build());
}
} else if (value.hasL() && !value.l().isEmpty() && value.l().get(0).hasM()) {
TableSchema<?> elementListSchema = getTableSchemaForListElement(context.tableSchema(), key);

List<AttributeValue> updatedList = value.l().stream()
.map(listItem -> listItem.hasM() ?
AttributeValue.builder().m(processNestedObject(listItem.m(), elementListSchema, currentInstant)).build() : listItem)
.collect(Collectors.toList());
updatedItems.put(key, AttributeValue.builder().l(updatedList).build());
}
});

Map<String, TableSchema<?>> stringTableSchemaMap = resolveSchemasPerPath(itemToTransform, context.tableSchema());

stringTableSchemaMap.forEach((path, schema) -> {
Collection<String> customMetadataObject = schema.tableMetadata()
.customMetadataObject(CUSTOM_METADATA_KEY, Collection.class).orElse(null);

if (customMetadataObject != null) {
customMetadataObject.forEach(
key -> insertTimestampInItemToTransform(updatedItems, reconstructCompositeKey(path, key),
schema.converterForAttribute(key), currentInstant));
}
});

if (customMetadataObject == null) {
if (updatedItems.isEmpty()) {
return WriteModification.builder().build();
}
Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
customMetadataObject.forEach(
key -> insertTimestampInItemToTransform(itemToTransform, key,
context.tableSchema().converterForAttribute(key)));

itemToTransform.putAll(updatedItems);

return WriteModification.builder()
.transformedItem(Collections.unmodifiableMap(itemToTransform))
.build();
}

private TableSchema<?> getTableSchemaForListElement(TableSchema<?> rootSchema, String key) {
TableSchema<?> listElementSchema;
try {
if (!key.contains(NESTED_OBJECT_UPDATE)) {
listElementSchema = TableSchema.fromClass(
Class.forName(rootSchema.converterForAttribute(key).type().rawClassParameters().get(0).rawClass().getName()));
} else {
String[] parts = NESTED_OBJECT_PATTERN.split(key);
TableSchema<?> currentSchema = rootSchema;

for (int i = 0; i < parts.length - 1; i++) {
Optional<? extends TableSchema<?>> nestedSchema = getNestedSchema(currentSchema, parts[i]);
if (nestedSchema.isPresent()) {
currentSchema = nestedSchema.get();
}
}
String attributeName = parts[parts.length - 1];
listElementSchema = TableSchema.fromClass(
Class.forName(currentSchema.converterForAttribute(attributeName).type().rawClassParameters().get(0).rawClass().getName()));
}
} catch (ClassNotFoundException e) {
throw new IllegalArgumentException("Class not found for field name: " + key, e);
}
return listElementSchema;
}

private Map<String, TableSchema<?>> resolveSchemasPerPath(Map<String, AttributeValue> attributesToSet,
TableSchema<?> rootSchema) {
Map<String, TableSchema<?>> schemaMap = new HashMap<>();
schemaMap.put("", rootSchema);

for (String key : attributesToSet.keySet()) {
String[] parts = NESTED_OBJECT_PATTERN.split(key);

String path = "";
TableSchema<?> currentSchema = rootSchema;

for (int i = 0; i < parts.length - 1; i++) {
path = path.isEmpty() ? parts[i] : path + "." + parts[i];

if (!schemaMap.containsKey(path)) {
Optional<? extends TableSchema<?>> nestedSchema = getNestedSchema(currentSchema, parts[i]);
if (nestedSchema.isPresent()) {
schemaMap.put(path, nestedSchema.get());
currentSchema = nestedSchema.get();
}
} else {
currentSchema = schemaMap.get(path);
}
}
}
return schemaMap;
}

private static String reconstructCompositeKey(String path, String attributeName) {
if (path == null || path.isEmpty()) {
return attributeName;
}
return String.join(NESTED_OBJECT_UPDATE, path.split("\\."))
+ NESTED_OBJECT_UPDATE + attributeName;
}

private Map<String, AttributeValue> processNestedObject(Map<String, AttributeValue> nestedMap, TableSchema<?> nestedSchema,
Instant currentInstant) {
Map<String, AttributeValue> updatedNestedMap = new HashMap<>(nestedMap);
Collection<String> customMetadataObject = nestedSchema.tableMetadata()
.customMetadataObject(CUSTOM_METADATA_KEY, Collection.class).orElse(null);

if (customMetadataObject!= null) {
customMetadataObject.forEach(
key -> insertTimestampInItemToTransform(updatedNestedMap, String.valueOf(key),
nestedSchema.converterForAttribute(key), currentInstant));
}

nestedMap.forEach((nestedKey, nestedValue) -> {
if (nestedValue.hasM()) {
updatedNestedMap.put(nestedKey,
AttributeValue.builder().m(processNestedObject(nestedValue.m(), nestedSchema,
currentInstant)).build());
} else if (nestedValue.hasL() && !nestedValue.l().isEmpty()
&& nestedValue.l().get(0).hasM()) {
try {
TableSchema<?> listElementSchema = TableSchema.fromClass(
Class.forName(nestedSchema.converterForAttribute(nestedKey).type().rawClassParameters().get(0).rawClass().getName()));
List<AttributeValue> updatedList = nestedValue.l().stream()
.map(listItem -> listItem.hasM() ?
AttributeValue.builder().m(processNestedObject(listItem.m(), listElementSchema, currentInstant)).build() : listItem)
.collect(Collectors.toList());
updatedNestedMap.put(nestedKey, AttributeValue.builder().l(updatedList).build());
} catch (ClassNotFoundException e) {
throw new IllegalArgumentException("Class not found for field name: " + nestedKey, e);
}
}
});
return updatedNestedMap;
}

private void insertTimestampInItemToTransform(Map<String, AttributeValue> itemToTransform,
String key,
AttributeConverter converter) {
itemToTransform.put(key, converter.transformFrom(clock.instant()));
AttributeConverter converter,
Instant instant) {
itemToTransform.put(key, converter.transformFrom(instant));
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -204,4 +204,15 @@ public static <T> List<T> getItemsFromSupplier(List<Supplier<T>> itemSupplierLis
public static boolean isNullAttributeValue(AttributeValue attributeValue) {
return attributeValue.nul() != null && attributeValue.nul();
}

/**
* Retrieves the {@link TableSchema} for a nested attribute within the given parent schema.
*
* @param parentSchema the schema of the parent bean class
* @param attributeName the name of the nested attribute
* @return an {@link Optional} containing the nested attribute's {@link TableSchema}, or empty if unavailable
*/
public static Optional<? extends TableSchema<?>> getNestedSchema(TableSchema<?> parentSchema, String attributeName) {
return parentSchema.converterForAttribute(attributeName).type().tableSchema();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -131,8 +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(tableSchema, transformation, nonKeyAttributes);
Expression conditionExpression = generateConditionExpressionIfExist(transformation, request);

Map<String, String> expressionNames = coalesceExpressionNames(updateExpression, conditionExpression);
Expand Down Expand Up @@ -275,7 +274,7 @@ public TransactWriteItem generateTransactWriteItem(TableSchema<T> tableSchema, O
* 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.
*/
private Expression generateUpdateExpressionIfExist(TableMetadata tableMetadata,
private Expression generateUpdateExpressionIfExist(TableSchema<T> tableSchema,
WriteModification transformation,
Map<String, AttributeValue> attributes) {
UpdateExpression updateExpression = null;
Expand All @@ -284,7 +283,7 @@ private Expression generateUpdateExpressionIfExist(TableMetadata tableMetadata,
}
if (!attributes.isEmpty()) {
List<String> nonRemoveAttributes = UpdateExpressionConverter.findAttributeNames(updateExpression);
UpdateExpression operationUpdateExpression = operationExpression(attributes, tableMetadata, nonRemoveAttributes);
UpdateExpression operationUpdateExpression = operationExpression(attributes, tableSchema, nonRemoveAttributes);
if (updateExpression == null) {
updateExpression = operationUpdateExpression;
} else {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,21 +15,24 @@

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

import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.getNestedSchema;
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.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import software.amazon.awssdk.annotations.SdkInternalApi;
import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
import software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils;
import software.amazon.awssdk.enhanced.dynamodb.internal.mapper.UpdateBehaviorTag;
import software.amazon.awssdk.enhanced.dynamodb.mapper.UpdateBehavior;
Expand Down Expand Up @@ -57,12 +60,12 @@ public static String ifNotExists(String key, String initValue) {
* Generates an UpdateExpression representing a POJO, with only SET and REMOVE actions.
*/
public static UpdateExpression operationExpression(Map<String, AttributeValue> itemMap,
TableMetadata tableMetadata,
TableSchema tableSchema,
List<String> nonRemoveAttributes) {

Map<String, AttributeValue> setAttributes = filterMap(itemMap, e -> !isNullAttributeValue(e.getValue()));
UpdateExpression setAttributeExpression = UpdateExpression.builder()
.actions(setActionsFor(setAttributes, tableMetadata))
.actions(setActionsFor(setAttributes, tableSchema))
.build();

Map<String, AttributeValue> removeAttributes =
Expand All @@ -78,13 +81,30 @@ public static UpdateExpression operationExpression(Map<String, AttributeValue> i
/**
* Creates a list of SET actions for all attributes supplied in the map.
*/
private static List<SetAction> setActionsFor(Map<String, AttributeValue> attributesToSet, TableMetadata tableMetadata) {
return attributesToSet.entrySet()
.stream()
.map(entry -> setValue(entry.getKey(),
entry.getValue(),
UpdateBehaviorTag.resolveForAttribute(entry.getKey(), tableMetadata)))
.collect(Collectors.toList());
private static List<SetAction> setActionsFor(Map<String, AttributeValue> attributesToSet, TableSchema tableSchema) {
List<SetAction> actions = new ArrayList<>();
for (Map.Entry<String, AttributeValue> entry : attributesToSet.entrySet()) {
String key = entry.getKey();
AttributeValue value = entry.getValue();

if (key.contains(NESTED_OBJECT_UPDATE)) {
TableSchema currentSchema = tableSchema;
List<String> pathFieldNames = Arrays.asList(PATTERN.split(key));
String attributeName = pathFieldNames.get(pathFieldNames.size()-1);

for(int i = 0; i < pathFieldNames.size() - 1; i++ ) {
Optional<? extends TableSchema<?>> nestedSchema = getNestedSchema(currentSchema, pathFieldNames.get(i));
if (nestedSchema.isPresent()) {
currentSchema = nestedSchema.get();
}
}

actions.add(setValue(key, value, UpdateBehaviorTag.resolveForAttribute(attributeName, currentSchema.tableMetadata())));
} else {
actions.add(setValue(key, value, UpdateBehaviorTag.resolveForAttribute(key, tableSchema.tableMetadata())));
}
}
return actions;
}

/**
Expand Down
Loading