Added support for @DynamoDbUpdateBehavior on attributes within nested objects

Author: anasatirbasa

.changes/next-release/feature-AmazonDynamoDBEnhancedClient-3b3908c.json

@@ -0,0 +1,6 @@
+{
+  "type": "feature",
+  "category": "Amazon DynamoDB Enhanced Client",
+  "contributor": "",
+  "description": "Added support for @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,12 +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.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 +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,10 +73,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 +148,179 @@ 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 (value.hasM() && value.m() != null) {
+                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 (firstElement != null && firstElement.hasM()) {
+                    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 (listItem != null && listItem.hasM()) {
+                                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 -> 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 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) {
+                if (!updated) {
+                    updatedNestedMap = new HashMap<>(nestedMap);
+                    updated = true;
+                }
+                insertTimestampInItemToTransform(updatedNestedMap, String.valueOf(key),
+                                                 nestedSchema.converterForAttribute(key), 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 (firstElement != null && firstElement.hasM()) {
+                    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 (listItem != null && listItem.hasM()) {
+                                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 +365,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 +380,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/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();
+    }
 }