Fix handling of UpdateBehavior.WRITE_IF_NOT_EXISTS on @DynamoDbAutoGeneratedUuid

Author: anasatirbasa

.changes/next-release/feature-AmazonDynamoDBEnhancedClient-d5e2088.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "Amazon DynamoDB Enhanced Client",
+    "contributor": "",
+    "description": "Fix handling of UpdateBehavior.WRITE_IF_NOT_EXISTS on @DynamoDbAutoGeneratedUuid"
+}

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

@@ -29,51 +29,53 @@
 import software.amazon.awssdk.enhanced.dynamodb.EnhancedType;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
-import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbUpdateBehavior;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+import software.amazon.awssdk.utils.StringUtils;
 import software.amazon.awssdk.utils.Validate;
 
 
 /**
- * This extension facilitates the automatic generation of a unique UUID (Universally Unique Identifier) for a specified attribute
- * every time a new record is written to the database. The generated UUID is obtained using the
- * {@link java.util.UUID#randomUUID()} method.
- * <p>
- * This extension is not loaded by default when you instantiate a
- * {@link software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient}. Therefore, you need to specify it in a custom
- * extension when creating the enhanced client.
- * <p>
- * Example to add AutoGeneratedUuidExtension along with default extensions is
- * {@snippet :
- * DynamoDbEnhancedClient.builder().extensions(Stream.concat(ExtensionResolver.defaultExtensions().stream(),
- * Stream.of(AutoGeneratedUuidExtension.create())).collect(Collectors.toList())).build();
- *}
- * </p>
- * <p>
- * Example to just add AutoGeneratedUuidExtension without default extensions is
- * {@snippet :
- *             DynamoDbEnhancedClient.builder().extensions(AutoGeneratedUuidExtension.create()).build();
- *}
- * </p>
- * <p>
- * To utilize the auto-generated UUID feature, first, create a field in your model that will store the UUID for the attribute.
- * This class field must be of type {@link java.lang.String}, and you need to tag it as the autoGeneratedUuidAttribute. If you are
- * using the {@link software.amazon.awssdk.enhanced.dynamodb.mapper.BeanTableSchema}, then you should use the
- * {@link software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbAutoGeneratedUuid} annotation. If you are using
- * the {@link software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableSchema}, then you should use the
- * {@link
- * software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedUuidExtension.AttributeTags#autoGeneratedUuidAttribute()}
- * static attribute tag.
- * </p>
- * <p>
- * Every time a new record is successfully put into the database, the specified attribute will be automatically populated with a
- * unique UUID generated using {@link java.util.UUID#randomUUID()}. If the UUID needs to be created only for `putItem` and should
- * not be generated for an `updateItem`, then
- * {@link software.amazon.awssdk.enhanced.dynamodb.mapper.UpdateBehavior#WRITE_IF_NOT_EXISTS} must be along with
- * {@link DynamoDbUpdateBehavior}
+ * This extension facilitates the automatic generation of a unique UUID (Universally Unique Identifier) for attributes tagged with
+ * {@code @DynamoDbAutoGeneratedUuid} or {@link AutoGeneratedUuidExtension.AttributeTags#autoGeneratedUuidAttribute()}. The
+ * generated UUID is obtained using {@link java.util.UUID#randomUUID()}.
  *
- * </p>
+ * <p>Usage:</p>
+ * <ul>
+ *   <li>This extension is not loaded by default; register it explicitly when building
+ *       a {@link software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient}.</li>
+ *   <li>The annotated attribute must be of type {@link String}.</li>
+ *   <li>If the attribute is <b>missing or empty</b> in the item being written, a UUID
+ *       will be generated and set in the outgoing request map.</li>
+ *   <li>If the attribute already has a non-empty value, it is preserved and not overwritten.</li>
+ * </ul>
+ *
+ * <p><b>Update behavior:</b></p>
+ * <ul>
+ *   <li>With the default {@link software.amazon.awssdk.enhanced.dynamodb.mapper.UpdateBehavior#WRITE_ALWAYS},
+ *       a missing attribute on an update will cause a new UUID to be generated and written.</li>
+ *   <li>With {@link software.amazon.awssdk.enhanced.dynamodb.mapper.UpdateBehavior#WRITE_IF_NOT_EXISTS},
+ *       a missing attribute on an update will cause a UUID to be generated in the outgoing map,
+ *       but DynamoDB will preserve the existing stored value thanks to the conditional
+ *       <code>if_not_exists(...)</code> expression the mapper generates.</li>
+ * </ul>
+ *
+ * <p><b>Difference between putItem and updateItem:</b></p>
+ * <ul>
+ *   <li>{@code putItem} always replaces the entire item. If the field is absent in the payload, the extension
+ *       will generate a new UUID, even when {@code WRITE_IF_NOT_EXISTS} is specified.</li>
+ *   <li>{@code updateItem} respects {@code WRITE_IF_NOT_EXISTS}: if the attribute already exists in DynamoDB,
+ *       the previously stored value is preserved and a new UUID will not overwrite it.</li>
+ * </ul>
+ *
+ * <p>Examples:</p>
+ * <pre>{@code
+ * DynamoDbEnhancedClient client = DynamoDbEnhancedClient.builder()
+ *     .dynamoDbClient(lowLevelClient)
+ *     .extensions(AutoGeneratedUuidExtension.create())
+ *     .build();
+ * }</pre>
  */
