Support AutoGeneratedTimestamp annotations 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 on attributes within nested objects."
+}

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

@@ -15,12 +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.EnhancedClientUtils.hasMap;
+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.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.function.Consumer;
 import software.amazon.awssdk.annotations.NotThreadSafe;
 import software.amazon.awssdk.annotations.SdkPublicApi;
@@ -30,6 +39,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,10 +74,23 @@
  * <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.
+ *
+ * <p>
+ * <b>Note:</b> This implementation uses a temporary cache keyed by {@link TableSchema} instance.
+ * When updating timestamps in nested objects or lists, the correct {@code TableSchema} must be used for each object.
+ * This cache ensures that each nested object is processed with its own schema, avoiding redundant lookups and ensuring
+ * all annotated timestamp fields are updated correctly.
+ * </p>
  */
 @SdkPublicApi
 @ThreadSafe
 public final class AutoGeneratedTimestampRecordExtension implements DynamoDbEnhancedClientExtension {
+
+    private static final String NESTED_OBJECT_UPDATE = "_NESTED_ATTR_UPDATE_";
     private static final String CUSTOM_METADATA_KEY = "AutoGeneratedTimestampExtension:AutoGeneratedTimestampAttribute";
     private static final AutoGeneratedTimestampAttribute
         AUTO_GENERATED_TIMESTAMP_ATTRIBUTE = new AutoGeneratedTimestampAttribute();
@@ -126,26 +149,187 @@ 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);
+        // Use TableSchema<?> instance as the cache key
+        Map<TableSchema<?>, TableSchema<?>> schemaInstanceCache = new HashMap<>();
 
