Added support for DynamoDbAutoGeneratedKey annotation

Author: anasatirbasa

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

@@ -42,6 +42,14 @@
  * Generates a random UUID (via {@link java.util.UUID#randomUUID()}) for any attribute tagged with
  * {@code @DynamoDbAutoGeneratedKey} when that attribute is missing or empty on a write (put/update).
  * <p>
+ * <b>Key Difference from @DynamoDbAutoGeneratedUuid:</b> This extension only generates UUIDs when the 
+ * attribute value is null or empty, preserving existing values. In contrast, {@code @DynamoDbAutoGeneratedUuid} 
+ * always generates new UUIDs regardless of existing values.
+ * <p>
+ * <b>Conflict Detection:</b> This extension cannot be used together with {@code @DynamoDbAutoGeneratedUuid} on the same
+ * attribute. If both annotations are applied to the same field, an {@link IllegalArgumentException} will be thrown
+ * at runtime to prevent unpredictable behavior based on extension load order.
+ * <p>
  * The annotation may be placed <b>only</b> on key attributes:
  * <ul>
  *   <li>Primary partition key (PK) or primary sort key (SK)</li>
@@ -51,6 +59,9 @@
  * <p><b>Validation:</b> The extension enforces this at runtime during {@link #beforeWrite} by comparing the
  * annotated attributes against the table's known key attributes. If an annotated attribute
  * is not a PK/SK or an GSI/LSI, an {@link IllegalArgumentException} is thrown.</p>
+ *
+ * <p><b>UpdateBehavior Limitations:</b> {@code @DynamoDbUpdateBehavior} has no effect on primary keys due to 
+ * DynamoDB's UpdateItem API requirements. It only affects secondary index keys.</p>
  */
 @SdkPublicApi
 @ThreadSafe
@@ -62,6 +73,12 @@ public final class AutoGeneratedKeyExtension implements DynamoDbEnhancedClientEx
     private static final String CUSTOM_METADATA_KEY =
         "software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension:AutoGeneratedKeyAttribute";
 
+    /**
+     * Metadata key used by AutoGeneratedUuidExtension to detect conflicts.
+     */
+    private static final String UUID_EXTENSION_METADATA_KEY =
+        "software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedUuidExtension:AutoGeneratedUuidAttribute";
+
     private static final AutoGeneratedKeyAttribute AUTO_GENERATED_KEY_ATTRIBUTE = new AutoGeneratedKeyAttribute();
 
     private AutoGeneratedKeyExtension() {
@@ -73,9 +90,10 @@ public static Builder builder() {
 
     /**
      * If this table has attributes tagged for auto-generation, insert a UUID value into the outgoing item for any such attribute
-     * that is currently missing/empty.
+     * that is currently missing/empty. Unlike {@code @DynamoDbAutoGeneratedUuid}, this preserves existing values.
      * <p>
-     * Also validates that the annotation is only used on PK/SK/GSI/LSI key attributes.
+     * Also validates that the annotation is only used on PK/SK/GSI/LSI key attributes and that there are no conflicts
+     * with @DynamoDbAutoGeneratedUuid.
      */
     @Override
     public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
@@ -87,6 +105,21 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
             return WriteModification.builder().build();
         }
 
+        // Check for conflicts with @DynamoDbAutoGeneratedUuid
+        Collection<String> uuidTaggedAttributes = context.tableMetadata()
+                                                         .customMetadataObject(UUID_EXTENSION_METADATA_KEY, Collection.class)
+                                                         .orElse(Collections.emptyList());
+
+        taggedAttributes.stream()
+                        .filter(uuidTaggedAttributes::contains)
+                        .findFirst()
+                        .ifPresent(attribute -> {
+                            throw new IllegalArgumentException(
+                                "Attribute '" + attribute + "' cannot have both @DynamoDbAutoGeneratedKey and "
+                                + "@DynamoDbAutoGeneratedUuid annotations. These annotations have conflicting behaviors "
+                                + "and cannot be used together on the same attribute.");
+                        });
+
         TableMetadata meta = context.tableMetadata();
         Set<String> allowedKeys = new HashSet<>();
 
@@ -100,15 +133,15 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
             meta.indexSortKey(indexName).ifPresent(allowedKeys::add);
         }
 
-        for (String attr : taggedAttributes) {
-            if (!allowedKeys.contains(attr)) {
-                throw new IllegalArgumentException(
-                    "@DynamoDbAutoGeneratedKey can only be applied to key attributes: " +
-                    "primary partition key, primary sort key, or GSI/LSI partition/sort keys." +
-                    "Invalid placement on attribute: " + attr);
-
-            }
-        }
+        taggedAttributes.stream()
+                        .filter(attr -> !allowedKeys.contains(attr))
+                        .findFirst()
+                        .ifPresent(attr -> {
+                            throw new IllegalArgumentException(
+                                "@DynamoDbAutoGeneratedKey can only be applied to key attributes: "
+                                + "primary partition key, primary sort key, or GSI/LSI partition/sort keys."
+                                + "Invalid placement on attribute: " + attr);
+                        });
 
         // Generate UUIDs for missing/empty annotated attributes
         Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());

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

@@ -39,6 +39,14 @@
  * every time a new record is written to the database. The generated UUID is obtained using the
  * {@link java.util.UUID#randomUUID()} method.
  * <p>
+ * <b>Key Difference from @DynamoDbAutoGeneratedKey:</b> This extension always generates new UUIDs on every write,
+ * regardless of existing values. In contrast, {@code @DynamoDbAutoGeneratedKey} only generates UUIDs when the 
+ * attribute value is null or empty, preserving existing values.
+ * <p>
+ * <b>Conflict Detection:</b> This extension cannot be used together with {@code @DynamoDbAutoGeneratedKey} on the same
+ * attribute. If both annotations are applied to the same field, an {@link IllegalArgumentException} will be thrown
+ * at runtime to prevent unpredictable behavior.
+ * <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.
@@ -79,6 +87,13 @@
 public final class AutoGeneratedUuidExtension implements DynamoDbEnhancedClientExtension {
     private static final String CUSTOM_METADATA_KEY =
         "software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedUuidExtension:AutoGeneratedUuidAttribute";
+    
+    /**
+     * Metadata key used by AutoGeneratedKeyExtension to detect conflicts.
+     */
+    private static final String KEY_EXTENSION_METADATA_KEY =
+        "software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension:AutoGeneratedKeyAttribute";
+    
     private static final AutoGeneratedUuidAttribute AUTO_GENERATED_UUID_ATTRIBUTE = new AutoGeneratedUuidAttribute();
 
     private AutoGeneratedUuidExtension() {
@@ -99,8 +114,6 @@ public static AutoGeneratedUuidExtension create() {
      */
     @Override
     public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
-
-
         Collection<String> customMetadataObject = context.tableMetadata()
                                                          .customMetadataObject(CUSTOM_METADATA_KEY, Collection.class)
                                                          .orElse(null);
@@ -109,6 +122,21 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
             return WriteModification.builder().build();
         }
 
+        // Check for conflicts with @DynamoDbAutoGeneratedKey
+        Collection<String> keyTaggedAttributes = context.tableMetadata()
+                                                        .customMetadataObject(KEY_EXTENSION_METADATA_KEY, Collection.class)
+                                                        .orElse(Collections.emptyList());
+
+        customMetadataObject.stream()
+                            .filter(keyTaggedAttributes::contains)
+                            .findFirst()
+                            .ifPresent(attribute -> {
+                                throw new IllegalArgumentException(
+                                    "Attribute '" + attribute + "' cannot have both @DynamoDbAutoGeneratedKey and "
+                                    + "@DynamoDbAutoGeneratedUuid annotations. These annotations have conflicting behaviors "
+                                    + "and cannot be used together on the same attribute.");
+                            });
+
         Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
         customMetadataObject.forEach(key -> insertUuidInItemToTransform(itemToTransform, key));
         return WriteModification.builder()

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

@@ -27,8 +27,8 @@
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbUpdateBehavior;
 
 /**
- * Annotation that marks a key attribute to be automatically populated with a random UUID if no value is provided during a
- * write operation (put or update). This annotation is intended to work specifically with key attributes.
+ * Annotation that marks a key attribute to be automatically populated with a random UUID if no value is provided during a write
+ * operation (put or update). This annotation is intended to work specifically with key attributes.
  *
  * <p>This annotation is designed for use with the V2 {@link software.amazon.awssdk.enhanced.dynamodb.mapper.BeanTableSchema}.
  * It is registered via {@link BeanTableSchemaAttributeTag} and its behavior is implemented by
@@ -48,19 +48,22 @@
  *   <li>On writes where the annotated attribute is null or empty, a new UUID value is generated
  *       using {@link java.util.UUID#randomUUID()}.</li>
  *   <li>If a value is already set on the attribute, that value is preserved and not replaced.</li>
+ *   <li>This behavior differs from {@code @DynamoDbAutoGeneratedUuid}, which always generates new UUIDs regardless of existing
+ *   values.</li>
  * </ul>
  *
- * <h3>Controlling regeneration on update</h3>
- * This annotation can be combined with {@link DynamoDbUpdateBehavior} to control regeneration behavior:
+ * <h3>Behavior with UpdateBehavior</h3>
+ * <p><strong>Primary Keys:</strong> {@link DynamoDbUpdateBehavior} has <strong>no effect</strong> on primary partition keys
+ * or primary sort keys. Primary keys are required for UpdateItem operations in DynamoDB and cannot be conditionally
+ * updated. UUIDs will be generated whenever the primary key attribute is missing or empty, regardless of any
+ * {@code UpdateBehavior} setting.</p>
+ *
+ * <p><strong>Secondary Index Keys:</strong> For GSI/LSI keys, {@link DynamoDbUpdateBehavior} can be used:
  * <ul>
- *   <li>{@link UpdateBehavior#WRITE_ALWAYS} (default) –
- *       Generate a new UUID whenever the attribute is missing during write.</li>
- *   <li>{@link UpdateBehavior#WRITE_IF_NOT_EXISTS} –
- *       Generate a UUID only the first time (on insert), and preserve that value on subsequent updates.</li>
- * </ul>
- * <p><strong>Important:</strong> {@link DynamoDbUpdateBehavior} only affects secondary index keys.
- * Primary keys cannot be null in DynamoDB and are required for UpdateItem operations, so UpdateBehavior
- * has no practical effect on primary keys.</p>
+ *   <li>{@link UpdateBehavior#WRITE_ALWAYS} (default) – Generate a new UUID whenever the attribute is missing during write.</li>
+ *   <li>{@link UpdateBehavior#WRITE_IF_NOT_EXISTS} – Generate a UUID only on the first write, preserving the value on
+ *   subsequent updates.</li>
+ * </ul></p>
  *
  * <h3>Type restriction</h3>
  * This annotation is only valid on attributes of type {@link String}.

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

@@ -251,6 +251,121 @@ public void autoGeneratedKey_onNonKey_throwsIllegalArgumentException() {
             .withMessageContaining("keyAttribute");
     }
 
+    @Test
+    public void conflictingAnnotations_throwsIllegalArgumentException() {
+        // Create a schema with both @DynamoDbAutoGeneratedKey and @DynamoDbAutoGeneratedUuid on the same attribute
+        StaticTableSchema<ItemWithKey> conflictingSchema =
+            StaticTableSchema
+                .builder(ItemWithKey.class)
+                .newItemSupplier(ItemWithKey::new)
+                .addAttribute(String.class, a -> a.name("id")
+                                                  .getter(ItemWithKey::getId)
+                                                  .setter(ItemWithKey::setId)
+                                                  .addTag(primaryPartitionKey())
+                                                  // Both annotations on the same attribute
+                                                  .addTag(AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())
+                                                  .addTag(AutoGeneratedUuidExtension.AttributeTags.autoGeneratedUuidAttribute()))
+                .addAttribute(String.class, a -> a.name("simpleString")
+                                                  .getter(ItemWithKey::getSimpleString)
+                                                  .setter(ItemWithKey::setSimpleString))
+                .build();
+
+        ItemWithKey item = new ItemWithKey();
+        item.setId(RECORD_ID);
+
+        Map<String, AttributeValue> items = conflictingSchema.itemToMap(item, true);
+
+        assertThatExceptionOfType(IllegalArgumentException.class)
+            .isThrownBy(() ->
+                            extension.beforeWrite(DefaultDynamoDbExtensionContext.builder()
+                                                                                 .items(items)
+                                                                                 .tableMetadata(conflictingSchema.tableMetadata())
+                                                                                 .operationName(OperationName.PUT_ITEM)
+                                                                                 .operationContext(PRIMARY_CONTEXT)
+                                                                                 .build())
+            )
+            .withMessage("Attribute 'id' cannot have both @DynamoDbAutoGeneratedKey and @DynamoDbAutoGeneratedUuid annotations. "
+                         + "These annotations have conflicting behaviors and cannot be used together on the same attribute.");
+    }
+
+    @Test
+    public void conflictingAnnotations_onSecondaryKey_throwsIllegalArgumentException() {
+        // Create a schema with both annotations on a GSI key
+        StaticTableSchema<ItemWithKey> conflictingGsiSchema =
+            StaticTableSchema
+                .builder(ItemWithKey.class)
+                .newItemSupplier(ItemWithKey::new)
+                .addAttribute(String.class, a -> a.name("id")
+                                                  .getter(ItemWithKey::getId)
+                                                  .setter(ItemWithKey::setId)
+                                                  .addTag(primaryPartitionKey()))
+                .addAttribute(String.class, a -> a.name("keyAttribute")
+                                                  .getter(ItemWithKey::getKeyAttribute)
+                                                  .setter(ItemWithKey::setKeyAttribute)
+                                                  .addTag(secondaryPartitionKey("gsi1"))
+                                                  // Both annotations on the same GSI key
+                                                  .addTag(AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())
+                                                  .addTag(AutoGeneratedUuidExtension.AttributeTags.autoGeneratedUuidAttribute()))
+                .addAttribute(String.class, a -> a.name("simpleString")
+                                                  .getter(ItemWithKey::getSimpleString)
+                                                  .setter(ItemWithKey::setSimpleString))
+                .build();
+
+        ItemWithKey item = new ItemWithKey();
+        item.setId(RECORD_ID);
+
+        Map<String, AttributeValue> items = conflictingGsiSchema.itemToMap(item, true);
+
+        assertThatExceptionOfType(IllegalArgumentException.class)
+            .isThrownBy(() ->
+                            extension.beforeWrite(
+                                DefaultDynamoDbExtensionContext.builder()
+                                                               .items(items)
+                                                               .tableMetadata(conflictingGsiSchema.tableMetadata())
+                                                               .operationName(OperationName.PUT_ITEM)
+                                                               .operationContext(PRIMARY_CONTEXT)
+                                                               .build()))
+            .withMessage("Attribute 'keyAttribute' cannot have both @DynamoDbAutoGeneratedKey and @DynamoDbAutoGeneratedUuid "
+                         + "annotations. "
+                         + "These annotations have conflicting behaviors and cannot be used together on the same attribute.");
+    }
+
+    @Test
+    public void conflictDetection_worksRegardlessOfExtensionOrder() {
+        // Verify that AutoGeneratedKeyExtension detects conflicts even when 
+        // AutoGeneratedUuidExtension metadata is already present
+        StaticTableSchema<ItemWithKey> conflictingSchema =
+            StaticTableSchema
+                .builder(ItemWithKey.class)
+                .newItemSupplier(ItemWithKey::new)
+                .addAttribute(String.class, a -> a.name("id")
+                                                  .getter(ItemWithKey::getId)
+                                                  .setter(ItemWithKey::setId)
+                                                  .addTag(primaryPartitionKey())
+                                                  .addTag(AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())
+                                                  .addTag(AutoGeneratedUuidExtension.AttributeTags.autoGeneratedUuidAttribute()))
+                .build();
+
+        ItemWithKey item = new ItemWithKey();
+        item.setId(RECORD_ID);
+
+        Map<String, AttributeValue> items = conflictingSchema.itemToMap(item, true);
+
+        // Test that the conflict is detected regardless of which extension runs first
+        assertThatExceptionOfType(IllegalArgumentException.class)
+            .isThrownBy(() ->
+                            extension.beforeWrite(
+                                DefaultDynamoDbExtensionContext.builder()
+                                                               .items(items)
+                                                               .tableMetadata(conflictingSchema.tableMetadata())
+                                                               .operationName(OperationName.PUT_ITEM)
+                                                               .operationContext(PRIMARY_CONTEXT)
+                                                               .build()))
+            .withMessage("Attribute 'id' cannot have both @DynamoDbAutoGeneratedKey and @DynamoDbAutoGeneratedUuid "
+                         + "annotations. "
+                         + "These annotations have conflicting behaviors and cannot be used together on the same attribute.");
+    }
+
     private static class ItemWithKey {
 
         private String id;