+
 @SdkPublicApi
 @ThreadSafe
 public final class AutoGeneratedUuidExtension implements DynamoDbEnhancedClientExtension {
@@ -110,15 +112,18 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
         }
 
         Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
-        customMetadataObject.forEach(key -> insertUuidInItemToTransform(itemToTransform, key));
+        customMetadataObject.forEach(key -> insertUuidIfMissing(itemToTransform, key));
         return WriteModification.builder()
                                 .transformedItem(Collections.unmodifiableMap(itemToTransform))
                                 .build();
     }
 
-    private void insertUuidInItemToTransform(Map<String, AttributeValue> itemToTransform,
-                                             String key) {
-        itemToTransform.put(key, AttributeValue.builder().s(UUID.randomUUID().toString()).build());
+    private static void insertUuidIfMissing(Map<String, AttributeValue> item, String key) {
+        AttributeValue existing = item.get(key);
+        boolean missing = existing == null || StringUtils.isEmpty(existing.s());
+        if (missing) {
+            item.put(key, AttributeValue.builder().s(UUID.randomUUID().toString()).build());
+        }
     }
 
     public static final class AttributeTags {

services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/extensions/AutoGeneratedUuidExtensionTest.java

@@ -66,11 +66,11 @@ public class AutoGeneratedUuidExtensionTest {
                          .build();
 
     @Test
-    public void beforeWrite_updateItemOperation_hasUuidInItem_doesNotCreateUpdateExpressionAndFilters() {
+    public void beforeWrite_withExistingUuid_doesNotOverwriteValue() {
         ItemWithUuid SimpleItem = new ItemWithUuid();
         SimpleItem.setId(RECORD_ID);
-        String uuidAttribute = String.valueOf(UUID.randomUUID());
-        SimpleItem.setUuidAttribute(uuidAttribute);
+        String initialUuid = String.valueOf(UUID.randomUUID());
+        SimpleItem.setUuidAttribute(initialUuid);
 
         Map<String, AttributeValue> items = ITEM_WITH_UUID_MAPPER.itemToMap(SimpleItem, true);
         assertThat(items).hasSize(2);
@@ -85,13 +85,15 @@ public void beforeWrite_updateItemOperation_hasUuidInItem_doesNotCreateUpdateExp
         Map<String, AttributeValue> transformedItem = result.transformedItem();
         assertThat(transformedItem).isNotNull().hasSize(2);
         assertThat(transformedItem).containsEntry("id", AttributeValue.fromS(RECORD_ID));
-        isValidUuid(transformedItem.get("uuidAttribute").s());
         assertThat(result.updateExpression()).isNull();
 
+        String uuidValue = result.transformedItem().get("uuidAttribute").s();
+        isValidUuid(uuidValue);
+        assertThat(uuidValue).isEqualTo(initialUuid);
     }
 
     @Test
-    public void beforeWrite_updateItemOperation_hasNoUuidInItem_doesNotCreatesUpdateExpressionAndFilters() {
+    public void beforeWrite_withMissingUuid_generatesNewValue() {
         ItemWithUuid SimpleItem = new ItemWithUuid();
         SimpleItem.setId(RECORD_ID);
 
@@ -108,8 +110,31 @@ public void beforeWrite_updateItemOperation_hasNoUuidInItem_doesNotCreatesUpdate
         Map<String, AttributeValue> transformedItem = result.transformedItem();
         assertThat(transformedItem).isNotNull().hasSize(2);
         assertThat(transformedItem).containsEntry("id", AttributeValue.fromS(RECORD_ID));
-        isValidUuid(transformedItem.get("uuidAttribute").s());
         assertThat(result.updateExpression()).isNull();
+
+        isValidUuid(result.transformedItem().get("uuidAttribute").s());
+        assertThat(items).doesNotContainKey("uuidAttribute");
+    }
+
+    @Test
+    public void beforeWrite_withEmptyUuid_generatesNewValue() {
+        ItemWithUuid item = new ItemWithUuid();
+        item.setId(RECORD_ID);
+        item.setUuidAttribute(""); // empty should be treated as missing
+
+        Map<String, AttributeValue> items = ITEM_WITH_UUID_MAPPER.itemToMap(item, true);
+
+        WriteModification result = atomicCounterExtension.beforeWrite(
+            DefaultDynamoDbExtensionContext.builder()
+                                           .items(items)
+                                           .tableMetadata(ITEM_WITH_UUID_MAPPER.tableMetadata())
+                                           .operationName(OperationName.UPDATE_ITEM)
+                                           .operationContext(PRIMARY_CONTEXT)
+                                           .build());
+
+        assertThat(result.updateExpression()).isNull();
+        String uuidValue = result.transformedItem().get("uuidAttribute").s();
+        isValidUuid(uuidValue);
     }
 
     @Test