Added support for DynamoDbAutoGeneratedKey annotation

.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,254 @@
+/*
+ * 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 static java.util.Collections.newSetFromMap;
+
+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.concurrent.ConcurrentHashMap;
+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 attributes tagged with {@code @DynamoDbAutoGeneratedKey}
+ * when the attribute value is missing or empty during write operations (put, update, batch write, or transact write).
+ *
+ * <p><b>Difference from {@code @DynamoDbAutoGeneratedUuid}:</b>
+ * This extension generates a UUID only when the attribute is null or empty, preserving existing values. In contrast,
+ * {@code @DynamoDbAutoGeneratedUuid} always generates a new UUID regardless of the current value.
+ *
+ * <p><b>Conflict Detection:</b>
+ * {@code @DynamoDbAutoGeneratedKey} and {@code @DynamoDbAutoGeneratedUuid} cannot be applied to the same attribute. If both
+ * annotations are present, an {@link IllegalArgumentException} is thrown during schema validation.
+ *
+ * <p><b>Supported Attributes:</b>
+ * The annotation may only be applied to key attributes:
+ * <ul>
+ *   <li>Primary partition key (PK)</li>
+ *   <li>Primary sort key (SK)</li>
+ *   <li>Partition or sort keys of secondary indexes (GSI or LSI)</li>
+ * </ul>
+ *
+ * <p><b>Validation Behavior:</b>
+ * Annotation conflict detection and key-placement validation are performed once per {@link TableMetadata}
+ * instance and cached to avoid repeated validation on subsequent writes.
+ *
+ * <p><b>UpdateBehavior Limitation:</b>
+ * {@code @DynamoDbUpdateBehavior} does not apply to primary keys due to DynamoDB API constraints.
+ * It only affects secondary index keys.
+ */
+@SdkPublicApi
+@ThreadSafe
+public final class AutoGeneratedKeyExtension implements DynamoDbEnhancedClientExtension {
+
+    /**
+     * Metadata keys used to record attributes annotated with {@code @DynamoDbAutoGeneratedKey} and
+     * {@code @DynamoDbAutoGeneratedUuid}. These are used during schema validation to detect annotation conflicts.
+     */
+    private static final String CUSTOM_METADATA_KEY =
+        "software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension:AutoGeneratedKeyAttribute";
+    private static final String UUID_EXTENSION_METADATA_KEY =
+        "software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedUuidExtension:AutoGeneratedUuidAttribute";
+
+    private static final StaticAttributeTag AUTO_GENERATED_KEY_ATTRIBUTE = new AutoGeneratedKeyAttribute();
+    private static final String AUTOGENERATED_KEY_ANNOTATION = "@DynamoDbAutoGeneratedKey";
+    private static final String AUTOGENERATED_UUID_ANNOTATION = "@DynamoDbAutoGeneratedUuid";
+
+    /**
+     * Stores the TableMetadata instances that have already been validated by this extension. Uses a ConcurrentHashMap to ensure
+     * thread-safe access during concurrent write operations.
+     */
+    private final Set<TableMetadata> validatedSchemas = newSetFromMap(new ConcurrentHashMap<>());
+
+    /**
+     * Caches the set of valid key attribute names per TableMetadata instance. Computed once per schema.
+     */
+    private final Map<TableMetadata, Set<String>> allowedKeysCache = new ConcurrentHashMap<>();
+
+    private AutoGeneratedKeyExtension() {
+    }
+
+    /**
+     * @return an Instance of {@link AutoGeneratedKeyExtension}
+     */
+    public static AutoGeneratedKeyExtension create() {
+        return new AutoGeneratedKeyExtension();
+    }
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /**
+     * Inserts a UUID for attributes tagged with {@code @DynamoDbAutoGeneratedKey} when the attribute is null or empty, preserving
+     * existing values.
+     *
+     * <p>Schema-level validation (annotation conflict detection and key-placement checks)
+     * is executed once per schema instance and cached.
+     */
+    @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();
+        }
+
+        TableMetadata tableMetadata = context.tableMetadata();
+        if (validatedSchemas.add(tableMetadata)) {
+            validateNoAutoGeneratedAnnotationConflict(tableMetadata);
+            validateAutoGeneratedKeyPlacement(tableMetadata, taggedAttributes);
+        }
+
+        Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
+        taggedAttributes.forEach(attr -> insertUuidIfMissing(itemToTransform, attr));
+
+        return WriteModification.builder()
+                                .transformedItem(Collections.unmodifiableMap(itemToTransform))
+                                .build();
+    }
+
+    /**
+     * Validates (once per TableMetadata instance) that @DynamoDbAutoGeneratedKey and @DynamoDbAutoGeneratedUuid are not applied
+     * to the same attribute.
+     */
+    private static void validateNoAutoGeneratedAnnotationConflict(TableMetadata tableMetadata) {
+        ExtensionsValidationUtils.validateNoAnnotationConflict(
+            tableMetadata,
+            CUSTOM_METADATA_KEY,
+            UUID_EXTENSION_METADATA_KEY,
+            AUTOGENERATED_KEY_ANNOTATION,
+            AUTOGENERATED_UUID_ANNOTATION
+        );
+    }
+
+    /**
+     * Validates that all attributes tagged with @DynamoDbAutoGeneratedKey are either primary keys or secondary index keys.
+     **/
+    private void validateAutoGeneratedKeyPlacement(TableMetadata tableMetadata,
+                                                   Collection<String> taggedAttributeNames) {
+
+        Set<String> allowedKeyAttributes = allowedKeysCache.computeIfAbsent(tableMetadata, metadata -> {
+
+            // Add primary keys
+            Set<String> keyAttributes = new HashSet<>(metadata.primaryKeys());
+
+            // Add all secondary index keys
+            metadata.indices().stream().map(IndexMetadata::name).forEach(indexName -> {
+                keyAttributes.addAll(metadata.indexPartitionKeys(indexName));
+                keyAttributes.addAll(metadata.indexSortKeys(indexName));
+            });
+
+            return keyAttributes;
+        });
+
+        taggedAttributeNames.stream()
+                            .filter(attrName -> !allowedKeyAttributes.contains(attrName))
+                            .findFirst()
+                            .ifPresent(invalidAttribute -> {
+                                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: " + invalidAttribute);
+                            });
+    }
+
+    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) {
+            return metadata -> metadata.addCustomMetadataObject(CUSTOM_METADATA_KEY, Collections.singleton(attributeName))
+                                       .markAttributeAsKey(attributeName, attributeValueType);
+        }
+    }
+}

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