-        if (customMetadataObject == null) {
+        itemToTransform.forEach((key, value) -> {
+            if (hasMap(value)) {
+                Optional<? extends TableSchema<?>> nestedSchemaOpt = getNestedSchema(context.tableSchema(), key);
+                if (nestedSchemaOpt.isPresent()) {
+                    TableSchema<?> nestedSchema = nestedSchemaOpt.get();
+                    TableSchema<?> cachedSchema = getOrCacheSchema(schemaInstanceCache, nestedSchema);
+                    Map<String, AttributeValue> processed =
+                        processNestedObject(value.m(), cachedSchema, currentInstant, schemaInstanceCache);
+                    updatedItems.put(key, AttributeValue.builder().m(processed).build());
+                }
+            } else if (value.hasL() && !value.l().isEmpty()) {
+                // Check first non-null element to determine if this is a list of maps
+                AttributeValue firstElement = value.l().stream()
+                                                   .filter(Objects::nonNull)
+                                                   .findFirst()
+                                                   .orElse(null);
+
+                if (hasMap(firstElement)) {
+                    TableSchema<?> elementListSchema = getTableSchemaForListElement(context.tableSchema(), key);
+                    if (elementListSchema != null) {
+                        TableSchema<?> cachedSchema = getOrCacheSchema(schemaInstanceCache, elementListSchema);
+                        Collection<AttributeValue> updatedList = new ArrayList<>(value.l().size());
+                        for (AttributeValue listItem : value.l()) {
+                            if (hasMap(listItem)) {
+                                updatedList.add(AttributeValue.builder()
+                                                              .m(processNestedObject(
+                                                                  listItem.m(),
+                                                                  cachedSchema,
+                                                                  currentInstant,
+                                                                  schemaInstanceCache))
+                                                              .build());
+                            } else {
+                                updatedList.add(listItem);
+                            }
+                        }
+                        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 -> {
+                        AttributeConverter<?> converter = schema.converterForAttribute(key);
+                        if (converter != null) {
+                            insertTimestampInItemToTransform(updatedItems, reconstructCompositeKey(path, key),
+                                                             converter, 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 static TableSchema<?> getOrCacheSchema(
+        Map<TableSchema<?>, TableSchema<?>> schemaInstanceCache, TableSchema<?> schema) {
+
+        TableSchema<?> cachedSchema = schemaInstanceCache.get(schema);
+        if (cachedSchema == null) {
+            schemaInstanceCache.put(schema, schema);
+            cachedSchema = schema;
+        }
+        return cachedSchema;
+    }
+
+    private Map<String, AttributeValue> processNestedObject(Map<String, AttributeValue> nestedMap, TableSchema<?> nestedSchema,
+                                                            Instant currentInstant,
+                                                            Map<TableSchema<?>, TableSchema<?>> schemaInstanceCache) {
+        Map<String, AttributeValue> updatedNestedMap = nestedMap;
+        boolean updated = false;
+
+        Collection<String> customMetadataObject = nestedSchema.tableMetadata()
+                                                              .customMetadataObject(CUSTOM_METADATA_KEY, Collection.class)
+                                                              .orElse(null);
+
+        if (customMetadataObject != null) {
+            for (String key : customMetadataObject) {
+                AttributeConverter<?> converter = nestedSchema.converterForAttribute(key);
+                if (converter != null) {
+                    if (!updated) {
+                        updatedNestedMap = new HashMap<>(nestedMap);
+                        updated = true;
+                    }
+                    insertTimestampInItemToTransform(updatedNestedMap, String.valueOf(key),
+                                                     converter, currentInstant);
+                }
+            }
+        }
+
+        for (Map.Entry<String, AttributeValue> entry : nestedMap.entrySet()) {
+            String nestedKey = entry.getKey();
+            AttributeValue nestedValue = entry.getValue();
+            if (nestedValue.hasM()) {
+                Optional<? extends TableSchema<?>> childSchemaOpt = getNestedSchema(nestedSchema, nestedKey);
+                if (childSchemaOpt.isPresent()) {
+                    TableSchema<?> childSchema = childSchemaOpt.get();
+                    TableSchema<?> cachedSchema = getOrCacheSchema(schemaInstanceCache, childSchema);
+                    Map<String, AttributeValue> processed = processNestedObject(
+                        nestedValue.m(), cachedSchema, currentInstant, schemaInstanceCache);
+
+                    if (!Objects.equals(processed, nestedValue.m())) {
+                        if (!updated) {
+                            updatedNestedMap = new HashMap<>(nestedMap);
+                            updated = true;
+                        }
+                        updatedNestedMap.put(nestedKey, AttributeValue.builder().m(processed).build());
+                    }
+                }
+            } else if (nestedValue.hasL() && !nestedValue.l().isEmpty()) {
+                // Check first non-null element to determine if this is a list of maps
+                AttributeValue firstElement = nestedValue.l().stream()
+                                                   .filter(Objects::nonNull)
+                                                   .findFirst()
+                                                   .orElse(null);
+                if (hasMap(firstElement)) {
+                    TableSchema<?> listElementSchema = getTableSchemaForListElement(nestedSchema, nestedKey);
+                    if (listElementSchema != null) {
+                        TableSchema<?> cachedSchema = getOrCacheSchema(schemaInstanceCache, listElementSchema);
+                        Collection<AttributeValue> updatedList = new ArrayList<>(nestedValue.l().size());
+                        boolean listModified = false;
+                        for (AttributeValue listItem : nestedValue.l()) {
+                            if (hasMap(listItem)) {
+                                AttributeValue updatedItem = AttributeValue.builder()
+                                                                           .m(processNestedObject(
+                                                                               listItem.m(),
+                                                                               cachedSchema,
+                                                                               currentInstant,
+                                                                               schemaInstanceCache))
+                                                                           .build();
+                                updatedList.add(updatedItem);
+                                if (!Objects.equals(updatedItem, listItem)) {
+                                    listModified = true;
+                                }
+                            } else {
+                                updatedList.add(listItem);
+                            }
+                        }
+                        if (listModified) {
+                            if (!updated) {
+                                updatedNestedMap = new HashMap<>(nestedMap);
+                                updated = true;
+                            }
+                            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));
     }
 
     /**
@@ -190,6 +374,7 @@ public <R> void validateType(String attributeName, EnhancedType<R> type,
             Validate.notNull(type, "type is null");
             Validate.notNull(type.rawClass(), "rawClass is null");
             Validate.notNull(attributeValueType, "attributeValueType is null");
+            validateAttributeName(attributeName);
 
             if (!type.rawClass().equals(Instant.class)) {
                 throw new IllegalArgumentException(String.format(
@@ -204,5 +389,15 @@ public Consumer<StaticTableMetadata.Builder> modifyMetadata(String attributeName
             return metadata -> metadata.addCustomMetadataObject(CUSTOM_METADATA_KEY, Collections.singleton(attributeName))
                                        .markAttributeAsKey(attributeName, attributeValueType);
         }
+
+        private static void validateAttributeName(String attributeName) {
+            if (attributeName.contains(NESTED_OBJECT_UPDATE)) {
+                throw new IllegalArgumentException(
+                    String.format(
+                        "Attribute name '%s' contains reserved marker '%s' and is not allowed.",
+                        attributeName,
+                        NESTED_OBJECT_UPDATE));
+            }
+        }
     }
 }

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

@@ -27,6 +27,10 @@
  * Denotes this attribute as recording the auto generated last updated timestamp for the record.
  * Every time a record with this attribute is written to the database it will update the attribute with current timestamp when
  * its updated.
+ * <p>
+ * Note: This annotation must not be applied to fields whose names contain the reserved marker "_NESTED_ATTR_UPDATE_".
+ * This marker is used internally by the Enhanced Client to represent flattened paths for nested attribute updates.
+ * If a field name contains this marker, an IllegalArgumentException will be thrown during schema registration.
  */
 @SdkPublicApi
 @Target({ElementType.METHOD})

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

@@ -28,7 +28,9 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.enhanced.dynamodb.AttributeConverter;
 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;
@@ -37,6 +39,8 @@
 import software.amazon.awssdk.enhanced.dynamodb.model.Page;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.ConsumedCapacity;
+import software.amazon.awssdk.utils.CollectionUtils;
+import software.amazon.awssdk.utils.StringUtils;
 
 @SdkInternalApi
 public final class EnhancedClientUtils {
@@ -204,4 +208,42 @@ public static <T> List<T> getItemsFromSupplier(List<Supplier<T>> itemSupplierLis
     public static boolean isNullAttributeValue(AttributeValue attributeValue) {
         return attributeValue.nul() != null && attributeValue.nul();
     }
+
+    public static boolean hasMap(AttributeValue attributeValue) {
+        return attributeValue != null && attributeValue.hasM();
+    }
+
+    /**
+     * Retrieves the nested {@link TableSchema} for an attribute from the parent schema.
+     * For parameterized types (e.g., Set, List, Map), extracts the first type parameter's schema.
+     *
+     * @param parentSchema the parent schema; must not be null
+     * @param attributeName the attribute name; must not be null or empty
+     * @return the nested schema, or empty if unavailable
+     */
+    public static Optional<? extends TableSchema<?>> getNestedSchema(TableSchema<?> parentSchema, CharSequence attributeName) {
+        if (parentSchema == null) {
+            throw new IllegalArgumentException("Parent schema cannot be null.");
+        }
+        if (StringUtils.isEmpty(attributeName)) {
+            throw new IllegalArgumentException("Attribute name cannot be null or empty.");
+        }
+
+        AttributeConverter<?> converter = parentSchema.converterForAttribute(attributeName);
+        if (converter == null) {
+            return Optional.empty();
+        }
+
+        EnhancedType<?> enhancedType = converter.type();
+        if (enhancedType == null) {
+            return Optional.empty();
+        }
+
+        List<EnhancedType<?>> rawClassParameters = enhancedType.rawClassParameters();
+        if (!CollectionUtils.isNullOrEmpty(rawClassParameters)) {
+            enhancedType = rawClassParameters.get(0);
+        }
+
+        return enhancedType.tableSchema();
+    }
 }