Add support for strategy-based UUID auto-generation (@DynamoDbAutoGeneratedUuid)

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

@@ -0,0 +1,6 @@
+{
+  "type": "feature",
+  "category": "Amazon DynamoDB Enhanced Client",
+  "contributor": "",
+  "description": "Added support for DynamoDbAutoGeneratedKey annotation"
+}

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

@@ -0,0 +1,219 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.enhanced.dynamodb.extensions;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.UUID;
+import java.util.function.Consumer;
+import software.amazon.awssdk.annotations.SdkPublicApi;
+import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.enhanced.dynamodb.AttributeValueType;
+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.IndexMetadata;
+import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
+import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+import software.amazon.awssdk.utils.StringUtils;
+import software.amazon.awssdk.utils.Validate;
+
+/**
+ * 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>
+ *   <li>Partition key or sort key of any secondary index (GSI or LSI)</li>
+ * </ul>
+ *
+ * <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
+public final class AutoGeneratedKeyExtension implements DynamoDbEnhancedClientExtension {
+
+    /**
+     * Custom metadata key under which we store the set of annotated attribute names.
+     */
+    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() {
+    }
+
+    public static Builder builder() {
+        return new 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. Unlike {@code @DynamoDbAutoGeneratedUuid}, this preserves existing values.
+     * <p>
+     * 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) {
+        Collection<String> taggedAttributes = context.tableMetadata()
+                                                     .customMetadataObject(CUSTOM_METADATA_KEY, Collection.class)
+                                                     .orElse(null);
+
+        if (taggedAttributes == null) {
+            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<>();
+
+        // ensure every @DynamoDbAutoGeneratedKey attribute is a PK/SK or GSI/LSI. If not, throw IllegalArgumentException
+        allowedKeys.add(meta.primaryPartitionKey());
+        meta.primarySortKey().ifPresent(allowedKeys::add);
+
+        for (IndexMetadata idx : meta.indices()) {
+            String indexName = idx.name();
+            allowedKeys.add(meta.indexPartitionKey(indexName));
+            meta.indexSortKey(indexName).ifPresent(allowedKeys::add);
+        }
+
+        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());
+        taggedAttributes.forEach(attr -> insertUuidIfMissing(itemToTransform, attr));
+
+        return WriteModification.builder()
+                                .transformedItem(Collections.unmodifiableMap(itemToTransform))
+                                .build();
+    }
+
+    private void insertUuidIfMissing(Map<String, AttributeValue> itemToTransform, String key) {
+        AttributeValue existing = itemToTransform.get(key);
+        if (Objects.isNull(existing) || StringUtils.isBlank(existing.s())) {
+            itemToTransform.put(key, AttributeValue.builder().s(UUID.randomUUID().toString()).build());
+        }
+    }
+
+    /**
+     * Static helpers used by the {@code @BeanTableSchemaAttributeTag}-based annotation tag.
+     */
+    public static final class AttributeTags {
+        private AttributeTags() {
+        }
+
+        /**
+         * @return a {@link StaticAttributeTag} that marks the attribute for auto-generated key behavior.
+         */
+        public static StaticAttributeTag autoGeneratedKeyAttribute() {
+            return AUTO_GENERATED_KEY_ATTRIBUTE;
+        }
+    }
+
+    /**
+     * Stateless builder.
+     */
+    public static final class Builder {
+        private Builder() {
+        }
+
+        public AutoGeneratedKeyExtension build() {
+            return new AutoGeneratedKeyExtension();
+        }
+    }
+
+    /**
+     * Validates Java type and records the tagged attribute into table metadata so {@link #beforeWrite} can find it at runtime.
+     */
+    private static final class AutoGeneratedKeyAttribute implements StaticAttributeTag {
+
+        @Override
+        public <R> void validateType(String attributeName,
+                                     EnhancedType<R> type,
+                                     AttributeValueType attributeValueType) {
+
+            Validate.notNull(type, "type is null");
+            Validate.notNull(type.rawClass(), "rawClass is null");
+            Validate.notNull(attributeValueType, "attributeValueType is null");
+
+            if (!type.rawClass().equals(String.class)) {
+                throw new IllegalArgumentException(String.format(
+                    "Attribute '%s' of Class type %s is not a suitable Java Class type to be used as a Auto Generated "
+                    + "Key attribute. Only String Class type is supported.", attributeName, type.rawClass()));
+            }
+        }
+
+        @Override
+        public Consumer<StaticTableMetadata.Builder> modifyMetadata(String attributeName,
+                                                                    AttributeValueType attributeValueType) {
+            // Record the names of the attributes annotated with @DynamoDbAutoGeneratedKey for later lookup in beforeWrite()
+            return metadata -> metadata.addCustomMetadataObject(
+                CUSTOM_METADATA_KEY, Collections.singleton(attributeName));
+        }
+    }
+}

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() {
@@ -109,6 +124,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

@@ -0,0 +1,78 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.enhanced.dynamodb.extensions.annotations;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import software.amazon.awssdk.annotations.SdkPublicApi;
+import software.amazon.awssdk.enhanced.dynamodb.internal.extensions.AutoGeneratedKeyTag;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.UpdateBehavior;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.BeanTableSchemaAttributeTag;
+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.
+ *
+ * <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
+ * {@link software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension}.</p>
+ *
+ * <h3>Where this annotation can be applied</h3>
+ * This annotation is only valid on attributes that serve as keys:
+ * <ul>
+ *   <li>The table's primary partition key or sort key</li>
+ *   <li>The partition key or sort key of a secondary index (GSI or LSI)</li>
+ * </ul>
+ * If applied to any other attribute, the {@code AutoGeneratedKeyExtension} will throw an
+ * {@link IllegalArgumentException} at runtime.
+ *
+ * <h3>How values are generated</h3>
+ * <ul>
+ *   <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>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 immutable in DynamoDB and cannot use conditional update behaviors like
+ * {@link UpdateBehavior#WRITE_IF_NOT_EXISTS}. 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 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}.
+ */
+@SdkPublicApi
+@Documented
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.METHOD, ElementType.FIELD})
+@BeanTableSchemaAttributeTag(AutoGeneratedKeyTag.class)
+public @interface DynamoDbAutoGeneratedKey {
+}

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

@@ -0,0 +1,33 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.enhanced.dynamodb.internal.extensions;
+
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension;
+import software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbAutoGeneratedKey;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
+
+@SdkInternalApi
+public final class AutoGeneratedKeyTag {
+
+    private AutoGeneratedKeyTag() {
+    }
+
+    public static StaticAttributeTag attributeTagFor(DynamoDbAutoGeneratedKey annotation) {
+        return AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute();
+    }
+
+}