@@ -15,18 +15,23 @@
 
 package software.amazon.awssdk.enhanced.dynamodb.extensions;
 
+import static java.util.Collections.newSetFromMap;
+
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
 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.TableMetadata;
 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;
@@ -39,6 +44,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.
@@ -77,9 +90,25 @@
 @SdkPublicApi
 @ThreadSafe
 public final class AutoGeneratedUuidExtension implements DynamoDbEnhancedClientExtension {
+
+    /**
+     * Custom metadata keys under which AutoGeneratedUuidExtension/AutoGeneratedKeyExtension record attributes annotated with
+     * {@code @DynamoDbAutoGeneratedUuid}/{@code @DynamoDbAutoGeneratedKey}. Used to detect conflicts during schema validation.
+     */
     private static final String CUSTOM_METADATA_KEY =
         "software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedUuidExtension:AutoGeneratedUuidAttribute";
+    private static final String AUTOGENERATED_KEY_EXTENSION_METADATA_KEY =
+        "software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension:AutoGeneratedKeyAttribute";
+
     private static final AutoGeneratedUuidAttribute AUTO_GENERATED_UUID_ATTRIBUTE = new AutoGeneratedUuidAttribute();
+    private static final String AUTOGENERATED_KEY_ANNOTATION = "@DynamoDbAutoGeneratedKey";
+    private static final String AUTOGENERATED_UUID_ANNOTATION = "@DynamoDbAutoGeneratedUuid";
+
+    /**
+     * Stores the TableMetadata instances that have already been validated by this extension. Uses a ConcurrentHashMap to ensure
+     * thread-safe access during concurrent write operations.
+     */
+    private final Set<TableMetadata> validatedSchemas = newSetFromMap(new ConcurrentHashMap<>());
 
     private AutoGeneratedUuidExtension() {
     }
@@ -109,13 +138,31 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
             return WriteModification.builder().build();
         }
 
+        TableMetadata metadata = context.tableMetadata();
+        validateNoAnnotationConflict(metadata);
+
         Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
         customMetadataObject.forEach(key -> insertUuidInItemToTransform(itemToTransform, key));
         return WriteModification.builder()
                                 .transformedItem(Collections.unmodifiableMap(itemToTransform))
                                 .build();
     }
 
+    /**
+     * Validates (once per TableMetadata instance) that @DynamoDbAutoGeneratedKey and @DynamoDbAutoGeneratedUuid are not applied
+     * to the same attribute.
+     */
+    private void validateNoAnnotationConflict(TableMetadata tableMetadata) {
+        if (validatedSchemas.add(tableMetadata)) {
+            ExtensionsValidationUtils.validateNoAnnotationConflict(
+                tableMetadata,
+                AUTOGENERATED_KEY_EXTENSION_METADATA_KEY,
+                CUSTOM_METADATA_KEY,
+                AUTOGENERATED_KEY_ANNOTATION,
+                AUTOGENERATED_UUID_ANNOTATION);
+        }
+    }
+
     private void insertUuidInItemToTransform(Map<String, AttributeValue> itemToTransform,
                                              String key) {
         itemToTransform.put(key, AttributeValue.builder().s(UUID.randomUUID().toString()).build());