Added support for @DynamoDbAutoGeneratedTimestampAttribute and @DynamoDbUpdateBehavior on attributes within nested objects

Author: anasatirbasa

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

@@ -206,9 +206,17 @@ private Map<String, AttributeValue> processNestedObject(Map<String, AttributeVal
 
         nestedMap.forEach((nestedKey, nestedValue) -> {
             if (nestedValue.hasM()) {
-                updatedNestedMap.put(nestedKey,
-                                     AttributeValue.builder().m(processNestedObject(nestedValue.m(), nestedSchema,
-                                                                                    currentInstant)).build());
+                // Resolve the child-attribute schema; default to the current (parent) schema if none is registered
+                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()) {
                 try {

services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/AutoGeneratedTimestampRecordTest.java

@@ -72,6 +72,19 @@ public class AutoGeneratedTimestampRecordTest extends LocalDynamoDbSyncTestBase
                                                            .tags(autoGeneratedTimestampAttribute()))
                          .build();
 
+
+    private static final TableSchema<SubNestedRecord> SUB_NESTED_TABLE_SCHEMA =
+        StaticTableSchema.builder(SubNestedRecord.class)
+                         .newItemSupplier(SubNestedRecord::new)
+                         .addAttribute(Instant.class, a -> a.name("subNestedTimeAttribute")
+                                                            .getter(SubNestedRecord::getSubNestedTimeAttribute)
+                                                            .setter(SubNestedRecord::setSubNestedTimeAttribute)
+                                                            .tags(autoGeneratedTimestampAttribute()))
+                         .addAttribute(String.class, a -> a.name("subNestedAttribute")
+                                                           .getter(SubNestedRecord::getSubNestedAttribute)
+                                                           .setter(SubNestedRecord::setSubNestedAttribute))
+                         .build();
+
     private static final TableSchema<NestedRecord> NESTED_TABLE_SCHEMA =
         StaticTableSchema.builder(NestedRecord.class)
                          .newItemSupplier(NestedRecord::new)
@@ -82,6 +95,11 @@ public class AutoGeneratedTimestampRecordTest extends LocalDynamoDbSyncTestBase
                          .addAttribute(String.class, a -> a.name("nestedAttribute")
                                                            .getter(NestedRecord::getNestedAttribute)
                                                            .setter(NestedRecord::setNestedAttribute))
+                         .addAttribute(
+                             EnhancedType.documentOf(SubNestedRecord.class, SUB_NESTED_TABLE_SCHEMA),
+                             a -> a.name("subNestedRecord")
+                                   .getter(NestedRecord::getSubNestedRecord)
+                                   .setter(NestedRecord::setSubNestedRecord))
                          .build();
 
     private static final TableSchema<Record> TABLE_SCHEMA =
@@ -162,14 +180,28 @@ public void deleteTable() {
 
     @Test
     public void putNewRecordSetsInitialAutoGeneratedTimestampIncludingNestedFields() {
-        Record item = new Record().setId("id").setAttribute("one")
-                                  .setNestedRecord(new NestedRecord().setNestedAttribute("attribute"));
+        NestedRecord nestedRecord = new NestedRecord()
+            .setNestedAttribute("nestedAttribute")
+            .setSubNestedRecord(new SubNestedRecord().setSubNestedAttribute("subNestedAttribute"));
+
+        Record item = new Record()
+            .setId("id")
+            .setAttribute("one")
+            .setNestedRecord(nestedRecord);
+
         mappedTable.putItem(r -> r.item(item));
         Record result = mappedTable.getItem(r -> r.key(k -> k.partitionValue("id")));
         GetItemResponse itemAsStoredInDDB = getItemAsStoredFromDDB();
         FlattenedRecord flattenedRecord = new FlattenedRecord().setGenerated(MOCKED_INSTANT_NOW);
-        NestedRecord expectedNestedRecord = new NestedRecord().setNestedTimeAttribute(MOCKED_INSTANT_NOW)
-            .setNestedAttribute("attribute");
+
+        SubNestedRecord expectedSubNestedRecord = new SubNestedRecord()
+            .setSubNestedTimeAttribute(MOCKED_INSTANT_NOW)
+            .setSubNestedAttribute("subNestedAttribute");
+        NestedRecord expectedNestedRecord = new NestedRecord()
+            .setNestedTimeAttribute(MOCKED_INSTANT_NOW)
+            .setNestedAttribute("nestedAttribute")
+            .setSubNestedRecord(expectedSubNestedRecord);
+
         Record expectedRecord = new Record().setId("id")
                                             .setAttribute("one")
                                             .setLastUpdatedDate(MOCKED_INSTANT_NOW)
@@ -181,7 +213,8 @@ public void putNewRecordSetsInitialAutoGeneratedTimestampIncludingNestedFields()
         assertThat(result, is(expectedRecord));
         // The data in DDB is stored in converted time format
         assertThat(itemAsStoredInDDB.item().get("convertedLastUpdatedDate").s(), is("13 01 2019 14:00:00"));
-        assertThat(itemAsStoredInDDB.item().get("nestedRecord").m().get("nestedTimeAttribute").s(), is(MOCKED_INSTANT_NOW.toString()));
+        assertThat(itemAsStoredInDDB.item().get("nestedRecord").m().get("nestedTimeAttribute").s(),
+ is(MOCKED_INSTANT_NOW.toString()));
     }
 
     @Test
@@ -666,6 +699,7 @@ public String toString() {
     private static class NestedRecord {
         private Instant nestedTimeAttribute;
         private String nestedAttribute;
+        private SubNestedRecord subNestedRecord;
 
         public Instant getNestedTimeAttribute() {
             return nestedTimeAttribute;
@@ -676,13 +710,24 @@ public NestedRecord setNestedTimeAttribute(Instant nestedTimeAttribute) {
             return this;
         }
 
-        public String getNestedAttribute() {return nestedAttribute;}
+        public String getNestedAttribute() {
+            return nestedAttribute;
+        }
 
         public NestedRecord setNestedAttribute(String nestedAttribute) {
             this.nestedAttribute = nestedAttribute;
             return this;
         }
 
+        public SubNestedRecord getSubNestedRecord() {
+            return subNestedRecord;
+        }
+
+        public NestedRecord setSubNestedRecord(SubNestedRecord subNestedRecord) {
+            this.subNestedRecord = subNestedRecord;
+            return this;
+        }
+
         @Override
         public boolean equals(Object o) {
             if (this == o) {
@@ -693,19 +738,21 @@ public boolean equals(Object o) {
             }
             NestedRecord that = (NestedRecord) o;
             return Objects.equals(nestedTimeAttribute, that.nestedTimeAttribute) &&
-                   Objects.equals(nestedAttribute, that.nestedAttribute);
+                   Objects.equals(nestedAttribute, that.nestedAttribute) &&
+                   Objects.equals(subNestedRecord, that.subNestedRecord);
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(nestedTimeAttribute, nestedAttribute);
+            return Objects.hash(nestedTimeAttribute, nestedAttribute, subNestedRecord);
         }
 
         @Override
         public String toString() {
             return "NestedRecord{" +
                    "nestedTimeAttribute=" + nestedTimeAttribute +
-                   "nestedAttribute=" + nestedAttribute +
+                   ", nestedAttribute='" + nestedAttribute + '\'' +
+                   ", subNestedRecord=" + subNestedRecord +
                    '}';
         }
     }
@@ -761,5 +808,52 @@ public String toString() {
         }
     }
 
+    private static class SubNestedRecord {
+        private Instant subNestedTimeAttribute;
+        private String subNestedAttribute;
+
+        public Instant getSubNestedTimeAttribute() {
+            return subNestedTimeAttribute;
+        }
+
+        public SubNestedRecord setSubNestedTimeAttribute(Instant subNestedTimeAttribute) {
+            this.subNestedTimeAttribute = subNestedTimeAttribute;
+            return this;
+        }
+
+        public String getSubNestedAttribute() {
+            return subNestedAttribute;
+        }
+
+        public SubNestedRecord setSubNestedAttribute(String subNestedAttribute) {
+            this.subNestedAttribute = subNestedAttribute;
+            return this;
+        }
 
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+            SubNestedRecord that = (SubNestedRecord) o;
+            return Objects.equals(subNestedTimeAttribute, that.subNestedTimeAttribute) &&
+                   Objects.equals(subNestedAttribute, that.subNestedAttribute);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(subNestedTimeAttribute, subNestedAttribute);
+        }
+
+        @Override
+        public String toString() {
+            return "SubNestedRecord{" +
+                   "subNestedTimeAttribute=" + subNestedTimeAttribute +
+                   ", subNestedAttribute='" + subNestedAttribute + '\'' +
+                   '}';
+        }
+    }
 }