Support AutoGeneratedTimestamp and UpdateBehavior annotation in nested objects

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

@@ -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."
+}

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

@@ -15,13 +15,21 @@
 
 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.extensions.utility.NestedRecordUtils.getTableSchemaForListElement;
+import static software.amazon.awssdk.enhanced.dynamodb.internal.extensions.utility.NestedRecordUtils.reconstructCompositeKey;
+import static software.amazon.awssdk.enhanced.dynamodb.internal.extensions.utility.NestedRecordUtils.resolveSchemasPerPath;
+
 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.stream.Collectors;
 import software.amazon.awssdk.annotations.NotThreadSafe;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
@@ -30,6 +38,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;
@@ -64,6 +73,10 @@
  * <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.
+ * If a nested object or list is {@code null}, no timestamp values will be generated for any of its annotated fields.
+ * The same timestamp value is used for both top-level attributes and all applicable nested fields.
  */
 @SdkPublicApi
 @ThreadSafe
@@ -126,26 +139,103 @@ 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();
+
+        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());
 
-        Collection<String> customMetadataObject = context.tableMetadata()
-                                       .customMetadataObject(CUSTOM_METADATA_KEY, Collection.class).orElse(null);
+        stringTableSchemaMap.forEach((path, schema) -> {
+            Collection<String> customMetadataObject = schema.tableMetadata()
+                                                            .customMetadataObject(CUSTOM_METADATA_KEY, Collection.class)
+                                                            .orElse(null);
 
-        if (customMetadataObject == null) {
+            if (customMetadataObject != null) {
+                customMetadataObject.forEach(
+                    key -> insertTimestampInItemToTransform(updatedItems, reconstructCompositeKey(path, key),
+                                                            schema.converterForAttribute(key), currentInstant));
+            }
+        });
+
+        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 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()) {
+                Optional<? extends TableSchema<?>> childSchemaOptional = getNestedSchema(nestedSchema, nestedKey);
+                TableSchema<?> schemaToUse = childSchemaOptional.isPresent() ? childSchemaOptional.get() : nestedSchema;
+                updatedNestedMap.put(nestedKey,
+                                     AttributeValue.builder()
+                                                   .m(processNestedObject(nestedValue.m(), schemaToUse, currentInstant))
+                                                   .build());
+
+            } else if (nestedValue.hasL() && !nestedValue.l().isEmpty() && nestedValue.l().get(0).hasM()) {
+                TableSchema<?> listElementSchema = getTableSchemaForListElement(nestedSchema, nestedKey);
+                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());
+            }
+        });
+        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));
     }
 
     /**

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

@@ -29,6 +29,7 @@
 import java.util.stream.Stream;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
+import software.amazon.awssdk.enhanced.dynamodb.EnhancedType;
 import software.amazon.awssdk.enhanced.dynamodb.Key;
 import software.amazon.awssdk.enhanced.dynamodb.OperationContext;
 import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
@@ -204,4 +205,24 @@ 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. When the attribute is a
+     * parameterized type (e.g., List<?>), it retrieves the schema of the first type parameter. Otherwise, it retrieves the schema
+     * directly from the attribute's enhanced type.
+     *
+     * @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) {
+        EnhancedType<?> enhancedType = parentSchema.converterForAttribute(attributeName).type();
+        List<EnhancedType<?>> rawClassParameters = enhancedType.rawClassParameters();
+
+        if (rawClassParameters != null && !rawClassParameters.isEmpty()) {
+            enhancedType = rawClassParameters.get(0);
+        }
+
+        return enhancedType.tableSchema();
+    }
 }

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

@@ -0,0 +1,140 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.enhanced.dynamodb.internal.extensions.utility;
+
+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.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.regex.Pattern;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
+import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+
+@SdkInternalApi
+public final class NestedRecordUtils {
+
+    private static final Pattern NESTED_OBJECT_PATTERN = Pattern.compile(NESTED_OBJECT_UPDATE);
+
+    private NestedRecordUtils() {
+    }
+
+    /**
+     * Resolves and returns the {@link TableSchema} for the element type of a list attribute from the provided root schema.
+     * <p>
+     * This method is useful when dealing with lists of nested objects in a DynamoDB-enhanced table schema, particularly in
+     * scenarios where the list is part of a flattened nested structure.
+     * <p>
+     * If the provided key contains the nested object delimiter (e.g., {@code _NESTED_ATTR_UPDATE_}), the method traverses the
+     * nested hierarchy based on that path to locate the correct schema for the target attribute. Otherwise, it directly resolves
+     * the list element type from the root schema using reflection.
+     *
+     * @param rootSchema The root {@link TableSchema} representing the top-level entity.
+     * @param key        The key representing the list attribute, either flat or nested (using a delimiter).
+     * @return The {@link TableSchema} representing the list element type of the specified attribute.
+     * @throws IllegalArgumentException If the list element class cannot be found via reflection.
+     */
+    public static TableSchema<?> getTableSchemaForListElement(TableSchema<?> rootSchema, String key) {
+        TableSchema<?> listElementSchema;
+        try {
+            if (!key.contains(NESTED_OBJECT_UPDATE)) {
+                Optional<? extends TableSchema<?>> staticSchema = getNestedSchema(rootSchema, key);
+                listElementSchema =
+                    staticSchema.isPresent()
+                    ? staticSchema.get()
+                    : 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;
+    }
+
+    /**
+     * Traverses the attribute keys representing flattened nested structures and resolves the corresponding {@link TableSchema}
+     * for each nested path.
+     * <p>
+     * The method constructs a mapping between each unique nested path (represented as dot-delimited strings) and the
+     * corresponding {@link TableSchema} object derived from the root schema. It supports resolving schemas for arbitrarily deep
+     * nesting, using the {@code _NESTED_ATTR_UPDATE_} pattern as a path delimiter.
+     * <p>
+     * This is typically used in update or transformation flows where fields from nested objects are represented as flattened keys
+     * in the attribute map (e.g., {@code parent_NESTED_ATTR_UPDATE_child}).
+     *
+     * @param attributesToSet A map of flattened attribute keys to values, where keys may represent paths to nested attributes.
+     * @param rootSchema      The root {@link TableSchema} of the top-level entity.
+     * @return A map where the key is the nested path (e.g., {@code "parent.child"}) and the value is the {@link TableSchema}
+     * corresponding to that level in the object hierarchy.
+     */
+    public static 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);
+
+            StringBuilder pathBuilder = new StringBuilder();
+            TableSchema<?> currentSchema = rootSchema;
+
+            for (int i = 0; i < parts.length - 1; i++) {
+                if (pathBuilder.length() > 0) {
+                    pathBuilder.append(".");
+                }
+                pathBuilder.append(parts[i]);
+
+                String path = pathBuilder.toString();
+
+                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;
+    }
+
+    public 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;
+    }
+}